[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

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

[Kuba Piecuch]

Hi Andrea,

On Sun Jul 26, 2026 at 6:47 AM UTC, Andrea Righi wrote:
> 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)

I don't think it matters much in terms of correctness, but to me it would
be more intuitive to read refresh_idle_masks first to ensure we're bootstrapped,
and then check the idle bit.

> +		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());

ops.running() doesn't have to run on the same CPU as @p, e.g. when changing
the priority of a task running on a remote CPU. I believe the correct thing
to do here is scx_bpf_test_and_clear_cpu_idle(scx_bpf_task_cpu(p)).

> +}
> +
>  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;
> +}
> +

This bootstrapping mechanism feels like a bit of a hack.
Couldn't we improve SCX itself to ensure the initial state of the idle masks
is accurate?

I was thinking we could enhance scx_idle_enable() by making it enable idle
tracking (currently idle tracking is controlled by the __scx_enabled static
branch), and then iterating over all CPUs, locking their rq locks and setting
their idle bit based on whether rq->curr == rq->idle. All this would happen
before calling ops.init(), so the BPF scheduler will be guaranteed to have an
accurate idle cpumask. WDYT?

>  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);
> +

Won't a WRITE_ONCE() suffice here? test_and_clear_bit() implies a full memory
barrier, so I don't think we need any extra synchronization once the read of
refresh_idle_masks is moved before scx_bpf_test_and_clear_cpu_idle() in
validate_local_idle_state().

>  	/* Pick an idle CPU from user-space */
>  	SCX_FAIL_IF(test_select_cpu_from_user(skel), "Failed to pick idle CPU");
>  

Thanks,
Kuba

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

[Andrea Righi]

Hi Kuba,

On Thu, Jul 30, 2026 at 03:19:08PM +0000, Kuba Piecuch wrote:
> Hi Andrea,
> 
> On Sun Jul 26, 2026 at 6:47 AM UTC, Andrea Righi wrote:
> > 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)
> 
> I don't think it matters much in terms of correctness, but to me it would
> be more intuitive to read refresh_idle_masks first to ensure we're bootstrapped,
> and then check the idle bit.

Ack. And since I read your other comments below, we can remove this condition
entirely if we move the idle-mask initialization before ops.init() in the SCX
core.

> 
> > +		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());
> 
> ops.running() doesn't have to run on the same CPU as @p, e.g. when changing
> the priority of a task running on a remote CPU. I believe the correct thing
> to do here is scx_bpf_test_and_clear_cpu_idle(scx_bpf_task_cpu(p)).

Ah yes, that's a mistake, we should definitely use scx_bpf_task_cpu(p).

> 
> > +}
> > +
> >  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;
> > +}
> > +
> 
> This bootstrapping mechanism feels like a bit of a hack.
> Couldn't we improve SCX itself to ensure the initial state of the idle masks
> is accurate?
> 
> I was thinking we could enhance scx_idle_enable() by making it enable idle
> tracking (currently idle tracking is controlled by the __scx_enabled static
> branch), and then iterating over all CPUs, locking their rq locks and setting
> their idle bit based on whether rq->curr == rq->idle. All this would happen
> before calling ops.init(), so the BPF scheduler will be guaranteed to have an
> accurate idle cpumask. WDYT?

Agreed, this is much cleaner. I'll send v2 as a two-patch series and move the
initialization into SCX.

> 
> >  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);
> > +
> 
> Won't a WRITE_ONCE() suffice here? test_and_clear_bit() implies a full memory
> barrier, so I don't think we need any extra synchronization once the read of
> refresh_idle_masks is moved before scx_bpf_test_and_clear_cpu_idle() in
> validate_local_idle_state().

Yes, WRITE_ONCE() should be sufficient for the current workload.

> 
> >  	/* Pick an idle CPU from user-space */
> >  	SCX_FAIL_IF(test_select_cpu_from_user(skel), "Failed to pick idle CPU");
> >  

Thanks!
-Andrea