[PATCH sched_ext/for-7.4-fixes] sched_ext: Fix stale pick when a wakeup resets the donor during dispatch

[Tejun Heo <tj@kernel.org>]

With proxy execution, do_pick_task_scx() passes rq->donor to dispatch_one()
as @prev. Unlike rq->curr, the donor can be moved to another CPU by a wakeup
while dispatch_one() has the rq lock released:

1. dispatch_one() releases the rq lock, e.g. to move a task from another
   CPU.
2. Another CPU wakes the donor, which proxy execution migrated to this CPU.
   proxy_needs_return() dequeues it and resets the donor to rq->curr, and
   the wakeup enqueues it, possibly on another CPU.
3. dispatch_one() sees @prev queued and returns SCX_DSP_PREV.
4. do_pick_task_scx() returns @prev, and this CPU runs a task that is queued
   on another rq.

The other verdicts are stale too. If the dispatch finds nothing before the
wakeup enqueues @prev, SCX_DSP_NONE picks idle, and put_prev_task_scx()
parks rq->curr on the local DSQ with nothing to reschedule the CPU, or trips
its SCX_OPS_ENQ_LAST warning. SCX_DSP_LOCAL preempts rq->curr even when the
dispatch would have kept it.

rq_modified_above() catches the reset only when rq->curr is in a class above
ext, and the dl_server pick skips the check.

The verdict, whatever it is, describes the old donor. Restart the pick when
@prev is no longer the donor, as dispatch_core_pick() already does whenever
the dispatch released the rq lock. The retry runs with rq->curr as the
donor, and proxy_needs_return() returns early for the current task, so the
donor can't change again.

7a3a6bfbd62a ("sched: Rework prev_balance() to avoid stale prev references")
closed the same window in prev_balance() by reading rq->donor instead of
taking @prev. That doesn't work here, as the dispatch never evaluated
rq->curr, which may not even be a sched_ext task.

Verified with a synthetic reproducer.

Fixes: c6fe97c34a1a ("sched: Allow enabling proxy exec with sched_ext")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 15123f181b5b..3119a1938c6e 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3936,6 +3936,16 @@ static enum scx_dsp_verdict dispatch_pick(struct rq *rq, struct rq_flags *rf,
 	verdict = dispatch_one(rq, prev);
 	rq_repin_lock(rq, rf);
 
+	/*
+	 * With proxy execution, @prev is the donor and need not be the running
+	 * task. A wakeup while dispatch_one() has the rq lock released can
+	 * reset the donor and move @prev to another CPU. The verdict was
+	 * computed for @prev and doesn't apply once the donor changed. Restart
+	 * the pick.
+	 */
+	if (unlikely(prev != rq->donor))
+		return SCX_DSP_RETRY;
+
 	return verdict;
 }
 
-- 
2.56.0