[PATCH sched/core] sched: Warn if a sched_class says yes on balance() but no on pick_task()

[Tejun Heo <tj@kernel.org>]

If a sched_class says yes on balance() but no on pick_task(),
__pick_next_task() will end up descending to lower sched classes which
haven't been balanced. This can lead to sub-optimal scheduling decisions
and, especially for sched_ext, stalls. Detect the condition and warn.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
This triggers immedately on the current tip/sched/core. Best to apply after
fixing the existing issue.

Thanks.

 kernel/sched/core.c |   22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5893,8 +5893,8 @@ static inline void schedule_debug(struct
 	schedstat_inc(this_rq()->sched_count);
 }
 
-static void prev_balance(struct rq *rq, struct task_struct *prev,
-			 struct rq_flags *rf)
+static const struct sched_class *
+prev_balance(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
 {
 	const struct sched_class *start_class = prev->sched_class;
 	const struct sched_class *class;
@@ -5919,8 +5919,10 @@ static void prev_balance(struct rq *rq,
 	 */
 	for_active_class_range(class, start_class, &idle_sched_class) {
 		if (class->balance && class->balance(rq, prev, rf))
-			break;
+			return class;
 	}
+
+	return &idle_sched_class;
 }
 
 /*
@@ -5929,7 +5931,7 @@ static void prev_balance(struct rq *rq,
 static inline struct task_struct *
 __pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
 {
-	const struct sched_class *class;
+	const struct sched_class *class, *balanced;
 	struct task_struct *p;
 
 	rq->dl_server = NULL;
@@ -5960,23 +5962,29 @@ __pick_next_task(struct rq *rq, struct t
 	}
 
 restart:
-	prev_balance(rq, prev, rf);
+	balanced = prev_balance(rq, prev, rf);
 
 	for_each_active_class(class) {
 		if (class->pick_next_task) {
 			p = class->pick_next_task(rq, prev);
 			if (p)
-				return p;
+				goto picked;
 		} else {
 			p = class->pick_task(rq);
 			if (p) {
 				put_prev_set_next_task(rq, prev, p);
-				return p;
+				goto picked;
 			}
 		}
 	}
 
 	BUG(); /* The idle class should always have a runnable task. */
+
+picked:
+	WARN_ONCE(sched_class_above(balanced, class),
+		  "%ps->balance() returned true but no task, %ps picked instead\n",
+		  balanced, class);
+	return p;
 }
 
 #ifdef CONFIG_SCHED_CORE