[PATCHSET v4 sched_ext/for-7.3] sched_ext: Fix idle CPU state initialization and validation

[PATCHSET v4 sched_ext/for-7.3] sched_ext: Fix idle CPU state initialization and validation

[Andrea Righi]

The built-in idle masks are initialized with all online CPUs marked idle, but
idle tracking currently starts only after sched_ext is fully enabled. This
leaves busy CPUs incorrectly advertised as idle during ops.init() and until
their next idle transition.

Moreover, the allowed_cpus selftest checks that a remotely selected CPU is no
longer present in the idle mask. An idle-to-idle re-pick can re-advertise the
CPU before the test performs this check, making the validation inherently racy.

Fix this by initializing the built-in idle masks empty. This makes the initial
state conservative, so busy CPUs are never incorrectly advertised as idle. Once
the scheduler is enabled and bypass is lifted, the existing idle re-notification
path reschedules every online CPU. Idle-to-idle re-picks then populate the masks
with CPUs that are actually idle.

Also rework the allowed_cpus kselftest to replace the racy remote-CPU check with
a stable local invariant: a CPU running a non-idle scheduling context must not
be advertised as idle.

Changes in v4:
 - Initialize the built-in idle masks empty instead of enabling idle tracking
   before ops.init() and refreshing every CPU under its rq lock (Tejun Heo)
 - Rely on the existing bypass idle re-notification path to populate the masks
   with CPUs that are actually idle
 - Link to v3: https://lore.kernel.org/all/20260731182406.3166853-1-arighi@nvidia.com/

Changes in v3:
 - Reuse the built-in idle-selection static key instead of introducing a
   dedicated idle-tracking key (Kuba Piecuch)
 - Check the local CPU-idle invariant from both ops.select_cpu() and
   ops.enqueue() (Kuba Piecuch)
 - Inspect the idle mask without modifying it and perform the check before
   calling scx_bpf_select_cpu_and() (Kuba Piecuch)
 - Link to v2: https://lore.kernel.org/all/20260731090334.2911948-1-arighi@nvidia.com/

Changes in v2:
 - Move idle-mask initialization from the selftest into the sched_ext
   core (Kuba Piecuch)
 - Add a dedicated idle-tracking static key so transitions are tracked
   before the scheduler is fully enabled (Kuba Piecuch)
 - Rework the allowed_cpus selftest to validate the stable local CPU-idle
   invariant
 - Link to v1: https://lore.kernel.org/all/20260726064754.378671-1-arighi@nvidia.com/

Andrea Righi (2):
      sched_ext: Initialize idle masks as busy
      selftests/sched_ext: Make allowed_cpus idle validation race-free

 kernel/sched/ext/idle.c                            | 16 +++----
 .../testing/selftests/sched_ext/allowed_cpus.bpf.c | 49 ++++++++++++++++++----
 2 files changed, 50 insertions(+), 15 deletions(-)

[PATCH 1/2] sched_ext: Initialize idle masks as busy

[Andrea Righi]

The built-in idle masks are reset with all online CPUs marked idle
before sched_ext is enabled. Busy CPUs can therefore be incorrectly
advertised as idle until their next idle transition.

Initialize the masks empty so that the initial state is conservative.
When bypass is lifted, every CPU is rescheduled and idle-to-idle
re-picks populate the masks with CPUs that are actually idle. Later
idle transitions keep the masks up to date.

Cc: Kuba Piecuch <jpiecuch@google.com>
Suggested-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/idle.c | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 3e9d6a44bf431..33bd51ab3cede 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -838,20 +838,20 @@ static void reset_idle_masks(struct sched_ext_ops *ops)
 	int node;
 
 	/*
-	 * Consider all online cpus idle. Should converge to the actual state
-	 * quickly.
+	 * Start with all CPUs marked busy. The idle masks are populated when
+	 * bypass is lifted and each idle CPU is forced through an idle re-pick.
+	 * This may temporarily omit idle CPUs but never advertises a busy CPU as
+	 * idle.
 	 */
 	if (!(ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)) {
-		cpumask_copy(idle_cpumask(NUMA_NO_NODE)->cpu, cpu_online_mask);
-		cpumask_copy(idle_cpumask(NUMA_NO_NODE)->smt, cpu_online_mask);
+		cpumask_clear(idle_cpumask(NUMA_NO_NODE)->cpu);
+		cpumask_clear(idle_cpumask(NUMA_NO_NODE)->smt);
 		return;
 	}
 
 	for_each_node(node) {
-		const struct cpumask *node_mask = cpumask_of_node(node);
-
-		cpumask_and(idle_cpumask(node)->cpu, cpu_online_mask, node_mask);
-		cpumask_and(idle_cpumask(node)->smt, cpu_online_mask, node_mask);
+		cpumask_clear(idle_cpumask(node)->cpu);
+		cpumask_clear(idle_cpumask(node)->smt);
 	}
 }
 
-- 
2.55.0

[PATCH 2/2] 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 a stable local invariant instead: a CPU executing
ops.select_cpu() or ops.enqueue() in a non-idle scheduling context must
not be advertised as idle. Read the idle mask without modifying it and
also validate selected CPUs against the requested domain and task
affinity.

Suggested-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 49 ++++++++++++++++---
 1 file changed, 42 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..9dd72d0da29b2 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -15,15 +15,48 @@ UEI_DEFINE(uei);
 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
 
 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);
+	const struct cpumask *idle;
+	struct task_struct *curr;
+	s32 cpu = bpf_get_smp_processor_id();
+	bool cpu_is_idle, 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();
+
+	idle = scx_bpf_get_idle_cpumask();
+	cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle);
+	scx_bpf_put_idle_cpumask(idle);
+
+	/*
+	 * 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 && cpu_is_idle)
+		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,
@@ -32,6 +65,7 @@ s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
 	const struct cpumask *allowed;
 	s32 cpu;
 
+	validate_local_idle_state();
 	allowed = cast_mask(allowed_cpumask);
 	if (!allowed) {
 		scx_bpf_error("allowed domain not initialized");
@@ -43,7 +77,7 @@ s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
 	 */
 	cpu = scx_bpf_select_cpu_and(p, prev_cpu, wake_flags, allowed, 0);
 	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;
@@ -59,6 +93,7 @@ void BPF_STRUCT_OPS(allowed_cpus_enqueue, struct task_struct *p, u64 enq_flags)
 
 	scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
 
+	validate_local_idle_state();
 	allowed = cast_mask(allowed_cpumask);
 	if (!allowed) {
 		scx_bpf_error("allowed domain not initialized");
@@ -71,7 +106,7 @@ 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);
 	}
 }
-- 
2.55.0

Re: [PATCHSET v4 sched_ext/for-7.3] sched_ext: Fix idle CPU state initialization and validation

[Kuba Piecuch]

On Mon Aug 3, 2026 at 6:08 AM UTC, Andrea Righi wrote:
> The built-in idle masks are initialized with all online CPUs marked idle, but
> idle tracking currently starts only after sched_ext is fully enabled. This
> leaves busy CPUs incorrectly advertised as idle during ops.init() and until
> their next idle transition.
>
> Moreover, the allowed_cpus selftest checks that a remotely selected CPU is no
> longer present in the idle mask. An idle-to-idle re-pick can re-advertise the
> CPU before the test performs this check, making the validation inherently racy.
>
> Fix this by initializing the built-in idle masks empty. This makes the initial
> state conservative, so busy CPUs are never incorrectly advertised as idle. Once
> the scheduler is enabled and bypass is lifted, the existing idle re-notification
> path reschedules every online CPU. Idle-to-idle re-picks then populate the masks
> with CPUs that are actually idle.
>
> Also rework the allowed_cpus kselftest to replace the racy remote-CPU check with
> a stable local invariant: a CPU running a non-idle scheduling context must not
> be advertised as idle.
>
> Changes in v4:
>  - Initialize the built-in idle masks empty instead of enabling idle tracking
>    before ops.init() and refreshing every CPU under its rq lock (Tejun Heo)
>  - Rely on the existing bypass idle re-notification path to populate the masks
>    with CPUs that are actually idle
>  - Link to v3: https://lore.kernel.org/all/20260731182406.3166853-1-arighi@nvidia.com/
>
> Changes in v3:
>  - Reuse the built-in idle-selection static key instead of introducing a
>    dedicated idle-tracking key (Kuba Piecuch)
>  - Check the local CPU-idle invariant from both ops.select_cpu() and
>    ops.enqueue() (Kuba Piecuch)
>  - Inspect the idle mask without modifying it and perform the check before
>    calling scx_bpf_select_cpu_and() (Kuba Piecuch)
>  - Link to v2: https://lore.kernel.org/all/20260731090334.2911948-1-arighi@nvidia.com/
>
> Changes in v2:
>  - Move idle-mask initialization from the selftest into the sched_ext
>    core (Kuba Piecuch)
>  - Add a dedicated idle-tracking static key so transitions are tracked
>    before the scheduler is fully enabled (Kuba Piecuch)
>  - Rework the allowed_cpus selftest to validate the stable local CPU-idle
>    invariant
>  - Link to v1: https://lore.kernel.org/all/20260726064754.378671-1-arighi@nvidia.com/
>
> Andrea Righi (2):
>       sched_ext: Initialize idle masks as busy
>       selftests/sched_ext: Make allowed_cpus idle validation race-free
>
>  kernel/sched/ext/idle.c                            | 16 +++----
>  .../testing/selftests/sched_ext/allowed_cpus.bpf.c | 49 ++++++++++++++++++----
>  2 files changed, 50 insertions(+), 15 deletions(-)

Thanks Andrea, looks good to me!

For the entire series:

Reviewed-by: Kuba Piecuch <jpiecuch@google.com>

Re: [PATCHSET v4 sched_ext/for-7.3] sched_ext: Fix idle CPU state initialization and validation

[Tejun Heo]

Hello,

> Andrea Righi (2):
>       sched_ext: Initialize idle masks as busy
>       selftests/sched_ext: Make allowed_cpus idle validation race-free

Applied 1-2 to sched_ext/for-7.3 with Kuba's Reviewed-by added.

Thanks.

--
tejun

Re: [PATCH 2/2] selftests/sched_ext: Make allowed_cpus idle validation race-free

[Kuba Piecuch]

On Fri Jul 31, 2026 at 3:26 PM UTC, Andrea Righi wrote:
> On Fri, Jul 31, 2026 at 11:12:09AM +0000, Kuba Piecuch wrote:
>> Could we check a stronger invariant by also checking that the bit in the idle
>> mask is set if we're running an idle task? We can get the idle cpumask through
>> scx_bpf_get_idle_cpumask() and check bits without clearing them using
>> bpf_cpumask_test_cpu().
>
> I don't think the other direction always holds: an idle CPU can be claimed by
> another BPF idle CPU selection helper, which can clear the idle bit before the
> CPU necessarily stops running the idle task. In that case, observing an idle
> task with a clear idle bit is legitimate.

Right, thanks for pointing that out.

[PATCH 2/2] 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 a stable local invariant instead: a CPU executing
ops.select_cpu() or ops.enqueue() in a non-idle scheduling context must
not be advertised as idle. Read the idle mask without modifying it and
also validate selected CPUs against the requested domain and task
affinity.

Suggested-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 49 ++++++++++++++++---
 1 file changed, 42 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..9dd72d0da29b2 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -15,15 +15,48 @@ UEI_DEFINE(uei);
 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
 
 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);
+	const struct cpumask *idle;
+	struct task_struct *curr;
+	s32 cpu = bpf_get_smp_processor_id();
+	bool cpu_is_idle, 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();
+
+	idle = scx_bpf_get_idle_cpumask();
+	cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle);
+	scx_bpf_put_idle_cpumask(idle);
+
+	/*
+	 * 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 && cpu_is_idle)
+		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,
@@ -32,6 +65,7 @@ s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
 	const struct cpumask *allowed;
 	s32 cpu;
 
+	validate_local_idle_state();
 	allowed = cast_mask(allowed_cpumask);
 	if (!allowed) {
 		scx_bpf_error("allowed domain not initialized");
@@ -43,7 +77,7 @@ s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
 	 */
 	cpu = scx_bpf_select_cpu_and(p, prev_cpu, wake_flags, allowed, 0);
 	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;
@@ -59,6 +93,7 @@ void BPF_STRUCT_OPS(allowed_cpus_enqueue, struct task_struct *p, u64 enq_flags)
 
 	scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
 
+	validate_local_idle_state();
 	allowed = cast_mask(allowed_cpumask);
 	if (!allowed) {
 		scx_bpf_error("allowed domain not initialized");
@@ -71,7 +106,7 @@ 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);
 	}
 }
-- 
2.55.0

Re: [PATCH 2/2] selftests/sched_ext: Make allowed_cpus idle validation race-free

[Andrea Righi]

On Fri, Jul 31, 2026 at 11:12:09AM +0000, Kuba Piecuch wrote:
> Hi Andrea,
> 
> On Fri Jul 31, 2026 at 8:59 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.
> 
> That invariant sounds like it should hold in many contexts, not just in
> ops.select_cpu(). Is there something preventing us from checking it in
> ops.enqueue() as well?

Yes, nothing prevents checking the local invariant from ops.enqueue() as well.
The CPU running the callback shouldn't be advertised as idle. I'll add this in
v3.

> 
> > Also validate both the requested domain and task affinity for selected
> > CPUs.
> >
> > Signed-off-by: Andrea Righi <arighi@nvidia.com>
> > ---
> >  .../selftests/sched_ext/allowed_cpus.bpf.c    | 43 ++++++++++++++++---
> >  1 file changed, 36 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..411a7edcb9605 100644
> > --- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> > +++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> > @@ -15,15 +15,43 @@ UEI_DEFINE(uei);
> >  private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
> >  
> >  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))
> > +		scx_bpf_error("running CPU %d should be marked as busy", cpu);
> 
> Could we check a stronger invariant by also checking that the bit in the idle
> mask is set if we're running an idle task? We can get the idle cpumask through
> scx_bpf_get_idle_cpumask() and check bits without clearing them using
> bpf_cpumask_test_cpu().

I don't think the other direction always holds: an idle CPU can be claimed by
another BPF idle CPU selection helper, which can clear the idle bit before the
CPU necessarily stops running the idle task. In that case, observing an idle
task with a clear idle bit is legitimate.

However, we can safely use scx_bpf_get_idle_cpumask() to perform the existing
check without modifying the mask.

Thanks,
-Andrea

Re: [PATCH 2/2] selftests/sched_ext: Make allowed_cpus idle validation race-free

[Kuba Piecuch]

Hi Andrea,

On Fri Jul 31, 2026 at 8:59 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.

That invariant sounds like it should hold in many contexts, not just in
ops.select_cpu(). Is there something preventing us from checking it in
ops.enqueue() as well?

> Also validate both the requested domain and task affinity for selected
> CPUs.
>
> Signed-off-by: Andrea Righi <arighi@nvidia.com>
> ---
>  .../selftests/sched_ext/allowed_cpus.bpf.c    | 43 ++++++++++++++++---
>  1 file changed, 36 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..411a7edcb9605 100644
> --- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> +++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
> @@ -15,15 +15,43 @@ UEI_DEFINE(uei);
>  private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
>  
>  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))
> +		scx_bpf_error("running CPU %d should be marked as busy", cpu);

Could we check a stronger invariant by also checking that the bit in the idle
mask is set if we're running an idle task? We can get the idle cpumask through
scx_bpf_get_idle_cpumask() and check bits without clearing them using
bpf_cpumask_test_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 +70,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,7 +100,7 @@ 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);
>  	}
>  }

Thanks,
Kuba

[PATCH 2/2] 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.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 43 ++++++++++++++++---
 1 file changed, 36 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..411a7edcb9605 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -15,15 +15,43 @@ UEI_DEFINE(uei);
 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
 
 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))
+		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 +70,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,7 +100,7 @@ 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);
 	}
 }
-- 
2.55.0