diff --git a/arch/s390/kernel/runtime_instr.c b/arch/s390/kernel/runtime_instr.c
index 429d3a782f1c..b9738ae2e1de 100644
--- a/arch/s390/kernel/runtime_instr.c
+++ b/arch/s390/kernel/runtime_instr.c
@@ -49,11 +49,13 @@ void exit_thread_runtime_instr(void)
 {
 	struct task_struct *task = current;
 
+	preempt_disable();
 	if (!task->thread.ri_cb)
 		return;
 	disable_runtime_instr();
 	kfree(task->thread.ri_cb);
 	task->thread.ri_cb = NULL;
+	preempt_enable();
 }
 
 SYSCALL_DEFINE1(s390_runtime_instr, int, command)
@@ -64,9 +66,7 @@ SYSCALL_DEFINE1(s390_runtime_instr, int, command)
 		return -EOPNOTSUPP;
 
 	if (command == S390_RUNTIME_INSTR_STOP) {
-		preempt_disable();
 		exit_thread_runtime_instr();
-		preempt_enable();
 		return 0;
 	}