[PATCH 3/4] sched_ext: Skip sub-disable teardown for never-linked sub-schedulers

[Tejun Heo]

A sub-scheduler enable can fail before scx_link_sched() links the sched into
the hierarchy, e.g. when the parent is already being disabled, and cleanup
still runs the full scx_sub_disable().

That is racy against root disable: drain_descendants() is the only ordering
between a sub's disable-time task walk and root disable's all-task teardown,
and an unlinked sub is invisible to it. Root's teardown can thus run between
the never-linked sub's drain and its walk, exiting every task to no
scheduler.

The walk then trips the membership WARN and re-homes the exited tasks onto
the dying hierarchy, a use-after-free.

Skip the cgroup ownership reset and the task walk if @sch was never linked,
indicated by the empty ->sibling as unlinking only happens later in the same
function. The membership WARN remains valid: a linked sub is always waited
on by an ancestor's drain.

Fixes: 337ec00b1d9c ("sched_ext: Implement cgroup sub-sched enabling and disabling")
Signed-off-by: Tejun Heo <tj@xxxxxxxxxx>
---
 kernel/sched/ext/ext.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 58898cd0727b..98dd7df88db4 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -5937,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);
@@ -5949,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));
 
@@ -6017,6 +6026,7 @@ static void scx_sub_disable(struct scx_sched *sch)
 	}
 	scx_task_iter_stop(&sti);
 
+dump:
 	scx_disable_dump(sch);
 
 	scx_cgroup_unlock();
-- 
2.55.0