[PATCH 08/12] sched_ext: Handle blocked donor migration with proxy execution

[Andrea Righi <arighi@nvidia.com>]

From: John Stultz <jstultz@google.com>

With proxy execution enabled, mutex-blocked donors stay runnable so
their scheduling context can execute the lock owner.

sched_ext may normally relocate a queued blocked donor: set_task_cpu()
updates both task_cpu() and wake_cpu, making the destination rq the
donor's new callback home. This remains valid if the donor was
previously proxy-migrated: the BPF scheduler is intentionally selecting
a new return destination, which the next proxy migration will preserve.
The usual migration-disabled, affinity, and rq-online checks apply.

An active donor cannot be moved normally because the source rq still
references it through rq->donor for scheduling-class operations and
runtime accounting. Moving it would associate the task with a different
rq while leaving those source-rq references in place. The core proxy
migration path avoids this by switching the rq donor to idle before
moving the scheduling context.

Check active-donor state in task_can_move_from_locked_rq(), where the
source rq is held and the result is authoritative. Keep it out of the
preliminary lockless DSQ filter, so a donor transition during the
rq-lock handoff is resolved by the locked recheck rather than consuming
a dispatch kick and skipping the task.

Allow normal sched_ext migration of blocked donors unless the donor is
active on its rq. In that case, defer migration until it is switched out
or leave it to the proxy machinery.

Keep blocked donors on the local DSQ so they remain visible to the proxy
pick path. This is a temporary solution: in a later commit, proxy donor
admission will be delegated to the BPF scheduler and donors will instead
be routed through ops.enqueue().

Signed-off-by: John Stultz <jstultz@google.com>
Co-developed-by: Andrea Righi <arighi@nvidia.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>

 interactive rebase in progress; onto c0414e1387386
---
 kernel/sched/ext/ext.c | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index a466adf8e1548..bf53f2be77f05 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2455,6 +2455,10 @@ static bool task_can_move_from_locked_rq(struct task_struct *p)
 	if (task_on_cpu(src_rq, p))
 		return false;
 
+	/* Don't move the current scheduling context off its source rq. */
+	if (task_current_donor(src_rq, p))
+		return false;
+
 	return !is_migration_disabled(p);
 }
 
@@ -3122,6 +3126,25 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 	if (p->scx.flags & SCX_TASK_QUEUED) {
 		set_task_runnable(rq, p);
 
+		/*
+		 * Mutex-blocked donors stay queued on the runqueue under proxy
+		 * execution, but the donor never runs as itself, proxy-exec
+		 * walks the blocked_on chain on the next __schedule() and runs
+		 * the lock owner in its place.
+		 *
+		 * Put the donor on the local DSQ directly so pick_next_task()
+		 * can still see it. find_proxy_task() will either run the chain
+		 * owner or deactivate the donor so the wakeup path can return it
+		 * and let BPF make a new dispatch decision once it is unblocked.
+		 *
+		 * This is preparatory code: a later patch will delegate blocked-donor
+		 * admission to the BPF scheduler.
+		 */
+		if (p->is_blocked) {
+			scx_dispatch_enqueue(sch, rq, &rq->scx.local_dsq, p, 0);
+			goto switch_class;
+		}
+
 		/*
 		 * If @p has slice left and is being put, @p is getting
 		 * preempted by a higher priority scheduler class or core-sched
-- 
2.55.0