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

[PATCHSET v2 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.

Make the built-in idle CPU masks reliable from ops.init() onward by enabling
idle tracking early and synchronizing every online CPU's state under its rq
lock. During this early phase, only the built-in masks are updated;
ops.update_idle() notifications remain suppressed until the scheduler is fully
enabled.

Also rework the allowed_cpus kselftest to replace the racy remote-CPU check with
a stable local invariant that relies on the newly guaranteed initial idle-mask
state.

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
 - Target sched_ext/for-7.3 instead of sched_ext/for-7.2-fixes
 - Link to v1: https://lore.kernel.org/all/20260726064754.378671-1-arighi@nvidia.com/

Andrea Righi (2):
      sched_ext: Initialize idle masks before ops.init()
      selftests/sched_ext: Make allowed_cpus idle validation race-free

 kernel/sched/ext/ext.h                             |  4 +-
 kernel/sched/ext/idle.c                            | 38 ++++++++++++++++++-
 .../testing/selftests/sched_ext/allowed_cpus.bpf.c | 43 ++++++++++++++++++----
 3 files changed, 75 insertions(+), 10 deletions(-)

[PATCH 1/2] sched_ext: Initialize idle masks before ops.init()

[Andrea Righi]

The built-in idle masks are reset with all online CPUs marked idle, but
idle state tracking starts only after the scheduler is fully enabled.
As a result, ops.init() can observe busy CPUs as idle, and those CPUs
remain incorrectly advertised until their next idle transition.

Enable idle tracking before ops.init() and refresh every online CPU under
its rq lock. Once a CPU is refreshed, later transitions keep its state
accurate. Keep ops.update_idle() notifications disabled until the
scheduler is fully enabled.

Suggested-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.h  |  4 +++-
 kernel/sched/ext/idle.c | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext/ext.h b/kernel/sched/ext/ext.h
index 0b7fc46aee08c..495e7023c9ee1 100644
--- a/kernel/sched/ext/ext.h
+++ b/kernel/sched/ext/ext.h
@@ -59,11 +59,13 @@ static inline void init_sched_ext_class(void) {}
 #endif	/* CONFIG_SCHED_CLASS_EXT */
 
 #ifdef CONFIG_SCHED_CLASS_EXT
+DECLARE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
+
 void __scx_update_idle(struct rq *rq, bool idle, bool do_notify);
 
 static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 {
-	if (scx_enabled())
+	if (static_branch_unlikely(&scx_idle_tracking_enabled))
 		__scx_update_idle(rq, idle, do_notify);
 }
 #else
diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 3e9d6a44bf431..6d81bb7c43966 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -14,6 +14,9 @@
 #include "idle.h"
 #include "sub.h"
 
+/* Enable/disable idle state tracking */
+DEFINE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
+
 /* Enable/disable built-in idle CPU selection policy */
 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
 
@@ -810,6 +813,15 @@ void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 	if (static_branch_likely(&scx_builtin_idle_enabled))
 		update_builtin_idle(cpu, idle);
 
+	/*
+	 * Idle tracking starts before the scheduler is enabled so that the
+	 * built-in idle masks are accurate when ops.init() runs. Suppress
+	 * ops.update_idle() notifications until the scheduler is fully
+	 * enabled.
+	 */
+	if (!scx_enabled())
+		return;
+
 	/*
 	 * ops.update_idle() fires on real idle transitions, indicated by
 	 * @do_notify and managed by put_prev_task_idle()/set_next_task_idle().
@@ -838,8 +850,8 @@ static void reset_idle_masks(struct sched_ext_ops *ops)
 	int node;
 
 	/*
-	 * Consider all online cpus idle. Should converge to the actual state
-	 * quickly.
+	 * Seed all online CPUs as idle. refresh_idle_masks() below corrects
+	 * their state before ops.init() runs.
 	 */
 	if (!(ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)) {
 		cpumask_copy(idle_cpumask(NUMA_NO_NODE)->cpu, cpu_online_mask);
@@ -855,6 +867,23 @@ static void reset_idle_masks(struct sched_ext_ops *ops)
 	}
 }
 
+static void refresh_idle_masks(void)
+{
+	int cpu;
+
+	/*
+	 * Idle tracking is already enabled and the online CPU set is stable.
+	 * Once a CPU is refreshed under its rq lock, subsequent transitions
+	 * keep its state up to date.
+	 */
+	for_each_online_cpu(cpu) {
+		struct rq *rq = cpu_rq(cpu);
+
+		scoped_guard(rq_lock_irqsave, rq)
+			update_builtin_idle(cpu, rq->curr == rq->idle);
+	}
+}
+
 void scx_idle_enable(struct sched_ext_ops *ops)
 {
 	if (!ops->update_idle || (ops->flags & SCX_OPS_KEEP_BUILTIN_IDLE))
@@ -868,10 +897,15 @@ void scx_idle_enable(struct sched_ext_ops *ops)
 		static_branch_disable_cpuslocked(&scx_builtin_idle_per_node);
 
 	reset_idle_masks(ops);
+	static_branch_enable_cpuslocked(&scx_idle_tracking_enabled);
+
+	if (static_branch_likely(&scx_builtin_idle_enabled))
+		refresh_idle_masks();
 }
 
 void scx_idle_disable(void)
 {
+	static_branch_disable(&scx_idle_tracking_enabled);
 	static_branch_disable(&scx_builtin_idle_enabled);
 	static_branch_disable(&scx_builtin_idle_per_node);
 }
-- 
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 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

Re: [PATCH 1/2] sched_ext: Initialize idle masks before ops.init()

[Kuba Piecuch]

Hi Andrea,

On Fri Jul 31, 2026 at 8:59 AM UTC, Andrea Righi wrote:
> @@ -59,11 +59,13 @@ static inline void init_sched_ext_class(void) {}
>  #endif	/* CONFIG_SCHED_CLASS_EXT */
>  
>  #ifdef CONFIG_SCHED_CLASS_EXT
> +DECLARE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
> +

I was originally thinking about reusing scx_builtin_idle_enabled here,
apologies if I wasn't clear enough.

My reasoning is: If the user is doing their own idle CPU tracking,
in which case scx_builtin_idle_enable will be disabled, what's the point
of SCX tracking idle CPUs?

Do you see a scenario where using one static branch is problematic?

Thanks,
Kuba

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

Re: [PATCH 1/2] sched_ext: Initialize idle masks before ops.init()

[Andrea Righi]

Hi Kuba,

On Fri, Jul 31, 2026 at 10:47:30AM +0000, Kuba Piecuch wrote:
> Hi Andrea,
> 
> On Fri Jul 31, 2026 at 8:59 AM UTC, Andrea Righi wrote:
> > @@ -59,11 +59,13 @@ static inline void init_sched_ext_class(void) {}
> >  #endif	/* CONFIG_SCHED_CLASS_EXT */
> >  
> >  #ifdef CONFIG_SCHED_CLASS_EXT
> > +DECLARE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
> > +
> 
> I was originally thinking about reusing scx_builtin_idle_enabled here,
> apologies if I wasn't clear enough.
> 
> My reasoning is: If the user is doing their own idle CPU tracking,
> in which case scx_builtin_idle_enable will be disabled, what's the point
> of SCX tracking idle CPUs?
> 
> Do you see a scenario where using one static branch is problematic?

I think you're right, I added a separate key because __scx_update_idle() also
delivers ops.update_idle() callbacks, but we can retain the existing
scx_enabled() check for that case and use scx_builtin_idle_enabled only for
early tracking. So I don't see any reason to add a separate static key.

I'll rework this in v3.

Thanks,
-Andrea

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

[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