[PATCH 3/3] selftests/sched_ext: Fix consume_immed test reliability

[zhidao su <soolaugust@gmail.com>]

The consume_immed test was failing because nr_consume_immed_reenq
stayed 0. Two issues:

1. Workers were spread across CPUs, so CPU 0's local DSQ rarely
   accumulated multiple tasks. Fix: pin all workers to CPU 0 and
   increase NUM_WORKERS to 8 to ensure USER_DSQ is always backlogged.

2. ops.dispatch() called scx_bpf_dsq_move_to_local() only once, so
   CPU 0's local DSQ would contain exactly 1 task after dispatch.
   The IMMED slow path requires dsq->nr > 1 at the time of insertion
   (in dsq_inc_nr()), so a single dispatch call never triggers it.

   Fix: call scx_bpf_dsq_move_to_local() in a loop (up to 4 times)
   within a single ops.dispatch() invocation. The second call finds
   dsq->nr already 1, so dsq->nr increments to 2, triggering
   schedule_reenq_local() and the IMMED slow path.

Signed-off-by: zhidao su <suzhidao@xiaomi.com>
---
 .../selftests/sched_ext/consume_immed.bpf.c   | 25 ++++++++++++++-----
 .../selftests/sched_ext/consume_immed.c       | 14 ++++++++++-
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/consume_immed.bpf.c b/tools/testing/selftests/sched_ext/consume_immed.bpf.c
index 9c7808f5abe1..e99bea0b2c24 100644
--- a/tools/testing/selftests/sched_ext/consume_immed.bpf.c
+++ b/tools/testing/selftests/sched_ext/consume_immed.bpf.c
@@ -11,9 +11,9 @@
  *                   explicit SCX_ENQ_IMMED in enq_flags (requires v2 kfunc)
  *
  * Worker threads belonging to test_tgid are inserted into USER_DSQ.
- * ops.dispatch() on CPU 0 consumes from USER_DSQ with SCX_ENQ_IMMED.
- * With multiple workers competing for CPU 0, dsq->nr > 1 triggers the
- * IMMED slow path (reenqueue with SCX_TASK_REENQ_IMMED).
+ * ops.dispatch() on CPU 0 consumes multiple tasks from USER_DSQ with
+ * SCX_ENQ_IMMED in a single dispatch call, causing dsq->nr to exceed 1
+ * and triggering the IMMED slow path (reenqueue with SCX_TASK_REENQ_IMMED).
  *
  * Requires scx_bpf_dsq_move_to_local___v2() (v7.1+) for enq_flags support.
  */
@@ -55,10 +55,23 @@ void BPF_STRUCT_OPS(consume_immed_enqueue, struct task_struct *p,
 
 void BPF_STRUCT_OPS(consume_immed_dispatch, s32 cpu, struct task_struct *prev)
 {
-	if (cpu == 0)
-		scx_bpf_dsq_move_to_local(USER_DSQ, SCX_ENQ_IMMED);
-	else
+	int i;
+
+	if (cpu != 0) {
 		scx_bpf_dsq_move_to_local(SCX_DSQ_GLOBAL, 0);
+		return;
+	}
+
+	/*
+	 * Move multiple tasks into CPU 0's local DSQ with SCX_ENQ_IMMED in a
+	 * single dispatch call. When the second task is inserted (dsq->nr > 1),
+	 * dsq_inc_nr() triggers the IMMED slow path via schedule_reenq_local(),
+	 * which calls ops.enqueue() with SCX_ENQ_REENQ | SCX_TASK_REENQ_IMMED.
+	 */
+	for (i = 0; i < 4; i++) {
+		if (!scx_bpf_dsq_move_to_local(USER_DSQ, SCX_ENQ_IMMED))
+			break;
+	}
 }
 
 s32 BPF_STRUCT_OPS_SLEEPABLE(consume_immed_init)
diff --git a/tools/testing/selftests/sched_ext/consume_immed.c b/tools/testing/selftests/sched_ext/consume_immed.c
index 7f9594cfa9cb..61cbc0fe3663 100644
--- a/tools/testing/selftests/sched_ext/consume_immed.c
+++ b/tools/testing/selftests/sched_ext/consume_immed.c
@@ -12,18 +12,30 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <pthread.h>
+#include <sched.h>
 #include <bpf/bpf.h>
 #include <scx/common.h>
 #include "consume_immed.bpf.skel.h"
 #include "scx_test.h"
 
-#define NUM_WORKERS		4
+/*
+ * Use more workers than CPUs, all pinned to CPU 0, so CPU 0's local DSQ
+ * accumulates multiple IMMED tasks at once, reliably triggering the slow path.
+ */
+#define NUM_WORKERS		8
 #define TEST_DURATION_SEC	3
 
 static volatile bool stop_workers;
 
 static void *worker_fn(void *arg)
 {
+	cpu_set_t cpuset;
+
+	/* Pin to CPU 0 to saturate its local DSQ */
+	CPU_ZERO(&cpuset);
+	CPU_SET(0, &cpuset);
+	pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+
 	while (!stop_workers) {
 		volatile unsigned long i;
 
-- 
2.43.0