[PATCH] sched_ext: don't rehome a dead task in scx_cgroup_task_migrated

[Tao Cui]

From: Tao Cui <cuitao@xxxxxxxxxx>

A task can exit between cgroup migration commit and the MIGRATED callback:
sched_ext_dead() marks it SCX_TASK_DEAD before cgroup_task_dead() removes it
from the migration list, so scx_cgroup_task_migrated() can pick up a dead
task and call scx_rehome_task(), which re-enables it and leaks the BPF
scheduler's per-task resources. The other scx_rehome_task() callers already
check for this; do the same here.

Fixes: bf9dee58ab56 ("sched_ext: Re-home tasks on cgroup migration")
Signed-off-by: Tao Cui <cuitao@xxxxxxxxxx>
---
 kernel/sched/ext/sub.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index b81254be1b04..8f8d21b85abe 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -2083,6 +2083,13 @@ static void scx_cgroup_task_migrated(struct cgroup_task_migrate_ctx *ctx)
 		return;
 
 	rq = task_rq_lock(p, &rf);
+
+	if (scx_get_task_state(p) == SCX_TASK_DEAD) {
+		/* sched_ext_dead() raced us */
+		task_rq_unlock(rq, p, &rf);
+		return;
+	}
+
 	scx_rehome_task(to, p);
 	task_rq_unlock(rq, p, &rf);
 }
-- 
2.43.0