[PATCH sched_ext/for-6.12-fixes] sched_ext: Call __balance_callbacks() from __scx_task_iter_rq_unlock()

[Tejun Heo <tj@kernel.org>]

While tasks are being iterated to be switched in and out of SCX,
__scx_task_iter_rq_unlock() is called to unlock rq. As
sched_class->switched_from() and friends may schedule balance callbacks,
they should be executed before moving onto the next task.

This brekage is currently theoretical as only RT and DL schedule balance
callbacks on class switches and SCX only switches from/to fair.

Make __balance_callbacks() global and call it from
__scx_task_iter_rq_unlock().

Signed-off-by: Tejun Heo <tj@kernel.org>
Suggested-by: Peter Zijlstra <peterz@infradead.org>
Link: http://lkml.kernel.org/r/20241030154551.GN14555@noisy.programming.kicks-ass.net
---
Hello,

If this looks okay, I'll route it through sched_ext/for-6.12-fixes.

Thanks.

 kernel/sched/core.c  |    8 +-------
 kernel/sched/ext.c   |    2 ++
 kernel/sched/sched.h |    5 +++++
 3 files changed, 8 insertions(+), 7 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5028,7 +5028,7 @@ struct balance_callback *splice_balance_
 	return __splice_balance_callbacks(rq, true);
 }
 
-static void __balance_callbacks(struct rq *rq)
+void __balance_callbacks(struct rq *rq)
 {
 	do_balance_callbacks(rq, __splice_balance_callbacks(rq, false));
 }
@@ -5044,12 +5044,6 @@ void balance_callbacks(struct rq *rq, st
 	}
 }
 
-#else
-
-static inline void __balance_callbacks(struct rq *rq)
-{
-}
-
 #endif
 
 static inline void
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1315,6 +1315,8 @@ static void scx_task_iter_start(struct s
 static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
 {
 	if (iter->locked) {
+		/* ->switched_from() may have scheduled balance callbacks */
+		__balance_callbacks(iter->rq);
 		task_rq_unlock(iter->rq, iter->locked, &iter->rf);
 		iter->locked = NULL;
 	}
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -3908,6 +3908,7 @@ extern void check_class_changed(struct r
 #ifdef CONFIG_SMP
 extern struct balance_callback *splice_balance_callbacks(struct rq *rq);
 extern void balance_callbacks(struct rq *rq, struct balance_callback *head);
+extern void __balance_callbacks(struct rq *rq);
 #else
 
 static inline struct balance_callback *splice_balance_callbacks(struct rq *rq)
@@ -3919,6 +3920,10 @@ static inline void balance_callbacks(str
 {
 }
 
+static inline void __balance_callbacks(struct rq *rq)
+{
+}
+
 #endif
 
 #ifdef CONFIG_SCHED_CLASS_EXT