[PATCH sched_ext/for-7.4-fixes] sched_ext: Avoid unnecessary IMMED reenqueues on idle CPUs

[Andrea Righi <arighi@nvidia.com>]

When two IMMED tasks wake on an idle CPU, the kernel can unnecessarily
send both back to BPF instead of keeping the first task ready to run.

The first wakeup inserts a task into the local DSQ and raises
rq->next_class to ext before the CPU leaves idle. The second insertion
queues a reenqueue scan. rq_is_open() sees next_class == ext and reports
the CPU unavailable, although it is still idle. The scan then reenqueues
the head along with the task behind it, losing a placement that could
have been honored on the next pick.

Keep the local DSQ head through the reenqueue scan when the CPU is idle
and, if core scheduling is enabled, its cookie matches. Skipped tasks
still consume the first position, so skipping the head because of
pending migration or scheduler ownership can't let a later IMMED task
stay queued even though it cannot run next.

A local reproducer runs in a four-vCPU virtme-ng guest: two futex
waiters pinned to CPU 0 are woken together from CPU 2, and the BPF
scheduler inserts their wakeups into the local DSQ with SCX_ENQ_IMMED.
Across 10000 pairs, kernel instrumentation counts 9776 unnecessary head
reenqueues while the CPU is idle and next_class is ext. With this fix,
this count is zero (IMMED tasks behind the head still return to BPF).

An additional reproducer exercises the same problem with core scheduling
enabled: it uses a 1 ms kernel mutex hold, a weighted-vruntime BPF
scheduler with SCX_OPS_ENQ_BLOCKED and SCX_OPS_ALWAYS_ENQ_IMMED,
distinct owner and waiter core cookies, and an additional futex wakeup
stream targeting the waiter CPU. Over 2000 measured trials per case, the
mutex waiter receives 154 IMMED reenqueues with a same-CPU owner and 187
with a cross-CPU owner. Both counts drop to zero with this fix.

Fixes: 98d709cba319 ("sched_ext: Implement SCX_ENQ_IMMED")
Reported-by: Tejun Heo <tj@kernel.org>
Link: https://lore.kernel.org/r/1548cc9309a6e0894e358262e0b3d4ee@kernel.org
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 1b098728676f3..ba45fd017afcd 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -429,6 +429,12 @@ static struct scx_dispatch_q *bypass_enq_target_dsq(struct scx_sched *sch, struc
 	return dsq;
 }
 
+static struct task_struct *first_local_task(struct rq *rq)
+{
+	return list_first_entry_or_null(&rq->scx.local_dsq.list,
+					struct task_struct, scx.dsq_list.node);
+}
+
 /**
  * rq_is_open - Is the rq available for immediate execution of an SCX task?
  * @rq: rq to test
@@ -461,6 +467,18 @@ static bool rq_is_open(struct rq *rq, u64 enq_flags)
 	 * @rq is either picking, in transition to, or running an SCX task.
 	 */
 
+	/*
+	 * A local insertion can raise next_class to ext while the donor is
+	 * still idle, including after proxy_resched_idle(). The local DSQ head
+	 * is next to run if its cookie matches, so keep it through IMMED reenq.
+	 */
+	if (rq->donor == rq->idle) {
+		struct task_struct *first = first_local_task(rq);
+
+		if (first && sched_cpu_cookie_match(rq, first))
+			return true;
+	}
+
 	/*
 	 * If we're in the dispatch path holding rq lock, $donor may or may not
 	 * be ready depending on whether the on-going dispatch decides to extend
@@ -3909,12 +3927,6 @@ static void kick_sync_wait_bal_cb(struct rq *rq)
 		goto retry;
 }
 
-static struct task_struct *first_local_task(struct rq *rq)
-{
-	return list_first_entry_or_null(&rq->scx.local_dsq.list,
-					struct task_struct, scx.dsq_list.node);
-}
-
 /*
  * Run dispatch for a pick. The follow-up work is queued by do_pick_task_scx()
  * once the pick is final.
@@ -5197,6 +5209,9 @@ static u32 reenq_local(struct scx_sched *sch, struct rq *rq, u64 reenq_flags)
 		u32 reason;
 
 		/*
+		 * Skipped tasks still consume the first position in the DSQ. Set
+		 * NOT_FIRST on every path that skips local_task_should_reenq().
+		 *
 		 * If @p is being migrated, @p's current CPU may not agree with
 		 * its allowed CPUs and the migration_cpu_stop is about to
 		 * deactivate and re-activate @p anyway. Skip re-enqueueing.
@@ -5207,11 +5222,15 @@ static u32 reenq_local(struct scx_sched *sch, struct rq *rq, u64 reenq_flags)
 		 * the current local DSQ for running tasks and thus are not
 		 * visible to the BPF scheduler.
 		 */
-		if (p->migration_pending)
+		if (p->migration_pending) {
+			reenq_flags |= SCX_REENQ_TSR_NOT_FIRST;
 			continue;
+		}
 
-		if (!scx_is_descendant(task_sch, sch))
+		if (!scx_is_descendant(task_sch, sch)) {
+			reenq_flags |= SCX_REENQ_TSR_NOT_FIRST;
 			continue;
+		}
 
 		if (!local_task_should_reenq(rq, p, &reenq_flags, &reason))
 			continue;
-- 
2.56.0