[PATCH v2 sched_ext/for-7.4] sched_ext: Keep proxy donors with slice left on the local DSQ

[Andrea Righi <arighi@nvidia.com>]

Commit ee172227d0dc ("sched_ext: Delegate proxy donor admission to BPF
schedulers") makes put_prev_task_scx() pass a retained proxy donor to
ops.enqueue() with SCX_ENQ_BLOCKED.

Some of these puts are only proxy bookkeeping; proxy_resched_idle()
drops the rq's donor reference before switching to idle when:
 - find_proxy_task() finds a remote mutex owner before the donor
   switches out,
 - proxy_migrate_task() detaches the donor before migration,
 - proxy_deactivate() blocks a donor whose owner cannot run.

In the first case, BPF has already selected a donor with slice left,
returning it to ops.enqueue() forces BPF to dispatch it again before
proxy resolution can continue. In the other two cases, the caller
deactivates the donor immediately after ops.enqueue(), undoing any
placement BPF makes.

Keep a donor with slice left at the head of the local DSQ instead, so
the next pick can resolve its owner, or deactivation can remove it
without an unnecessary BPF handoff.

For donors with slice left, this leaves BPF to handle meaningful
placement decisions rather than transient proxy-bookkeeping puts.

Fixes: ee172227d0dc ("sched_ext: Delegate proxy donor admission to BPF schedulers")
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
Changes in v2:
 - Keep blocked donors with slice left on the local DSQ regardless of
   IMMED and drop SCX_RQ_PROXY_PICK_PENDING (Tejun Heo)
 - Fold blocked-donor fallback into the normal enqueue path and restrict
   SCX_ENQ_LAST handling to unblocked tasks (Tejun Heo)
 - Link to v1: https://lore.kernel.org/r/20261001191216.2391359-1-arighi@nvidia.com

 kernel/sched/ext/ext.c      | 66 ++++++++++++++++++++++---------------
 kernel/sched/ext/internal.h | 16 ++++++---
 2 files changed, 50 insertions(+), 32 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 96c904b396019..aba035213bfc8 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1574,7 +1574,8 @@ static void dsq_inc_nr(struct scx_dispatch_q *dsq, struct task_struct *p, u64 en
 
 		/*
 		 * If @rq already had other tasks or the current task is not
-		 * done yet, @p can't go on the CPU immediately. Re-enqueue.
+		 * done yet, @p can't go on the CPU immediately. Check whether
+		 * it needs to be reenqueued.
 		 */
 		if (unlikely(dsq->nr > 1 || !rq_is_open(rq, enq_flags)))
 			scx_schedule_reenq_local(rq, 0);
@@ -2563,8 +2564,17 @@ static void wakeup_preempt_scx(struct rq *rq, struct task_struct *p, int wake_fl
 		    p->is_blocked) {
 			struct scx_sched *sch = scx_task_sched(p);
 
-			if (sch && (sch->ops.flags & SCX_OPS_ENQ_BLOCKED))
+			if (sch && (sch->ops.flags & SCX_OPS_ENQ_BLOCKED)) {
+				/*
+				 * A proxy-migrated donor is dequeued before this callback.
+				 * Recheck a local IMMED donor still on its wake_cpu after
+				 * ttwu_runnable() clears is_blocked.
+				 */
+				if ((p->scx.flags & SCX_TASK_IMMED) &&
+				    p->scx.dsq == &rq->scx.local_dsq)
+					scx_schedule_reenq_local(rq, 0);
 				resched_curr(rq);
+			}
 		}
 		return;
 	}
@@ -3318,6 +3328,9 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 	 *
 	 * - A non-IMMED HEAD task can get queued in front of an IMMED task
 	 *   between the IMMED queueing and the subsequent scheduling event.
+	 *
+	 * A blocked IMMED donor may make this scan a no-op. Skipping it does
+	 * not schedule another scan, so avoid separate accounting for it.
 	 */
 	if (unlikely(rq->scx.local_dsq.nr > 1 && rq->scx.nr_immed))
 		scx_schedule_reenq_local(rq, 0);
@@ -3516,37 +3529,31 @@ 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);
 
-		/* Delegate retained donor admission to its owning BPF scheduler. */
-		if (p->is_blocked) {
-			/*
-			 * If the donor is the same and only the mutex owner
-			 * changes, avoid triggering another ops.enqueue(): the
-			 * BPF scheduler has already admitted the donor, so it
-			 * can continue running.
-			 */
-			if (next == p)
-				goto switch_class;
-
-			if (WARN_ON_ONCE(!sch))
-				goto switch_class;
-			WARN_ON_ONCE(!(sch->ops.flags & SCX_OPS_ENQ_BLOCKED));
-			scx_do_enqueue_task(rq, p, 0, -1);
+		/*
+		 * If the donor is the same and only the mutex owner changes,
+		 * avoid triggering another ops.enqueue(): the BPF scheduler has
+		 * already admitted the donor, so it can continue running.
+		 */
+		if (p->is_blocked && next == p)
 			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
 		 * forcing a different task. Leave it at the head of the local
-		 * DSQ unless it was an IMMED task. IMMED tasks should not
-		 * linger on a busy CPU, reenqueue them to the BPF scheduler.
+		 * DSQ unless it was an unblocked IMMED task. Such tasks should not
+		 * linger on a busy CPU, so reenqueue them to the BPF scheduler.
+		 *
+		 * A blocked donor's progress depends on its mutex owner. Moving the
+		 * donor elsewhere does not move its owner, so keep it local even if
+		 * it is IMMED.
 		 *
 		 * An open rescue must keep @p on the local DSQ even if the
 		 * scheduler zeroed the slice in ops.stopping() above.
 		 */
 		if ((p->scx.slice || unlikely(p == scx_rescuee(rq))) &&
 		    !scx_bypassing(sch, cpu_of(rq))) {
-			if (p->scx.flags & SCX_TASK_IMMED) {
+			if ((p->scx.flags & SCX_TASK_IMMED) && !p->is_blocked) {
 				p->scx.flags |= SCX_TASK_REENQ_PREEMPTED;
 				scx_do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
 			} else {
@@ -3563,6 +3570,8 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 						enq_flags |= SCX_ENQ_HEAD;
 				} else {
 					enq_flags |= SCX_ENQ_HEAD;
+					if (p->scx.flags & SCX_TASK_IMMED)
+						enq_flags |= SCX_ENQ_IMMED;
 				}
 
 				scx_dispatch_enqueue(sch, rq, &rq->scx.local_dsq, p, 0, 0,
@@ -3582,7 +3591,8 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 		 * locally runnable and can legitimately go idle with @p still
 		 * runnable (see do_pick_task_scx()).
 		 */
-		if (next && sched_class_above(&ext_sched_class, next->sched_class) &&
+		if (!p->is_blocked &&
+		    next && sched_class_above(&ext_sched_class, next->sched_class) &&
 		    scx_task_can_stay_on_cpu(rq, p)) {
 			WARN_ON_ONCE(!sched_core_enabled(rq) &&
 				     !(sch->ops.flags & SCX_OPS_ENQ_LAST));
@@ -4811,10 +4821,12 @@ static void process_ddsp_deferred_locals(struct rq *rq)
  * - %SCX_REENQ_TSR_RQ_OPEN: Set by reenq_local() before the walk if
  *   rq_is_open() is true.
  *
- * An IMMED task is kept (returns %false) only if it's the first task in the DSQ
- * AND the current task is done — i.e. it will execute immediately. All other
- * IMMED tasks are reenqueued. This means if a non-IMMED task sits at the head,
- * every IMMED task behind it gets reenqueued.
+ * An unblocked IMMED task is kept (returns %false) only if it's the first task
+ * in the DSQ AND the current task is done, so it will execute immediately.
+ * Other unblocked IMMED tasks are reenqueued. Blocked proxy donors are not
+ * reenqueued for IMMED alone, even if they are not first or the rq is busy.
+ * If a non-IMMED task sits at the head, every unblocked IMMED task behind it
+ * gets reenqueued.
  *
  * Reenqueued tasks go through ops.enqueue() with %SCX_ENQ_REENQ |
  * %SCX_TASK_REENQ_IMMED. If the BPF scheduler dispatches back to the same local
@@ -4835,7 +4847,7 @@ static bool local_task_should_reenq(struct rq *rq, struct task_struct *p,
 
 	*reason = SCX_TASK_REENQ_KFUNC;
 
-	if ((p->scx.flags & SCX_TASK_IMMED) &&
+	if ((p->scx.flags & SCX_TASK_IMMED) && !p->is_blocked &&
 	    (!first || !(*reenq_flags & SCX_REENQ_TSR_RQ_OPEN))) {
 		__scx_add_event(scx_task_sched(p), SCX_EV_REENQ_IMMED, 1);
 		*reason = SCX_TASK_REENQ_IMMED;
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index d65fec631bdf7..dfea7b1f95f98 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1780,14 +1780,14 @@ enum scx_enq_flags {
 	SCX_ENQ_PREEMPT_LAZY	= 1LLU << 35,
 
 	/*
-	 * Only allowed on local DSQs. Guarantees that the task either gets
-	 * on the CPU immediately and stays on it, or gets reenqueued back
-	 * to the BPF scheduler. It will never linger on a local DSQ or be
-	 * silently put back after preemption.
+	 * Only allowed on local DSQs. Guarantees that an unblocked task either
+	 * gets on the CPU immediately and stays on it, or gets reenqueued back
+	 * to the BPF scheduler. A blocked proxy donor can stay on the local DSQ
+	 * with slice left because its progress depends on its mutex owner.
 	 *
 	 * The protection persists until the next fresh enqueue - it
 	 * survives SAVE/RESTORE cycles, slice extensions and preemption.
-	 * If the task can't stay on the CPU for any reason, it gets
+	 * If an unblocked task can't stay on the CPU for any reason, it gets
 	 * reenqueued back to the BPF scheduler.
 	 *
 	 * Exiting and migration-disabled tasks bypass ops.enqueue() and
@@ -1829,6 +1829,12 @@ enum scx_enq_flags {
 	/*
 	 * The task is blocked on a mutex and is being kept runnable as a proxy
 	 * donor. Only passed to ops.enqueue() when %SCX_OPS_ENQ_BLOCKED is set.
+	 *
+	 * Blocking on the mutex does not enqueue the task by itself. A donor put
+	 * with slice left stays at the head of the local DSQ, including IMMED
+	 * donors. It is passed to ops.enqueue() when its slice runs out
+	 * (including an SCX preemption that zeros it), or proxy execution moves
+	 * it to the CPU of the mutex owner.
 	 */
 	SCX_ENQ_BLOCKED		= 1LLU << 42,
 
-- 
2.55.0