[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

[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