[PATCH 09/14] sched_ext: Handle proxy-exec races in remote DSQ transfers

[Andrea Righi <arighi@nvidia.com>]

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

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

Recheck these states after acquiring the source rq lock. If the transfer
can no longer proceed, park the task on the source rq's reject DSQ and
re-enqueue it through its owning scheduler. This preserves the BPF
scheduler's placement policy and keeps descendant tasks within their
sub-scheduler's cap grants.

Implement the scx_proxy_resolved() hook to drain the parked tasks once
proxy resolution has settled and the outgoing owner has switched out.

Without this change and 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.

This is a preparatory change to support proxy execution with sched_ext.

Suggested-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 include/linux/sched/ext.h |   5 ++
 kernel/sched/ext/ext.c    | 154 +++++++++++++++++++++++++++++++++-----
 kernel/sched/sched.h      |   1 +
 3 files changed, 140 insertions(+), 20 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index 91b255fdd16bc..4d16c225e9af1 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -135,6 +135,9 @@ enum scx_ent_flags {
 	 * IMMED	reenqueued due to failed ENQ_IMMED
 	 * PREEMPTED	preempted while running
 	 * CAP		sub-sched cap miss, see p->scx.reenq_reason_*
+	 * MIGRATION_DISABLED
+	 *		migration-disabled during a remote DSQ transfer
+	 * PROXY	physically executing or donating during a remote DSQ transfer
 	 */
 	SCX_TASK_REENQ_REASON_SHIFT = 12,
 	SCX_TASK_REENQ_REASON_BITS = 3,
@@ -145,6 +148,8 @@ enum scx_ent_flags {
 	SCX_TASK_REENQ_IMMED	= 2 << SCX_TASK_REENQ_REASON_SHIFT,
 	SCX_TASK_REENQ_PREEMPTED = 3 << SCX_TASK_REENQ_REASON_SHIFT,
 	SCX_TASK_REENQ_CAP	= 4 << SCX_TASK_REENQ_REASON_SHIFT,
+	SCX_TASK_REENQ_MIGRATION_DISABLED = 5 << SCX_TASK_REENQ_REASON_SHIFT,
+	SCX_TASK_REENQ_PROXY	= 6 << SCX_TASK_REENQ_REASON_SHIFT,
 
 	/* iteration cursor, not a task */
 	SCX_TASK_CURSOR		= 1 << 31,
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 2fcd303b44894..95aca029a6e57 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1069,8 +1069,17 @@ static void schedule_deferred_locked(struct rq *rq)
 	schedule_deferred(rq);
 }
 
+/*
+ * Proxy resolution happens before rq->curr is switched. Queue deferred work
+ * on the rq so that an outgoing proxy owner has cleared on_cpu by the time
+ * reject_dsq is drained.
+ */
 void scx_proxy_resolved(struct rq *rq)
 {
+	lockdep_assert_rq_held(rq);
+
+	if (rq->scx.flags & SCX_RQ_PROXY_REENQ)
+		schedule_deferred_locked(rq);
 }
 
 void schedule_dsq_reenq(struct scx_sched *sch, struct scx_dispatch_q *dsq,
@@ -1449,9 +1458,15 @@ static void rq_owned_post_enq(struct scx_sched *sch, struct rq *rq,
 {
 	call_task_dequeue(sch, rq, p, 0);
 
-	/* rejected: kick the deferred reenq, skip wakeup/preemption */
+	/*
+	 * Proxy-active tasks must remain parked until proxy resolution. Other
+	 * rejects can be reenqueued immediately.
+	 */
 	if (unlikely(dsq->id == SCX_DSQ_REJECT)) {
-		schedule_deferred_locked(rq);
+		if (p->scx.reject_reason == SCX_TASK_REENQ_PROXY)
+			rq->scx.flags |= SCX_RQ_PROXY_REENQ;
+		else
+			schedule_deferred_locked(rq);
 		return;
 	}
 
@@ -2350,8 +2365,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,
@@ -2359,11 +2376,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));
 
@@ -2410,6 +2422,60 @@ 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 u32 task_move_reject_reason(struct task_struct *p)
+{
+	struct rq *src_rq = task_rq(p);
+
+	lockdep_assert_rq_held(src_rq);
+
+	if (!sched_proxy_exec())
+		return SCX_TASK_REENQ_NONE;
+
+	/* @p may be rq->curr under another task's scheduling context. */
+	if (task_on_cpu(src_rq, p))
+		return SCX_TASK_REENQ_PROXY;
+
+	/*
+	 * Reject only BPF-directed migration. proxy_migrate_task() may still
+	 * move a blocked donor's scheduling context to its lock owner's CPU.
+	 */
+	if (is_migration_disabled(p))
+		return SCX_TASK_REENQ_MIGRATION_DISABLED;
+
+	/* Don't move an active scheduling context off its source rq. */
+	if (task_current_donor(src_rq, p))
+		return SCX_TASK_REENQ_PROXY;
+
+	return SCX_TASK_REENQ_NONE;
+}
+
+/*
+ * Park a task whose remote transfer raced with proxy execution. Reenqueueing
+ * from the source rq makes the task's owning scheduler choose its placement
+ * again and preserves sub-scheduler containment.
+ */
+static void scx_reject_task(struct scx_sched *sch, struct rq *rq,
+			    struct task_struct *p, u64 enq_flags, u32 reason)
+{
+	lockdep_assert_rq_held(rq);
+	WARN_ON_ONCE(reason != SCX_TASK_REENQ_MIGRATION_DISABLED &&
+		     reason != SCX_TASK_REENQ_PROXY);
+	WARN_ON_ONCE(p->scx.reject_reason);
+
+	p->scx.holding_cpu = -1;
+	p->scx.reject_reason = reason;
+	p->scx.flags &= ~SCX_TASK_IMMED;
+	enq_flags &= ~(SCX_ENQ_IMMED | SCX_ENQ_PREEMPT);
+
+	scx_dispatch_enqueue(sch, rq, &rq->scx.reject_dsq, p, enq_flags);
+}
+
 /**
  * unlink_dsq_and_switch_rq_lock() - Unlink task and switch to its rq lock
  * @p: target task
@@ -2467,6 +2533,23 @@ 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)) {
+		u32 reject_reason = task_move_reject_reason(p);
+
+		/*
+		 * Proxy execution may have changed @p's running or
+		 * migration-disabled state while switching rq locks without
+		 * clearing holding_cpu. Park it on the source rq and let its
+		 * owning scheduler choose its placement again.
+		 */
+		if (unlikely(reject_reason)) {
+			p->scx.dsq = NULL;
+			scx_reject_task(sch, src_rq, p,
+					enq_flags | SCX_ENQ_CLEAR_OPSS,
+					reject_reason);
+			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 {
@@ -2497,6 +2580,7 @@ static struct rq *move_task_between_dsqs(struct scx_sched *sch,
 					 struct scx_dispatch_q *dst_dsq)
 {
 	struct rq *src_rq = task_rq(p), *dst_rq;
+	u32 reject_reason;
 
 	BUG_ON(src_dsq->id == SCX_DSQ_LOCAL);
 	lockdep_assert_held(&src_dsq->lock);
@@ -2504,6 +2588,15 @@ 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);
+		reject_reason = src_rq != dst_rq ?
+			task_move_reject_reason(p) : SCX_TASK_REENQ_NONE;
+		if (unlikely(reject_reason)) {
+			dispatch_dequeue_locked(p, src_dsq);
+			raw_spin_unlock(&src_dsq->lock);
+			scx_reject_task(sch, src_rq, p, enq_flags,
+					reject_reason);
+			return src_rq;
+		}
 		if (src_rq != dst_rq &&
 		    unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
 			dst_dsq = find_global_dsq(sch, task_cpu(p));
@@ -2659,6 +2752,11 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 	if (likely(p->scx.holding_cpu == raw_smp_processor_id()) &&
 	    !WARN_ON_ONCE(src_rq != task_rq(p))) {
 		bool fallback = false;
+		u32 reject_reason;
+
+		reject_reason = src_rq != dst_rq ?
+			task_move_reject_reason(p) : SCX_TASK_REENQ_NONE;
+
 		/*
 		 * If @p is staying on the same rq, there's no need to go
 		 * through the full deactivate/activate cycle. Optimize by
@@ -2668,9 +2766,14 @@ 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))) {
-			p->scx.holding_cpu = -1;
+		} else if (unlikely(reject_reason)) {
 			fallback = true;
+			scx_reject_task(sch, src_rq, p, enq_flags,
+					reject_reason);
+		} else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq,
+							       true))) {
+			fallback = true;
+			p->scx.holding_cpu = -1;
 			scx_dispatch_enqueue(sch, src_rq, find_global_dsq(sch, task_cpu(p)),
 					     p, enq_flags | SCX_ENQ_GDSQ_FALLBACK);
 		} else {
@@ -4382,35 +4485,41 @@ static void process_deferred_reenq_users(struct rq *rq)
 }
 
 /*
- * Drain @rq->scx.reject_dsq and reenqueue each task so that its owning BPF
- * scheduler chooses placement again.
- *
- * A task can be re-rejected repeatedly, and there's no repeat limit here. The
- * private list below prevents a task from being revisited in the same round.
+ * Drain ready tasks from @rq->scx.reject_dsq and reenqueue them so that their
+ * owning BPF schedulers choose placement again. Proxy-active tasks remain
+ * parked until proxy resolution schedules another drain after switch-out.
  */
 static void scx_reenq_reject(struct rq *rq)
 {
 	LIST_HEAD(tasks);
 	struct task_struct *p, *n;
+	bool proxy_pending = false;
 
 	lockdep_assert_rq_held(rq);
 
-	if (list_empty(&rq->scx.reject_dsq.list))
+	if (list_empty(&rq->scx.reject_dsq.list)) {
+		rq->scx.flags &= ~SCX_RQ_PROXY_REENQ;
 		return;
+	}
 
 	/*
-	 * Move tasks to a private list so a task re-rejected by
+	 * Move ready tasks to a private list so a task re-rejected by
 	 * scx_do_enqueue_task() below isn't revisited this round.
 	 */
 	list_for_each_entry_safe(p, n, &rq->scx.reject_dsq.list, scx.dsq_list.node) {
 		u32 reason = p->scx.reject_reason;
 
 		/* migration_pending tasks should have bypassed to local DSQ */
-		if (WARN_ON_ONCE(p->migration_pending))
-			continue;
+		WARN_ON_ONCE(p->migration_pending);
 		if (WARN_ON_ONCE(!reason))
 			continue;
 
+		if (reason == SCX_TASK_REENQ_PROXY &&
+		    (task_on_cpu(rq, p) || task_current_donor(rq, p))) {
+			proxy_pending = true;
+			continue;
+		}
+
 		scx_dispatch_dequeue(rq, p);
 		p->scx.reject_reason = SCX_TASK_REENQ_NONE;
 
@@ -4421,6 +4530,11 @@ static void scx_reenq_reject(struct rq *rq)
 		list_add_tail(&p->scx.dsq_list.node, &tasks);
 	}
 
+	if (proxy_pending)
+		rq->scx.flags |= SCX_RQ_PROXY_REENQ;
+	else
+		rq->scx.flags &= ~SCX_RQ_PROXY_REENQ;
+
 	list_for_each_entry_safe(p, n, &tasks, scx.dsq_list.node) {
 		list_del_init(&p->scx.dsq_list.node);
 
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 00a7f864ac23d..fdcd723c29a3c 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -789,6 +789,7 @@ enum scx_rq_flags {
 	SCX_RQ_BAL_CB_PENDING	= 1 << 6, /* must queue a cb after dispatching */
 	SCX_RQ_SUB_IDLE_RENOTIFY	= 1 << 7, /* sub-scheds are owed update_idle() */
 	SCX_RQ_ROOT_IDLE_RENOTIFY	= 1 << 8, /* the root is owed update_idle() */
+	SCX_RQ_PROXY_REENQ	= 1 << 9, /* proxy-rejected tasks need reenqueue */
 
 	SCX_RQ_IN_WAKEUP	= 1 << 16,
 	SCX_RQ_IN_BALANCE	= 1 << 17,
-- 
2.55.0