[PATCH v2] 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. 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@xxxxxxxxxxxxxxx
Signed-off-by: Hui Su <sh_def@xxxxxxx>
---
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@xxxxxxx/
---
 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;
 	} else {
 		cpu = scx_select_cpu_dfl(p, prev_cpu, wake_flags,

base-commit: 7b63ef2d55f24519e7e9e5f4d15dbea03f126e40
-- 
2.55.0