[PATCH sched_ext/for-7.3-fixes] sched_ext: Fix spurious aborts in scx_bpf_dsq_move() on ownership change races

[Tejun Heo <tj@kernel.org>]

scx_dsq_move() verifies that the task belongs to the calling scheduler
before taking any locks and aborts the scheduler on mismatch. The task can
lose the sched association at any point: It can run and fully exit, which
clears the association, or get rehomed to a different sub-sched. Both are
benign races, but the early ownership check escalates them into scheduler
aborts.

Move the ownership check below the cursor-lost check. Every ownership change
dequeues the task first, so a task that is still on the iterated DSQ under
the lock while owned elsewhere indicates a genuine violation and should
abort.

Also fix two stale comments still referencing sched_ext_free(), which has
been renamed to sched_ext_dead().

Fixes: bb4d9fd55158 ("sched_ext: scx_dsq_move() should validate the task belongs to the right scheduler")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c |   21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -7694,7 +7694,7 @@ static void scx_root_enable_workfn(struc
 	/*
 	 * Enable ops for every task. Fork is excluded by scx_fork_rwsem
 	 * preventing new tasks from being added. No need to exclude tasks
-	 * leaving as sched_ext_free() can handle both prepped and enabled
+	 * leaving as sched_ext_dead() can handle both prepped and enabled
 	 * tasks. Prep all tasks first and then enable them with preemption
 	 * disabled.
 	 *
@@ -7786,7 +7786,7 @@ static void scx_root_enable_workfn(struc
 
 	/*
 	 * We're fully committed and can't fail. The task READY -> ENABLED
-	 * transitions here are synchronized against sched_ext_free() through
+	 * transitions here are synchronized against sched_ext_dead() through
 	 * scx_tasks_lock.
 	 */
 	percpu_down_write(&scx_fork_rwsem);
@@ -9004,12 +9004,6 @@ static bool scx_dsq_move(struct bpf_iter
 	if (unlikely(READ_ONCE(sch->aborting)))
 		return false;
 
-	if (unlikely(!scx_task_on_sched(sch, p))) {
-		scx_error(sch, "scx_bpf_dsq_move[_vtime]() on %s[%d] but the task belongs to a different scheduler",
-			  p->comm, p->pid);
-		return false;
-	}
-
 	/*
 	 * Can be called from either ops.dispatch() holding the dispatched rq's
 	 * lock or any context where no rq lock is held. If latter, lock @p's
@@ -9040,6 +9034,17 @@ static bool scx_dsq_move(struct bpf_iter
 		raw_spin_unlock(&src_dsq->lock);
 		goto out;
 	}
+
+	/*
+	 * @p has been on $src_dsq and can't move anymore. If @p is not on @sch,
+	 * the caller didn't have authority over @p at the time of the call.
+	 */
+	if (unlikely(!scx_task_on_sched(sch, p))) {
+		scx_error(sch, "scx_bpf_dsq_move[_vtime]() on %s[%d] but the task belongs to a different scheduler",
+			  p->comm, p->pid);
+		raw_spin_unlock(&src_dsq->lock);
+		goto out;
+	}
 
 	/* @p is still on $src_dsq and stable, determine the destination */
 	dst_dsq = find_dsq_for_dispatch(sch, locked_rq ?: this_rq(), dsq_id, task_cpu(p));