[PATCH sched_ext/for-7.3] selftests/sched_ext: Make numa idle validation race-free

[Cheng-Yang Chou <yphbchou0911@gmail.com>]

A CPU returned by scx_bpf_pick_idle_cpu_node() can be re-advertised as
idle by an idle-to-idle re-pick before the BPF program validates the
selection, and the scx_bpf_pick_any_cpu_node() fallback doesn't claim
the CPU at all. Asserting that the picked CPU is absent from the node's
idle cpumask is therefore inherently racy.

Follow the same approach as commit 12da4723b679 ("selftests/sched_ext:
Make allowed_cpus idle validation race-free") and validate a stable
local invariant instead: a CPU executing ops.select_cpu() in a non-idle
scheduling context must not be advertised as idle in its node's idle
cpumask. Keep the node-membership validation of the picked CPU, which
is stable.

Signed-off-by: Cheng-Yang Chou <yphbchou0911@gmail.com>
---
 tools/testing/selftests/sched_ext/numa.bpf.c | 28 +++++++++++++++-----
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/numa.bpf.c b/tools/testing/selftests/sched_ext/numa.bpf.c
index 6b4515c28aa0..679b51d38089 100644
--- a/tools/testing/selftests/sched_ext/numa.bpf.c
+++ b/tools/testing/selftests/sched_ext/numa.bpf.c
@@ -19,16 +19,31 @@ UEI_DEFINE(uei);
 
 const volatile unsigned int __COMPAT_SCX_PICK_IDLE_IN_NODE;
 
-static bool is_cpu_idle(s32 cpu, int node)
+static void validate_local_idle_state(void)
 {
 	const struct cpumask *idle_cpumask;
-	bool idle;
+	struct task_struct *curr;
+	s32 cpu = bpf_get_smp_processor_id();
+	int node = __COMPAT_scx_bpf_cpu_node(cpu);
+	bool cpu_is_idle, curr_is_idle;
+
+	bpf_rcu_read_lock();
+	curr = scx_bpf_cpu_curr(cpu);
+	curr_is_idle = curr && (curr->flags & PF_IDLE);
+	bpf_rcu_read_unlock();
 
 	idle_cpumask = __COMPAT_scx_bpf_get_idle_cpumask_node(node);
-	idle = bpf_cpumask_test_cpu(cpu, idle_cpumask);
+	cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle_cpumask);
 	scx_bpf_put_cpumask(idle_cpumask);
 
-	return idle;
+	/*
+	 * Unlike a remote picked CPU, the local CPU cannot go through an
+	 * idle re-pick while this callback is running. If it is running a
+	 * non-idle scheduling context, it must not be advertised as idle
+	 * in its node's idle cpumask.
+	 */
+	if (!curr_is_idle && cpu_is_idle)
+		scx_bpf_error("running CPU %d should be marked as busy", cpu);
 }
 
 s32 BPF_STRUCT_OPS(numa_select_cpu,
@@ -38,6 +53,8 @@ s32 BPF_STRUCT_OPS(numa_select_cpu,
 	int node = __COMPAT_scx_bpf_cpu_node(task_cpu);
 	s32 cpu;
 
+	validate_local_idle_state();
+
 	/*
 	 * We could just use __COMPAT_scx_bpf_pick_any_cpu_node() here,
 	 * since it already tries to pick an idle CPU within the node
@@ -59,9 +76,6 @@ s32 BPF_STRUCT_OPS(numa_select_cpu,
 	if (cpu < 0 && !bpf_cpumask_test_cpu(task_cpu, p->cpus_ptr))
 		return prev_cpu;
 
-	if (is_cpu_idle(cpu, node))
-		scx_bpf_error("CPU %d should be marked as busy", cpu);
-
 	if (__COMPAT_scx_bpf_cpu_node(cpu) != node)
 		scx_bpf_error("CPU %d should be in node %d", cpu, node);
 
-- 
2.43.0