sched/core: Robustify preemption leak checks

When we warn about a preempt_count leak; reset the preempt_count to
the known good value such that the problem does not ripple forward.

This is most important on x86 which has a per cpu preempt_count that is
not saved/restored (after this series). So if you schedule with an
invalid (!2*PREEMPT_DISABLE_OFFSET) preempt_count the next task is
messed up too.

Enforcing this invariant limits the borkage to just the one task.

Signed-off-by: Peter Zijlstra (Intel) <[e-mail]>
Reviewed-by: Frederic Weisbecker <[e-mail]>
Reviewed-by: Thomas Gleixner <[e-mail]>
Reviewed-by: Steven Rostedt <[e-mail]>
Cc: Linus Torvalds <[e-mail]>
Cc: Mike Galbraith <[e-mail]>
Cc: Peter Zijlstra <[e-mail]>
Cc: [e-mail]
Signed-off-by: Ingo Molnar <[e-mail]>

kernel/exit.c

···
 	smp_mb();
 	raw_spin_unlock_wait(&tsk->pi_lock);
 
-	if (unlikely(in_atomic()))
+	if (unlikely(in_atomic())) {
 		pr_info("note: %s[%d] exited with preempt_count %d\n",
 			current->comm, task_pid_nr(current),
 			preempt_count());
+		preempt_count_set(PREEMPT_ENABLED);
+	}
 
 	/* sync mm's RSS info before statistics gathering */
 	if (tsk->mm)

kernel/sched/core.c

···
 	BUG_ON(unlikely(task_stack_end_corrupted(prev)));
 #endif
 
-	if (unlikely(in_atomic_preempt_off()))
+	if (unlikely(in_atomic_preempt_off())) {
 		__schedule_bug(prev);
+		preempt_count_set(PREEMPT_DISABLED);
+	}
 	rcu_sleep_check();
 
 	profile_hit(SCHED_PROFILING, __builtin_return_address(0));