[PATCH] sched_ext: Fix sched_ext_dead() race with scx_root_enable_workfn()

[zhidao su <soolaugust@gmail.com>]

In CONFIG_EXT_SUB_SCHED, scx_task_sched(p) returns p->scx.sched instead
of scx_root.  scx_root_enable_workfn() initializes tasks in two steps:

  scx_init_task(sch, p, false)   /* state=INIT, p->scx.sched still NULL */
  scx_set_task_sched(p, sch)     /* p->scx.sched = sch                  */

Between these two steps, a concurrent sched_ext_dead() can call
scx_task_sched(p), get NULL, and pass it to scx_disable_and_exit_task()
which crashes in SCX_HAS_OP(NULL, ...).

In sched_ext_dead(), skip scx_disable_and_exit_task() when state=INIT and
p->scx.sched is still NULL, and reset state to NONE.  In
scx_root_enable_workfn(), after scx_init_task() returns, skip
scx_set_task_sched() if sched_ext_dead() has already removed @p from
scx_tasks.

Fixes: 073d4f0667b0 ("sched_ext: Refactor task init/exit helpers")
Signed-off-by: zhidao su <suzhidao@xiaomi.com>
---
 kernel/sched/ext.c | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index f7b1b16e81a5..1872af65b103 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -3887,9 +3887,21 @@ void sched_ext_dead(struct task_struct *p)
 		struct rq_flags rf;
 		struct rq *rq;
 
-		rq = task_rq_lock(p, &rf);
-		scx_disable_and_exit_task(scx_task_sched(p), p);
-		task_rq_unlock(rq, p, &rf);
+		/*
+		 * scx_root_enable_workfn() may be concurrently initializing @p.
+		 * scx_init_task() sets state=INIT before scx_set_task_sched()
+		 * sets p->scx.sched.  If we race that window, p->scx.sched is
+		 * still NULL; skip scx_disable_and_exit_task() and reset state
+		 * to NONE so @p leaves SCX cleanly.
+		 */
+		if (scx_get_task_state(p) == SCX_TASK_INIT &&
+		    !rcu_access_pointer(p->scx.sched)) {
+			scx_set_task_state(p, SCX_TASK_NONE);
+		} else {
+			rq = task_rq_lock(p, &rf);
+			scx_disable_and_exit_task(scx_task_sched(p), p);
+			task_rq_unlock(rq, p, &rf);
+		}
 	}
 }
 
@@ -6937,6 +6949,17 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 			goto err_disable_unlock_all;
 		}
 
+		/*
+		 * sched_ext_dead() may have already cleaned up @p while locks
+		 * were dropped in scx_task_iter_unlock(); skip it.
+		 */
+		scoped_guard(raw_spinlock_irq, &scx_tasks_lock) {
+			if (list_empty(&p->scx.tasks_node)) {
+				put_task_struct(p);
+				continue;
+			}
+		}
+
 		scx_set_task_sched(p, sch);
 		scx_set_task_state(p, SCX_TASK_READY);
 
-- 
2.43.0