Re: [PATCH v2] sched_ext: idle: Respect task affinity in restricted CPU selection

[Andrea Righi <arighi@nvidia.com>]

Hi Hui,

On Sun, Oct 11, 2026 at 10:17:00PM +0900, Hui Su wrote:
> The restricted select_cpu kfunc path handles single-CPU and
> migration-disabled tasks without going through scx_select_cpu_dfl().
> 
> It currently uses prev_cpu as the candidate. However, prev_cpu is only a
> hint and may no longer be usable after an affinity change. When the caller
> supplies @allowed, the restricted path can also claim and return an idle
> CPU which isn't in p->cpus_ptr.
> 
> Derive the candidate from task state instead. Use the sole CPU in
> p->cpus_ptr for single-CPU tasks and task_cpu(p) for migration-disabled
> tasks. Validate the candidate against the task's effective affinity and
> the caller-provided mask before claiming its idle state.
> 
> The validation must happen before scx_idle_test_and_clear_cpu(), as a
> successful call consumes the CPU's idle bit. A later placement validation
> therefore cannot undo the idle-state side effect.
> 
> Single-CPU and migration-disabled tasks skip ops.select_cpu(), so this
> path can instead be reached when the CPU-selection kfunc is called from
> other contexts such as ops.enqueue(). The normal ops.select_cpu() path is
> unchanged.
> 
> Fixes: 683d2d0faba1 ("sched_ext: idle: Introduce scx_bpf_select_cpu_and()")
> Cc: stable@vger.kernel.org
> Signed-off-by: Hui Su <sh_def@163.com>
> ---
> Testing:
> - x86_64 KVM/QEMU private BPF test-run harness:
>   - The baseline returned CPU1 outside the task affinity; v2 returned
>     -EBUSY.
>   - With a stale prev_cpu, the baseline returned -EBUSY; v2 selected the
>     task's idle CPU0.
>   - A caller mask excluding the task's candidate CPU returned -EBUSY with v2.
> - The stock sched_ext allowed_cpus selftest passed.
> - Built vmlinux and bzImage with CONFIG_PREEMPT_RCU=y.
> - Built kernel/sched/build_policy.o with CONFIG_PREEMPT_RCU=n and CONFIG_SMP=n.
> 
> Changes in v2:
> - Derive the restricted-path candidate from task state instead of prev_cpu.
> - Document the idle-bit side effect and the ops.enqueue() reachability.
> 
> v1: https://lore.kernel.org/lkml/20261011025256.4180374-1-sh_def@163.com/
> ---
>  kernel/sched/ext/idle.c | 20 ++++++++++++++------
>  1 file changed, 14 insertions(+), 6 deletions(-)
> 
> diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
> index aa9fb6de0ad6..1ccea60e9012 100644
> --- a/kernel/sched/ext/idle.c
> +++ b/kernel/sched/ext/idle.c
> @@ -958,6 +958,7 @@ static s32 select_cpu_from_kfunc(struct scx_sched *sch, struct task_struct *p,
>  {
>  	unsigned long irq_flags;
>  	bool we_locked = false;
> +	bool migration_disabled;
>  	s32 cpu;
>  
>  	if (!scx_cpu_valid(sch, prev_cpu, NULL))
> @@ -995,14 +996,21 @@ static s32 select_cpu_from_kfunc(struct scx_sched *sch, struct task_struct *p,
>  	/*
>  	 * This may also be called from ops.enqueue(), so we need to handle
>  	 * per-CPU tasks as well. For these tasks, we can skip all idle CPU
> -	 * selection optimizations and simply check whether the previously
> -	 * used CPU is idle and within the allowed cpumask.
> +	 * selection optimizations and simply check whether the task's only
> +	 * usable CPU is idle and within the allowed cpumask.
>  	 */
> -	if (p->nr_cpus_allowed == 1 || is_bpf_migration_disabled(p)) {
> -		if (cpumask_test_cpu(prev_cpu, allowed ?: p->cpus_ptr) &&
> -		    scx_idle_test_and_clear_cpu(prev_cpu))
> -			cpu = prev_cpu;
> +	migration_disabled = is_bpf_migration_disabled(p);
> +
> +	if (p->nr_cpus_allowed == 1 || migration_disabled) {
> +		if (migration_disabled)
> +			cpu = task_cpu(p);
>  		else
> +			cpu = cpumask_first(p->cpus_ptr);
> +
> +		if ((migration_disabled &&
> +		     !cpumask_test_cpu(cpu, p->cpus_ptr)) ||
> +		    (allowed && !cpumask_test_cpu(cpu, allowed)) ||
> +		    !scx_idle_test_and_clear_cpu(cpu))
>  			cpu = -EBUSY;

Sorry, I suggested this, but I just realized that we break SCX_PICK_IDLE_IN_NODE
if we do this. What if we factor the check with an helper? Something like this:

static s32 select_cpu_restricted(struct task_struct *p, s32 prev_cpu,
				 const struct cpumask *allowed, u64 flags,
				 bool migration_disabled)
{
	s32 cpu;

	if (migration_disabled)
		cpu = task_cpu(p);
	else
		cpu = cpumask_first(p->cpus_ptr);

	if (!cpumask_test_cpu(cpu, p->cpus_ptr))
		return -EBUSY;

	if (allowed && !cpumask_test_cpu(cpu, allowed))
		return -EBUSY;

	if (flags & SCX_PICK_IDLE_IN_NODE) {
		int node = scx_cpu_node_if_enabled(prev_cpu);

		if (scx_cpu_node_if_enabled(cpu) != node)
			return -EBUSY;
	}

	if (!scx_idle_test_and_clear_cpu(cpu))
		return -EBUSY;

	return cpu;
}

Then the caller becomes:

	migration_disabled = is_bpf_migration_disabled(p);

	if (p->nr_cpus_allowed == 1 || migration_disabled)
		cpu = select_cpu_restricted(p, prev_cpu, allowed, flags, migration_disabled);
	else
		cpu = scx_select_cpu_dfl(p, prev_cpu, wake_flags, allowed ?: p->cpus_ptr, flags);

Thanks,
-Andrea