[PATCH 06/12] sched_ext: Fix proxy-exec race in consume_remote_task()

[Andrea Righi]

Without proxy execution, the DSQ lock and holding_cpu handshake ensure
that a task cannot be dequeued or start running during the rq-lock
handoff without clearing holding_cpu.

Proxy execution is an exception. A task can start physically executing
as a lock owner while its own scheduling context remains on the DSQ. Its
on-CPU or migration-disabled state can therefore change without clearing
holding_cpu.

Keep the existing lockless destination validation unchanged. After
switching to the source rq lock, recheck only the on-CPU and
migration-disabled states when proxy execution is enabled. If either
changed, cancel the transfer and requeue the task on the global DSQ
without treating the race as a BPF scheduler error.

A rejected task may be requeued onto the DSQ currently being consumed.
Snapshot the DSQ sequence and skip tasks added after the consume attempt
started to prevent repeatedly selecting the same task. Keep this retry
bounding confined to proxy execution so the non-proxy path remains
unchanged.

With proxy execution enabled, stress-ng --pipeherd can trigger this race
and migrate an active execution context, leading to
sleeping-while-atomic warnings and subsequent lockdep corruption.

Signed-off-by: Andrea Righi <arighi@xxxxxxxxxx>
---
 kernel/sched/ext/ext.c | 79 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 69 insertions(+), 10 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 30d5965d09108..b5fd91d9bca4f 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2372,8 +2372,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,
@@ -2381,11 +2383,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));
 
@@ -2432,6 +2429,28 @@ static bool task_can_run_on_remote_rq(struct scx_sched *sch,
 	return true;
 }
 
+/*
+ * Proxy execution can change @p's execution and migration-disabled state
+ * without touching its DSQ entry or clearing holding_cpu. Check those states
+ * with @p's rq locked. Without proxy execution, the holding_cpu handshake is
+ * sufficient and this must not affect the existing migration path.
+ */
+static bool task_can_move_from_locked_rq(struct task_struct *p)
+{
+	struct rq *src_rq = task_rq(p);
+
+	lockdep_assert_rq_held(src_rq);
+
+	if (!sched_proxy_exec())
+		return true;
+
+	/* @p may be rq->curr under another task's proxy scheduling context. */
+	if (task_on_cpu(src_rq, p))
+		return false;
+
+	return !is_migration_disabled(p);
+}
+
 /**
  * unlink_dsq_and_switch_rq_lock() - Unlink task and switch to its rq lock
  * @p: target task
@@ -2489,6 +2508,27 @@ 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)) {
+		/*
+		 * Proxy execution may have changed @p's running or migration-disabled
+		 * state while switching rq locks without clearing holding_cpu. 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.
+		 */
+		if (unlikely(!task_can_move_from_locked_rq(p))) {
+			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 {
@@ -2527,7 +2567,8 @@ 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(p) ||
+			     !task_can_run_on_remote_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;
@@ -2570,7 +2611,9 @@ bool scx_consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
 			    struct scx_dispatch_q *dsq, u64 enq_flags)
 {
 	struct task_struct *p;
-retry:
+	bool proxy_exec = sched_proxy_exec();
+	u32 scan_seq = 0;
+
 	/*
 	 * The caller can't expect to successfully consume a task if the task's
 	 * addition to @dsq isn't guaranteed to be visible somehow. Test
@@ -2580,10 +2623,23 @@ bool scx_consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
 		return false;
 
 	raw_spin_lock(&dsq->lock);
+	if (proxy_exec)
+		scan_seq = READ_ONCE(dsq->seq);
 
+retry:
 	nldsq_for_each_task(p, dsq) {
 		struct rq *task_rq = task_rq(p);
 
+		/*
+		 * consume_remote_task() may drop @dsq->lock and requeue a
+		 * rejected task onto this same DSQ. Don't reconsider tasks queued
+		 * after this consume attempt started. This bounds the retry loop
+		 * even when the rejected task remains on-CPU.
+		 */
+		if (proxy_exec &&
+		    unlikely(u32_before(scan_seq, p->scx.dsq_seq)))
+			continue;
+
 		/*
 		 * This loop can lead to multiple lockup scenarios, e.g. the BPF
 		 * scheduler can put an enormous number of affinitized tasks into
@@ -2605,6 +2661,8 @@ bool scx_consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
 		if (task_can_run_on_remote_rq(sch, p, rq, false)) {
 			if (likely(consume_remote_task(sch, rq, p, enq_flags, dsq, task_rq)))
 				return true;
+			/* consume_remote_task() returns with @dsq unlocked. */
+			raw_spin_lock(&dsq->lock);
 			goto retry;
 		}
 	}
@@ -2690,7 +2748,8 @@ 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(p) ||
+				    !task_can_run_on_remote_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