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

4 files changed, 58 insertions, 16 deletions

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index 20b2343aa344..87e353f7e011 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -244,11 +244,11 @@ struct sched_ext_entity {
 	 * to %SCHED_EXT with -%EACCES.
 	 *
 	 * Can be set from ops.init_task() while the BPF scheduler is being
-	 * loaded (!scx_init_task_args->fork). If set and the task's policy is
-	 * already %SCHED_EXT, the task's policy is rejected and forcefully
-	 * reverted to %SCHED_NORMAL. The number of such events are reported
-	 * through /sys/kernel/debug/sched_ext::nr_rejected. Setting this flag
-	 * during fork is not allowed.
+	 * loaded. If set and the task's policy is already %SCHED_EXT, the
+	 * task's policy is rejected and forcefully reverted to %SCHED_NORMAL.
+	 * The number of such events are reported through
+	 * /sys/kernel/sched_ext/nr_rejected. Setting this flag from any other
+	 * ops.init_task() invocation, such as during fork, fails the scheduler.
 	 */
 	bool			disallow;	/* reject switching into SCX */
 
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index e3fa7b2fac9d..18183062f751 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3606,6 +3606,9 @@ static int __scx_init_task(struct scx_sched *sch, struct task_struct *p, bool fo
 		} else if (unlikely(fork)) {
 			scx_error(sch, "ops.init_task() set task->scx.disallow for %s[%d] during fork",
 				  p->comm, p->pid);
+		} else if (unlikely(scx_enable_state() != SCX_ENABLING)) {
+			scx_error(sch, "ops.init_task() set task->scx.disallow for %s[%d] outside the enable path",
+				  p->comm, p->pid);
 		} else {
 			struct rq *rq;
 			struct rq_flags rf;
@@ -4577,20 +4580,25 @@ static struct cgroup *root_cgroup(void)
 	return &cgrp_dfl_root.cgrp;
 }
 
+/*
+ * cgroup_lock() must nest outside the rwsem write side: a writer waiting
+ * for cgroup_mutex deadlocks with cgroup teardown, which holds it while
+ * draining a set_* file write blocked on the rwsem behind the writer.
+ */
 static void scx_cgroup_lock(void)
 {
+	cgroup_lock();
 #ifdef CONFIG_EXT_GROUP_SCHED
 	percpu_down_write(&scx_cgroup_ops_rwsem);
 #endif
-	cgroup_lock();
 }
 
 static void scx_cgroup_unlock(void)
 {
-	cgroup_unlock();
 #ifdef CONFIG_EXT_GROUP_SCHED
 	percpu_up_write(&scx_cgroup_ops_rwsem);
 #endif
+	cgroup_unlock();
 }
 #else	/* CONFIG_EXT_GROUP_SCHED || CONFIG_EXT_SUB_SCHED */
 static inline struct cgroup *root_cgroup(void) { return NULL; }
@@ -5929,6 +5937,15 @@ static void scx_sub_disable(struct scx_sched *sch)
 	percpu_down_write(&scx_fork_rwsem);
 	scx_cgroup_lock();
 
+	/*
+	 * An enable that failed before scx_link_sched() never owned a cgroup or
+	 * task and won't be waited on by an ancestor's drain_descendants().
+	 * Nothing to reparent and walking the tasks can misbehave as the task
+	 * ownership invariant (either owned by self or parent) does not hold.
+	 */
+	if (list_empty(&sch->sibling))
+		goto dump;
+
 	set_cgroup_sched(sch_cgroup(sch), parent);
 
 	scx_task_iter_start(&sti, sch->cgrp);
@@ -5941,8 +5958,8 @@ static void scx_sub_disable(struct scx_sched *sch)
 			continue;
 
 		/*
-		 * By the time control reaches here, all descendant schedulers
-		 * should already have been disabled.
+		 * By the time control reaches here, all linked descendant
+		 * schedulers should have been disabled.
 		 */
 		WARN_ON_ONCE(!scx_task_on_sched(sch, p));
 
@@ -5993,15 +6010,22 @@ static void scx_sub_disable(struct scx_sched *sch)
 			/*
 			 * $p is initialized for $parent and still attached to
 			 * @sch. Disable and exit for @sch, switch over to
-			 * $parent, override the state to READY to account for
-			 * $p having already been initialized, and then enable.
+			 * $parent and override the state to READY to account
+			 * for $p having already been initialized.
 			 */
 			scx_disable_and_exit_task(sch, p);
 			scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
 			scx_set_task_state(p, SCX_TASK_INIT);
 			scx_set_task_sched(p, parent);
 			scx_set_task_state(p, SCX_TASK_READY);
-			scx_enable_task(parent, p);
+
+			/*
+			 * A task on a non-ext class, possible under an
+			 * %SCX_OPS_SWITCH_PARTIAL root, stays READY and is
+			 * enabled by switching_to_scx() if it switches over.
+			 */
+			if (p->sched_class == &ext_sched_class)
+				scx_enable_task(parent, p);
 		}
 
 		task_rq_unlock(rq, p, &rf);
@@ -6009,6 +6033,7 @@ static void scx_sub_disable(struct scx_sched *sch)
 	}
 	scx_task_iter_stop(&sti);
 
+dump:
 	scx_disable_dump(sch);
 
 	scx_cgroup_unlock();
@@ -7708,10 +7733,14 @@ static void scx_sub_enable_workfn(struct kthread_work *work)
 
 			/*
 			 * $p is now only initialized for @sch and READY, which
-			 * is what we want. Assign it to @sch and enable.
+			 * is what we want. Assign it to @sch and, if it's on
+			 * the ext class, enable. A non-ext task, possible under
+			 * an %SCX_OPS_SWITCH_PARTIAL root, stays READY and is
+			 * enabled by switching_to_scx() if it switches over.
 			 */
 			scx_set_task_sched(p, sch);
-			scx_enable_task(sch, p);
+			if (p->sched_class == &ext_sched_class)
+				scx_enable_task(sch, p);
 
 			p->scx.flags &= ~SCX_TASK_SUB_INIT;
 		}
diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 8e8c6201b7df..6f93cc32b650 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -554,8 +554,10 @@ s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
 		    cpu_rq(cpu)->scx.local_dsq.nr == 0 &&
 		    (!(flags & SCX_PICK_IDLE_IN_NODE) || (waker_node == node)) &&
 		    !cpumask_empty(idle_cpumask(waker_node)->cpu)) {
-			if (cpumask_test_cpu(cpu, allowed))
+			if (cpumask_test_cpu(cpu, allowed)) {
+				scx_idle_test_and_clear_cpu(cpu);
 				goto out_unlock;
+			}
 		}
 	}
 
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);