[PATCH 4/4] selftests/sched_ext: consume_immed: fix reliability with CPU affinity

[zhidao su <soolaugust@gmail.com>]

Two bugs prevented the IMMED slow path from triggering reliably:

1. Dispatch consumed at most one task per ops.dispatch() call.  The IMMED
   slow path requires dsq->nr > 1 in the local DSQ after a
   scx_bpf_dsq_move_to_local() call.  With only one call per dispatch,
   the local DSQ never accumulated two tasks.

   Fix: loop up to 8 times in dispatch, draining USER_DSQ into the local
   DSQ.  On the second iteration the local DSQ already has one task;
   inserting a second raises dsq->nr to 2, triggering schedule_reenq_local
   and calling ops.enqueue() with SCX_TASK_REENQ_IMMED.

2. Workers were not pinned to CPU 0, so the scheduler spread them across
   all CPUs.  USER_DSQ typically had only 0-1 tasks when CPU 0 dispatched,
   so the loop's second iteration rarely found a second task.

   Fix: pin all NUM_WORKERS threads to CPU 0 via
   pthread_attr_setaffinity_np().  With all workers competing for CPU 0,
   USER_DSQ always has a backlog, and every dispatch loop reliably
   accumulates 2+ tasks.

Also move nr_consume_immed_reenq read to after bpf_link__destroy() so
the kernel has fully flushed all counter updates before we check the
assertion.

Fixes: c50dcf533149 ("selftests/sched_ext: Add tests for SCX_ENQ_IMMED and scx_bpf_dsq_reenq()")
Signed-off-by: zhidao su <suzhidao@xiaomi.com>
---
 .../selftests/sched_ext/consume_immed.bpf.c   | 19 +++++++++++++++----
 .../selftests/sched_ext/consume_immed.c       | 15 ++++++++++++---
 2 files changed, 27 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..a2a4ee2ee95b 100644
--- a/tools/testing/selftests/sched_ext/consume_immed.bpf.c
+++ b/tools/testing/selftests/sched_ext/consume_immed.bpf.c
@@ -55,10 +55,21 @@ 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
-		scx_bpf_dsq_move_to_local(SCX_DSQ_GLOBAL, 0);
+	int i;
+
+	if (cpu != 0)
+		return;
+
+	/*
+	 * Drain USER_DSQ into the local DSQ with SCX_ENQ_IMMED. Once two or
+	 * more tasks accumulate in the local DSQ, dsq->nr > 1 triggers the
+	 * IMMED slow path (schedule_reenq_local), re-enqueuing IMMED tasks
+	 * through ops.enqueue() with SCX_TASK_REENQ_IMMED.
+	 */
+	for (i = 0; i < 8; 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..41e66cd5e879 100644
--- a/tools/testing/selftests/sched_ext/consume_immed.c
+++ b/tools/testing/selftests/sched_ext/consume_immed.c
@@ -12,6 +12,7 @@
 #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"
@@ -77,20 +78,28 @@ static enum scx_test_status run(void *ctx)
 
 	stop_workers = false;
 	for (i = 0; i < NUM_WORKERS; i++) {
-		SCX_FAIL_IF(pthread_create(&workers[i], NULL, worker_fn, NULL),
+		pthread_attr_t attr;
+		cpu_set_t cpuset;
+
+		pthread_attr_init(&attr);
+		CPU_ZERO(&cpuset);
+		CPU_SET(0, &cpuset);
+		pthread_attr_setaffinity_np(&attr, sizeof(cpuset), &cpuset);
+		SCX_FAIL_IF(pthread_create(&workers[i], &attr, worker_fn, NULL),
 			    "Failed to create worker %d", i);
+		pthread_attr_destroy(&attr);
 	}
 
 	sleep(TEST_DURATION_SEC);
 
-	reenq = skel->bss->nr_consume_immed_reenq;
-
 	stop_workers = true;
 	for (i = 0; i < NUM_WORKERS; i++)
 		pthread_join(workers[i], NULL);
 
 	bpf_link__destroy(link);
 
+	reenq = skel->bss->nr_consume_immed_reenq;
+
 	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_UNREG));
 	SCX_GT(reenq, 0);
 
-- 
2.43.0