[PATCH 1/2] sched_ext: Don't queue balance callbacks from a core-sched pick

[Andrea Righi <arighi@nvidia.com>]

Commit ff0545e0c263 ("sched_ext: Don't leave balance callbacks queued
across a pick retry") queues the deferred work as a balance callback
only when do_pick_task_scx() returns a final pick. However, under core
scheduling, a pick that returns a task isn't final either:
pick_next_task() goes on to pick the SMT siblings, and their dispatch
can release the shared rq lock with the callback still queued,
triggering the rq_pin_lock() warning.

Moreover, dispatch_core_pick() returns early when another dispatch is
in flight on the rq, and do_pick_task_scx() then consumes
%SCX_RQ_BAL_CB_PENDING set by that dispatch while it's still running.

This can be reproduced running scx_qmap -I on SMT siblings with
coresched-tagged workloads that frequently sleep, mixed with untagged
CPU-bound tasks.

Fix this by always deferring the pending work to the irq_work path from
dispatch_core_pick() under core scheduling, for both the local and the
sibling picks, and by leaving %SCX_RQ_BAL_CB_PENDING alone in
do_pick_task_scx(), so that it's handled by the dispatch that set it.
Sample the core scheduling state once at the start of the pick, since
core scheduling can get enabled while a non-core dispatch has the rq lock
released, and the follow-up work must match the dispatch path that ran.

Fixes: 4c95380701f5 ("sched/ext: Fold balance_scx() into pick_task_scx()")
Cc: stable@vger.kernel.org # v6.19+
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c | 35 ++++++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 0e93a5c45ede4..0c5ed97c4d080 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3390,15 +3390,20 @@ static enum scx_dsp_verdict dispatch_core_pick(struct rq *rq, struct rq_flags *r
 
 	verdict = dispatch_one(rq, prev);
 
-	if (cpu_of(rq) != smp_processor_id() &&
-	    unlikely(rq->scx.flags & SCX_RQ_BAL_CB_PENDING)) {
-		/*
-		 * Balance callbacks must run in the context that queued them,
-		 * so they can't be queued on another CPU's rq. Run the deferred
-		 * work directly instead.
-		 */
+	/*
+	 * A pick that returns a task isn't final under core scheduling: the
+	 * selection goes on to pick the siblings, and their dispatch can
+	 * release the rq lock, which can't happen with balance callbacks
+	 * queued. Don't queue the deferred work as a balance callback, for this
+	 * CPU or a sibling, and run it from irq_work instead.
+	 *
+	 * do_pick_task_scx() leaves %SCX_RQ_BAL_CB_PENDING alone under core
+	 * scheduling, so that the early return above leaves it to the dispatch
+	 * in flight, which handles it here.
+	 */
+	if (unlikely(rq->scx.flags & SCX_RQ_BAL_CB_PENDING)) {
 		rq->scx.flags &= ~SCX_RQ_BAL_CB_PENDING;
-		run_deferred(rq);
+		schedule_deferred(rq);
 	}
 
 	rq_repin_lock(rq, rf);
@@ -3421,6 +3426,7 @@ static struct task_struct *
 do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
 {
 	struct task_struct *prev = rq->curr;
+	bool core_sched = sched_core_enabled(rq);
 	enum scx_dsp_verdict verdict;
 	struct task_struct *p;
 
@@ -3429,7 +3435,12 @@ do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
 
 	rq_modified_begin(rq, &ext_sched_class);
 
-	if (sched_core_enabled(rq))
+	/*
+	 * Core scheduling can get enabled while dispatch_pick() has the rq lock
+	 * released. Stick to @core_sched, so that the follow-up work below
+	 * matches the dispatch that ran.
+	 */
+	if (core_sched)
 		verdict = dispatch_core_pick(rq, rf, prev);
 	else
 		verdict = dispatch_pick(rq, rf, prev);
@@ -3495,7 +3506,9 @@ do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
 	 * for the next pick.
 	 */
 	if (p != RETRY_TASK && (p || !force_scx)) {
-		maybe_queue_balance_callback(rq);
+		/* under core scheduling, see dispatch_core_pick() */
+		if (!core_sched)
+			maybe_queue_balance_callback(rq);
 
 		/*
 		 * Defer to a balance callback which can drop rq lock and enable
@@ -3508,7 +3521,7 @@ do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
 			queue_balance_callback(rq, &rq->scx.kick_sync_bal_cb,
 					       kick_sync_wait_bal_cb);
 		}
-	} else if (rq->scx.flags & SCX_RQ_BAL_CB_PENDING) {
+	} else if (!core_sched && (rq->scx.flags & SCX_RQ_BAL_CB_PENDING)) {
 		rq->scx.flags &= ~SCX_RQ_BAL_CB_PENDING;
 		schedule_deferred(rq);
 	}
-- 
2.56.0