[PATCH 2/2] sched_ext: Don't leave the kick_sync wait queued across a core-sched pick

[Andrea Righi <arighi@nvidia.com>]

The SCX_KICK_WAIT wait is deferred to rq->scx.kick_sync_bal_cb, which a
final pick queues on the local rq. Under core scheduling, a pick that
returns a task isn't final: pick_next_task() goes on to pick the SMT
siblings, and their dispatch can release the shared rq lock with the
callback still queued. A CPU acquiring the rq lock in that window then
finds a pending balance callback and triggers the warning in
rq_pin_lock().

Unlike the deferred work, the kick_sync wait can't be moved to irq_work,
since it needs to wait with IRQs enabled. However, during a core-wide
selection the rq lock is only released by a dispatch, which already
restarts the selection through lock_drop_seq.

Fix this by taking the kick_sync callback off the local rq before each
dispatch of the selection. If the rq lock isn't released, queue it
again right away. Otherwise, set kick_sync_pending again, so that the
next final pick of this CPU queues it, as done for a non-core pick that
returns RETRY_TASK.

Fixes: 415cb193bb97 ("sched_ext: Fix SCX_KICK_WAIT deadlock by deferring wait to balance callback")
Cc: stable@vger.kernel.org # v7.0+
Suggested-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c | 43 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 0c5ed97c4d080..2348d5b72eb26 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3369,6 +3369,28 @@ static enum scx_dsp_verdict dispatch_pick(struct rq *rq, struct rq_flags *rf,
 }
 
 #ifdef CONFIG_SCHED_CORE
+/*
+ * Take @head off @rq's balance callbacks if it's queued. Returns whether it
+ * was queued.
+ */
+static bool dequeue_balance_callback(struct rq *rq, struct balance_callback *head)
+{
+	struct balance_callback **pprev;
+
+	lockdep_assert_rq_held(rq);
+
+	for (pprev = &rq->balance_callback; *pprev; pprev = &(*pprev)->next) {
+		if (*pprev != head)
+			continue;
+
+		*pprev = head->next;
+		head->next = NULL;
+		return true;
+	}
+
+	return false;
+}
+
 /*
  * Dispatch for a pick when core scheduling is enabled. The selection picks for
  * all SMT siblings and the rq_i->core_pick state it builds must stay atomic
@@ -3379,8 +3401,10 @@ static enum scx_dsp_verdict dispatch_pick(struct rq *rq, struct rq_flags *rf,
 static enum scx_dsp_verdict dispatch_core_pick(struct rq *rq, struct rq_flags *rf,
 					       struct task_struct *prev)
 {
+	struct rq *this_rq = this_rq();
 	enum scx_dsp_verdict verdict;
 	u32 seq = rq->scx.lock_drop_seq;
+	bool kick_sync_queued;
 
 	/* another dispatch is in flight on @rq, let that handle it */
 	if (rq->scx.flags & SCX_RQ_IN_DISPATCH)
@@ -3388,6 +3412,16 @@ static enum scx_dsp_verdict dispatch_core_pick(struct rq *rq, struct rq_flags *r
 
 	rq_unpin_lock(rq, rf);
 
+	/*
+	 * dispatch_one() can release the rq lock while the selection is in
+	 * progress (see below), so this CPU's kick_sync wait queued by an
+	 * earlier pick of the selection can't stay queued across it. Take it
+	 * off and queue it again if the rq lock isn't released. Otherwise, the
+	 * selection restarts, so leave it pending for the next final pick of
+	 * this CPU.
+	 */
+	kick_sync_queued = dequeue_balance_callback(this_rq, &this_rq->scx.kick_sync_bal_cb);
+
 	verdict = dispatch_one(rq, prev);
 
 	/*
@@ -3409,8 +3443,15 @@ static enum scx_dsp_verdict dispatch_core_pick(struct rq *rq, struct rq_flags *r
 	rq_repin_lock(rq, rf);
 
 	/* if dispatch_one() released the rq lock, restart the selection */
-	if (rq->scx.lock_drop_seq != seq)
+	if (rq->scx.lock_drop_seq != seq) {
+		if (kick_sync_queued)
+			this_rq->scx.kick_sync_pending = true;
 		return SCX_DSP_RETRY;
+	}
+
+	if (kick_sync_queued)
+		queue_balance_callback(this_rq, &this_rq->scx.kick_sync_bal_cb,
+				       kick_sync_wait_bal_cb);
 
 	return verdict;
 }
-- 
2.56.0