[PATCH sched_ext/for-7.2-fixes] selftests/sched_ext: Make allowed_cpus idle validation race-free

[PATCH sched_ext/for-7.2-fixes] selftests/sched_ext: Make allowed_cpus idle validation race-free

[Andrea Righi]

A remotely selected CPU can be re-advertised as idle by an idle-to-idle
re-pick before the BPF program validates the selection. Checking that
the selected CPU remains absent from the idle mask is therefore
inherently racy.

Validate the stable local invariant instead: a CPU running a non-idle
scheduling context in ops.select_cpu() must not be advertised as idle.
Also validate both the requested domain and task affinity for selected
CPUs.

Moreover, bootstrap the test by running a task on every active CPU while
ops.running() refreshes the initial idle state. This ensures that the
idle masks are properly initialized before strict validation begins.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 51 ++++++++++++++++---
 .../selftests/sched_ext/allowed_cpus.c        | 38 ++++++++++++++
 2 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
index 35923e74a2ec3..4a14b05065453 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -13,17 +13,46 @@ char _license[] SEC("license") = "GPL";
 UEI_DEFINE(uei);
 
 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
+volatile bool refresh_idle_masks;
 
 static void
-validate_idle_cpu(const struct task_struct *p, const struct cpumask *allowed, s32 cpu)
+validate_local_idle_state(void)
 {
-	if (scx_bpf_test_and_clear_cpu_idle(cpu))
-		scx_bpf_error("CPU %d should be marked as busy", cpu);
+	struct task_struct *curr;
+	s32 cpu = bpf_get_smp_processor_id();
+	bool curr_is_idle;
 
-	if (bpf_cpumask_subset(allowed, p->cpus_ptr) &&
-	    !bpf_cpumask_test_cpu(cpu, allowed))
+	bpf_rcu_read_lock();
+	curr = scx_bpf_cpu_curr(cpu);
+	curr_is_idle = curr && (curr->flags & PF_IDLE);
+	bpf_rcu_read_unlock();
+
+	/*
+	 * Unlike a remote selected 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.
+	 */
+	if (!curr_is_idle && scx_bpf_test_and_clear_cpu_idle(cpu) && !refresh_idle_masks)
+		scx_bpf_error("running CPU %d should be marked as busy", cpu);
+}
+
+static void
+validate_selected_cpu(const struct task_struct *p, s32 cpu)
+{
+	const struct cpumask *allowed = cast_mask(allowed_cpumask);
+
+	if (!allowed) {
+		scx_bpf_error("allowed domain not initialized");
+		return;
+	}
+
+	if (!bpf_cpumask_test_cpu(cpu, allowed))
 		scx_bpf_error("CPU %d not in the allowed domain for %d (%s)",
 			      cpu, p->pid, p->comm);
+
+	if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr))
+		scx_bpf_error("CPU %d not in the affinity mask for %d (%s)",
+			      cpu, p->pid, p->comm);
 }
 
 s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
@@ -42,8 +71,9 @@ s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
 	 * Select an idle CPU strictly within the allowed domain.
 	 */
 	cpu = scx_bpf_select_cpu_and(p, prev_cpu, wake_flags, allowed, 0);
+	validate_local_idle_state();
 	if (cpu >= 0) {
-		validate_idle_cpu(p, allowed, cpu);
+		validate_selected_cpu(p, cpu);
 		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
 
 		return cpu;
@@ -71,11 +101,17 @@ void BPF_STRUCT_OPS(allowed_cpus_enqueue, struct task_struct *p, u64 enq_flags)
 	 */
 	cpu = scx_bpf_select_cpu_and(p, prev_cpu, 0, allowed, 0);
 	if (cpu >= 0) {
-		validate_idle_cpu(p, allowed, cpu);
+		validate_selected_cpu(p, cpu);
 		scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE);
 	}
 }
 
+void BPF_STRUCT_OPS(allowed_cpus_running, struct task_struct *p)
+{
+	if (refresh_idle_masks)
+		scx_bpf_test_and_clear_cpu_idle(bpf_get_smp_processor_id());
+}
+
 s32 BPF_STRUCT_OPS_SLEEPABLE(allowed_cpus_init)
 {
 	struct bpf_cpumask *mask;
@@ -138,6 +174,7 @@ SEC(".struct_ops.link")
 struct sched_ext_ops allowed_cpus_ops = {
 	.select_cpu		= (void *)allowed_cpus_select_cpu,
 	.enqueue		= (void *)allowed_cpus_enqueue,
+	.running		= (void *)allowed_cpus_running,
 	.init			= (void *)allowed_cpus_init,
 	.exit			= (void *)allowed_cpus_exit,
 	.name			= "allowed_cpus",
diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.c b/tools/testing/selftests/sched_ext/allowed_cpus.c
index 093f285ab4bae..eb1708e55982b 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.c
@@ -3,6 +3,7 @@
  * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
  */
 #include <bpf/bpf.h>
+#include <sched.h>
 #include <scx/common.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -47,14 +48,51 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
 	return 0;
 }
 
+/*
+ * Run this task once on every CPU while ops.running() repairs the bootstrap
+ * idle state. Once a CPU has been refreshed, subsequent idle transitions keep
+ * its state up to date.
+ */
+static int refresh_idle_masks(void)
+{
+	cpu_set_t original, one;
+	int cpu, ret = 0;
+
+	if (sched_getaffinity(0, sizeof(original), &original))
+		return -errno;
+
+	for (cpu = 0; cpu < CPU_SETSIZE; cpu++) {
+		if (!CPU_ISSET(cpu, &original))
+			continue;
+
+		CPU_ZERO(&one);
+		CPU_SET(cpu, &one);
+		if (sched_setaffinity(0, sizeof(one), &one)) {
+			ret = -errno;
+			break;
+		}
+
+		sched_yield();
+	}
+
+	if (sched_setaffinity(0, sizeof(original), &original) && !ret)
+		ret = -errno;
+
+	return ret;
+}
+
 static enum scx_test_status run(void *ctx)
 {
 	struct allowed_cpus *skel = ctx;
 	struct bpf_link *link;
 
+	skel->bss->refresh_idle_masks = true;
 	link = bpf_map__attach_struct_ops(skel->maps.allowed_cpus_ops);
 	SCX_FAIL_IF(!link, "Failed to attach scheduler");
 
+	SCX_FAIL_IF(refresh_idle_masks(), "Failed to refresh idle CPU state");
+	__atomic_store_n(&skel->bss->refresh_idle_masks, false, __ATOMIC_RELEASE);
+
 	/* Pick an idle CPU from user-space */
 	SCX_FAIL_IF(test_select_cpu_from_user(skel), "Failed to pick idle CPU");
 
-- 
2.55.0

Re: [PATCH sched_ext/for-7.2-fixes] selftests/sched_ext: Make allowed_cpus idle validation race-free

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential issue(s) to consider:
- [Medium] The selftest uses a fixed-size `cpu_set_t` for `sched_getaffinity()`, which fails with `-EINVAL` on systems with a large number of configured CPUs (e.g., `nr_cpu_ids > 1024`), causing the test to fail unconditionally.
- [Medium] The test's bootstrap phase incorrectly assumes full access to all system CPUs, leading to false positive test failures when run in restricted environments (e.g., taskset or cgroups).
--

--- Patch [1]: [PATCH sched_ext/for-7.2-fixes] selftests/sched_ext: Make allowed_cpus idle validation race-free ---
commit 2ede7dc703cde549d449386a7878e3392244c530
Author: Andrea Righi <arighi@nvidia.com>

selftests/sched_ext: Make allowed_cpus idle validation race-free

This patch updates the allowed_cpus sched_ext selftest to make idle CPU
validation race-free. It shifts validation to the stable local invariant in
ops.select_cpu() and adds a bootstrap phase to refresh the idle masks.

> diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> index 35923e74a2ec3..4a14b05065453 100644
> --- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> +++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
[ ... ]
> @@ -13,17 +13,46 @@ char _license[] SEC("license") = "GPL";
>  UEI_DEFINE(uei);
>  
>  private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
> +volatile bool refresh_idle_masks;
>  
>  static void
> -validate_idle_cpu(const struct task_struct *p, const struct cpumask *allowed, s32 cpu)
> +validate_local_idle_state(void)
>  {
> -	if (scx_bpf_test_and_clear_cpu_idle(cpu))
> -		scx_bpf_error("CPU %d should be marked as busy", cpu);
> +	struct task_struct *curr;
> +	s32 cpu = bpf_get_smp_processor_id();
> +	bool curr_is_idle;
>  
> -	if (bpf_cpumask_subset(allowed, p->cpus_ptr) &&
> -	    !bpf_cpumask_test_cpu(cpu, allowed))
> +	bpf_rcu_read_lock();
> +	curr = scx_bpf_cpu_curr(cpu);
> +	curr_is_idle = curr && (curr->flags & PF_IDLE);
> +	bpf_rcu_read_unlock();
> +
> +	/*
> +	 * Unlike a remote selected 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.
> +	 */
> +	if (!curr_is_idle && scx_bpf_test_and_clear_cpu_idle(cpu) && !refresh_idle_masks)
> +		scx_bpf_error("running CPU %d should be marked as busy", cpu);

[Severity: Medium]
Could this validation fail if it runs on CPUs that userspace skips during the
bootstrap phase?

Because the BPF scheduler is attached host-wide, background tasks waking up on
the skipped CPUs could trigger ops.select_cpu(), executing this check. If
these CPUs were not refreshed because they weren't in the initial affinity
mask, their state might be stale and trigger a false positive error here.

> +}
[ ... ]
> diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.c b/tools/testing/selftests/sched_ext/allowed_cpus.c
> index 093f285ab4bae..eb1708e55982b 100644
> --- a/tools/testing/selftests/sched_ext/allowed_cpus.c
> +++ b/tools/testing/selftests/sched_ext/allowed_cpus.c
[ ... ]
> @@ -47,14 +48,51 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
>  	return 0;
>  }
>  
> +/*
> + * Run this task once on every CPU while ops.running() repairs the bootstrap
> + * idle state. Once a CPU has been refreshed, subsequent idle transitions keep
> + * its state up to date.
> + */
> +static int refresh_idle_masks(void)
> +{
> +	cpu_set_t original, one;
> +	int cpu, ret = 0;
> +
> +	if (sched_getaffinity(0, sizeof(original), &original))
> +		return -errno;

[Severity: Medium]
Will this fail unconditionally on systems with a large number of CPUs?

The sched_getaffinity() syscall explicitly checks if the provided buffer length
in bits is less than nr_cpu_ids. Because cpu_set_t is typically fixed at 1024
bits, a system configured with nr_cpu_ids > 1024 would return -EINVAL here and
abort the test suite.

> +
> +	for (cpu = 0; cpu < CPU_SETSIZE; cpu++) {
> +		if (!CPU_ISSET(cpu, &original))
> +			continue;

[Severity: Medium]
Does skipping CPUs outside the initial affinity mask leave their idle state out
of sync?

In environments that use cgroups or tasksets to restrict the test's CPU
affinity, these skipped CPUs won't have their idle masks refreshed, which
could lead to the false positive test failures in validate_local_idle_state()
mentioned above.

> +
> +		CPU_ZERO(&one);
> +		CPU_SET(cpu, &one);
> +		if (sched_setaffinity(0, sizeof(one), &one)) {
> +			ret = -errno;
> +			break;
> +		}
> +
> +		sched_yield();
> +	}
[ ... ]

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260726064754.378671-1-arighi@nvidia.com?part=1