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

[Kuba Piecuch <jpiecuch@google.com>]

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