[PATCH 05/11] sched_ext: Fix TOCTOU race in consume_remote_task()

[Andrea Righi <arighi@nvidia.com>]

When pulling a task from a non-local DSQ, scx_consume_dispatch_q()
filters candidates without holding the task's source rq lock;
consume_remote_task() then unlinks the selected task, drops the DSQ and
destination rq locks and acquires the source rq lock before migrating
it.

Whether the task may be remotely migrated can change across this
handoff. In particular, proxy execution allows a task to execute as
rq->curr under another task's scheduling context while its own
scheduling context remains on a DSQ. The task can therefore become
on-CPU or migration disabled without dequeueing its DSQ entry or
clearing holding_cpu.

With proxy execution enabled, a mutex-intensive workload such as
stress-ng --pipeherd 0 can easily trigger this race when a donor becomes
active on its source rq during the unlocked window. Migrating the active
execution context can make it resume with IRQ and preemption state from
the wrong scheduling path, triggering sleeping-while-atomic warnings and
subsequent lockdep corruption.

Fix this by splitting the validation into two helpers:
task_can_run_on_remote_rq() remains the preliminary lockless filter used
while scanning a DSQ; task_can_move_from_locked_rq() is used at actual
migration points, it requires the task's source rq lock, rejects the
task if it is still on-CPU and then applies the existing destination
checks.

After consume_remote_task() has switched to the source rq lock, run the
locked check again immediately before move_remote_task_to_local_dsq().
If the check fails, the task has already been removed from its original
DSQ, so cancel the transfer and enqueue it on the global DSQ. Use a
non-enforcing check, because the failure results from a kernel-side
race, not an invalid BPF placement request.

Deferring this check until the source rq is locked both closes the
migration race and ensures the task remains discoverable by an idle
CPU.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c | 56 +++++++++++++++++++++++++++++++++++-------
 1 file changed, 47 insertions(+), 9 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index ecef4cd9be1fa..08773f0aca5a6 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2364,8 +2364,10 @@ static void move_remote_task_to_local_dsq(struct scx_sched *sch,
  * - The BPF scheduler is bypassed while the rq is offline and we can always say
  *   no to the BPF scheduler initiated migrations while offline.
  *
- * The caller must ensure that @p and @rq are on different CPUs.
- * If enforce == true, caller must hold @p's rq lock.
+ * The caller must ensure that @p and @rq are on different CPUs. If @enforce is
+ * true, report violations attributable to BPF-directed migrations. The caller
+ * must hold @p's rq lock to avoid reporting a transient race as a scheduler
+ * error.
  */
 static bool task_can_run_on_remote_rq(struct scx_sched *sch,
 				      struct task_struct *p, struct rq *rq,
@@ -2373,11 +2375,6 @@ static bool task_can_run_on_remote_rq(struct scx_sched *sch,
 {
 	s32 cpu = cpu_of(rq);
 
-	/*
-	 * To prevent races with @p still running on its old CPU while switching
-	 * out, make sure we're holding @p's rq lock so as not to risk
-	 * erroneously killing the BPF scheduler.
-	 */
 	if (enforce)
 		lockdep_assert_rq_held(task_rq(p));
 
@@ -2424,6 +2421,25 @@ static bool task_can_run_on_remote_rq(struct scx_sched *sch,
 	return true;
 }
 
+/*
+ * Final migration check with @p's rq locked. Unlike the lockless DSQ scan,
+ * this can safely reject transient execution state without losing a kick.
+ */
+static bool task_can_move_from_locked_rq(struct scx_sched *sch,
+					 struct task_struct *p, struct rq *dst_rq,
+					 bool enforce)
+{
+	struct rq *src_rq = task_rq(p);
+
+	lockdep_assert_rq_held(src_rq);
+
+	/* @p may be rq->curr under another task's proxy scheduling context. */
+	if (task_on_cpu(src_rq, p))
+		return false;
+
+	return task_can_run_on_remote_rq(sch, p, dst_rq, enforce);
+}
+
 /**
  * unlink_dsq_and_switch_rq_lock() - Unlink task and switch to its rq lock
  * @p: target task
@@ -2481,6 +2497,28 @@ static bool consume_remote_task(struct scx_sched *sch, struct rq *this_rq,
 				struct scx_dispatch_q *dsq, struct rq *src_rq)
 {
 	if (unlink_dsq_and_switch_rq_lock(p, dsq, this_rq, src_rq)) {
+		/*
+		 * Whether @p may be migrated to @this_rq may have changed while
+		 * switching rq locks. This is a kernel-side race, not an invalid BPF
+		 * placement request, so don't abort the scheduler on failure. Fall
+		 * back to the global DSQ, where normal consumption filters can select
+		 * a valid destination without forcing the task onto the source local
+		 * DSQ.
+		 */
+		if (unlikely(!task_can_move_from_locked_rq(sch, p, this_rq, false))) {
+			p->scx.dsq = NULL;
+			p->scx.holding_cpu = -1;
+			scx_dispatch_enqueue(sch, src_rq,
+					     find_global_dsq(sch, task_cpu(p)), p,
+					     enq_flags | SCX_ENQ_CLEAR_OPSS |
+					     SCX_ENQ_GDSQ_FALLBACK);
+			if (sched_class_above(p->sched_class,
+					      src_rq->donor->sched_class))
+				resched_curr(src_rq);
+			switch_rq_lock(src_rq, this_rq);
+			return false;
+		}
+
 		move_remote_task_to_local_dsq(sch, p, enq_flags, src_rq, this_rq);
 		return true;
 	} else {
@@ -2519,7 +2557,7 @@ static struct rq *move_task_between_dsqs(struct scx_sched *sch,
 	if (dst_dsq->id == SCX_DSQ_LOCAL) {
 		dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
 		if (src_rq != dst_rq &&
-		    unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
+		    unlikely(!task_can_move_from_locked_rq(sch, p, dst_rq, true))) {
 			dst_dsq = find_global_dsq(sch, task_cpu(p));
 			dst_rq = src_rq;
 			enq_flags |= SCX_ENQ_GDSQ_FALLBACK;
@@ -2682,7 +2720,7 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 			p->scx.holding_cpu = -1;
 			scx_dispatch_enqueue(sch, dst_rq, &dst_rq->scx.local_dsq, p,
 					     enq_flags);
-		} else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
+		} else if (unlikely(!task_can_move_from_locked_rq(sch, p, dst_rq, true))) {
 			p->scx.holding_cpu = -1;
 			fallback = true;
 			scx_dispatch_enqueue(sch, src_rq, find_global_dsq(sch, task_cpu(p)),
-- 
2.55.0