Merge tag 'sched_ext-for-7.2-rc6-fixes' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext

Pull sched_ext fixes from Tejun Heo:

 - More lifecycle fixes for the new sub-scheduler support: a failed
   enable could tear down a never-linked sub-scheduler in a way that
   races the root scheduler's disable and leads to a use-after-free,
   tasks that were not on the ext class could still get the enable
   callback, and a policy-rejection path silently rewrote a running
   task's scheduling policy instead of aborting the scheduler.

 - Scheduler enable/disable could deadlock with cgroup removal and a
   concurrent cgroup weight write through kernfs. Fixed by reordering
   lock acquisition.

 - Sync wakeups could leave the waker CPU incorrectly marked idle in the
   built-in idle-CPU tracking.

 - A selftest fix for sleeping tasks whose CPU affinity changes before
   wakeup.

* tag 'sched_ext-for-7.2-rc6-fixes' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext:
  selftests/sched_ext: Handle sleeping task affinity changes in numa test
  sched_ext: Mark waker CPU busy when selected in WAKE_SYNC case
  sched_ext: Don't enable non-ext tasks in the sub-sched task loops
  sched_ext: Skip sub-disable teardown for never-linked sub-schedulers
  sched_ext: Take cgroup_lock() first in scx_cgroup_lock()
  sched_ext: Reject setting disallow from init_task outside the enable path

1 files changed, 12 insertions, 1 deletions

diff --git a/tools/testing/selftests/sched_ext/numa.bpf.c b/tools/testing/selftests/sched_ext/numa.bpf.c
index 78cc49a7f9a6..6b4515c28aa0 100644
--- a/tools/testing/selftests/sched_ext/numa.bpf.c
+++ b/tools/testing/selftests/sched_ext/numa.bpf.c
@@ -34,7 +34,8 @@ static bool is_cpu_idle(s32 cpu, int node)
 s32 BPF_STRUCT_OPS(numa_select_cpu,
 		   struct task_struct *p, s32 prev_cpu, u64 wake_flags)
 {
-	int node = __COMPAT_scx_bpf_cpu_node(scx_bpf_task_cpu(p));
+	s32 task_cpu = scx_bpf_task_cpu(p);
+	int node = __COMPAT_scx_bpf_cpu_node(task_cpu);
 	s32 cpu;
 
 	/*
@@ -48,6 +49,16 @@ s32 BPF_STRUCT_OPS(numa_select_cpu,
 		cpu = __COMPAT_scx_bpf_pick_any_cpu_node(p->cpus_ptr, node,
 						__COMPAT_SCX_PICK_IDLE_IN_NODE);
 
+	/*
+	 * @task_cpu may be outside of p->cpus_ptr if @p's affinity
+	 * changed while it was sleeping. This means it's possible for
+	 * p->cpus_ptr to not include any CPUs from @node.
+	 * If we failed to find a cpu in @node, check if @task_cpu
+	 * is outside of p->cpus_ptr and just return @prev_cpu if it is.
+	 */
+	if (cpu < 0 && !bpf_cpumask_test_cpu(task_cpu, p->cpus_ptr))
+		return prev_cpu;
+
 	if (is_cpu_idle(cpu, node))
 		scx_bpf_error("CPU %d should be marked as busy", cpu);