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

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

[Hui Su]

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, even though prev_cpu is only a
hint and may no longer be usable after an affinity change. With a
caller-provided @allowed mask, this can also return an idle CPU outside
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 it against p->cpus_ptr and @allowed before claiming its
idle state.

Since the derived candidate may be on a different NUMA node from
prev_cpu, also preserve SCX_PICK_IDLE_IN_NODE by rejecting cross-node
candidates.

Perform all checks before scx_idle_test_and_clear_cpu(), as a successful
call consumes the CPU's idle bit.

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>
---
Tested on x86_64 KVM/QEMU:
- Private BPF harness passed affinity, stale-prev_cpu, caller-mask, and
  cross-NUMA SCX_PICK_IDLE_IN_NODE cases; v2 returned CPU2 across nodes,
  while v3 returned -EBUSY.
- Build coverage passed with CONFIG_PREEMPT_RCU=y/n and CONFIG_SMP=n.

Changes in v3:
- Factor restricted selection into select_cpu_restricted().
- Preserve SCX_PICK_IDLE_IN_NODE and add cross-NUMA validation.

v2: https://lore.kernel.org/lkml/20261011131700.1351541-1-sh_def@163.com/

 kernel/sched/ext/idle.c | 50 ++++++++++++++++++++++++++++++++---------
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index aa9fb6de0ad6..28ed9e327085 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -952,12 +952,43 @@ static bool is_bpf_migration_disabled(const struct task_struct *p)
 	return p->migration_disabled;
 }
 
+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;
+}
+
 static s32 select_cpu_from_kfunc(struct scx_sched *sch, struct task_struct *p,
 				 s32 prev_cpu, u64 wake_flags,
 				 const struct cpumask *allowed, u64 flags)
 {
 	unsigned long irq_flags;
 	bool we_locked = false;
+	bool migration_disabled;
 	s32 cpu;
 
 	if (!scx_cpu_valid(sch, prev_cpu, NULL))
@@ -994,20 +1025,17 @@ 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.
+	 * per-CPU tasks as well. For these tasks, skip general idle CPU
+	 * selection and validate the task's only usable CPU directly.
 	 */
-	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;
-		else
-			cpu = -EBUSY;
-	} else {
+	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);
-	}
 
 	if (we_locked)
 		raw_spin_unlock_irqrestore(&p->pi_lock, irq_flags);

base-commit: 7b63ef2d55f24519e7e9e5f4d15dbea03f126e40

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

[Andrea Righi]

Hi Hui,

On Mon, Oct 12, 2026 at 02:31:00AM +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, even though prev_cpu is only a
> hint and may no longer be usable after an affinity change. With a
> caller-provided @allowed mask, this can also return an idle CPU outside
> 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 it against p->cpus_ptr and @allowed before claiming its
> idle state.
> 
> Since the derived candidate may be on a different NUMA node from
> prev_cpu, also preserve SCX_PICK_IDLE_IN_NODE by rejecting cross-node
> candidates.
> 
> Perform all checks before scx_idle_test_and_clear_cpu(), as a successful
> call consumes the CPU's idle bit.
 
> 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>

Reviewed-by: Andrea Righi <arighi@nvidia.com>

Thanks,
-Andrea

> ---
> Tested on x86_64 KVM/QEMU:
> - Private BPF harness passed affinity, stale-prev_cpu, caller-mask, and
>   cross-NUMA SCX_PICK_IDLE_IN_NODE cases; v2 returned CPU2 across nodes,
>   while v3 returned -EBUSY.
> - Build coverage passed with CONFIG_PREEMPT_RCU=y/n and CONFIG_SMP=n.
> 
> Changes in v3:
> - Factor restricted selection into select_cpu_restricted().
> - Preserve SCX_PICK_IDLE_IN_NODE and add cross-NUMA validation.
> 
> v2: https://lore.kernel.org/lkml/20261011131700.1351541-1-sh_def@163.com/
> 
>  kernel/sched/ext/idle.c | 50 ++++++++++++++++++++++++++++++++---------
>  1 file changed, 39 insertions(+), 11 deletions(-)
> 
> diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
> index aa9fb6de0ad6..28ed9e327085 100644
> --- a/kernel/sched/ext/idle.c
> +++ b/kernel/sched/ext/idle.c
> @@ -952,12 +952,43 @@ static bool is_bpf_migration_disabled(const struct task_struct *p)
>  	return p->migration_disabled;
>  }
>  
> +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;
> +}
> +
>  static s32 select_cpu_from_kfunc(struct scx_sched *sch, struct task_struct *p,
>  				 s32 prev_cpu, u64 wake_flags,
>  				 const struct cpumask *allowed, u64 flags)
>  {
>  	unsigned long irq_flags;
>  	bool we_locked = false;
> +	bool migration_disabled;
>  	s32 cpu;
>  
>  	if (!scx_cpu_valid(sch, prev_cpu, NULL))
> @@ -994,20 +1025,17 @@ 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.
> +	 * per-CPU tasks as well. For these tasks, skip general idle CPU
> +	 * selection and validate the task's only usable CPU directly.
>  	 */
> -	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;
> -		else
> -			cpu = -EBUSY;
> -	} else {
> +	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);
> -	}
>  
>  	if (we_locked)
>  		raw_spin_unlock_irqrestore(&p->pi_lock, irq_flags);
> 
> base-commit: 7b63ef2d55f24519e7e9e5f4d15dbea03f126e40
> 
>