[PATCH sched_ext/for-7.4] sched_ext: Keep proxy donors local through core-sched retries

[Andrea Righi <arighi@nvidia.com>]

An abandoned blocked-donor pick can leave SCX_RQ_PROXY_PENDING set if
the next core-sched pick selects a higher-priority class without calling
pick_task_scx(). A cookie mismatch can then force the CPU idle, while
put_prev_task_scx() mistakes the stale flag for a proxy bookkeeping put.

A valid proxy put of an IMMED donor also schedules a local reenqueue
scan. If an SCX task wakes before the scan runs, rq->next_class becomes
ext again and rq_is_open() reports the CPU closed. The scan then sends
the parked donor back to BPF even though it is first on the local DSQ
and the CPU is switching to idle.

Require the donor's cookie to match before treating a put as a proxy
put. Mark donors parked by a proxy put in dsq_flags, clear the mark
when they leave the DSQ, and let rq_is_open() recognize a matching
marked donor at the head of the local DSQ while the rq donor is idle.

Fixes: 7ab553756463 ("sched_ext: Keep proxy donors with slice left on the local DSQ")
Reported-by: Tejun Heo <tj@kernel.org>
Link: https://lore.kernel.org/r/atCQB_nF5di8pgUN@gpd4
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 include/linux/sched/ext.h |  1 +
 kernel/sched/ext/ext.c    | 28 +++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index ae696d11dac9a..4267304823cf4 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -157,6 +157,7 @@ enum scx_ent_flags {
 /* scx_entity.dsq_flags */
 enum scx_ent_dsq_flags {
 	SCX_TASK_DSQ_ON_PRIQ	= 1 << 0, /* task is queued on the priority queue of a dsq */
+	SCX_TASK_DSQ_PROXY_IDLE	= 1 << 1, /* proxy donor parked through idle */
 };
 
 enum scx_dsq_lnode_flags {
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 4fa294f810f4b..2465e792db2fb 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -437,6 +437,23 @@ static bool rq_is_open(struct rq *rq, u64 enq_flags)
 	if (sched_class_above(&ext_sched_class, rq->next_class))
 		return true;
 
+	/*
+	 * proxy_resched_idle() can put a blocked donor at the head of the local
+	 * DSQ and switch @rq's donor to idle. An SCX wakeup then raises
+	 * next_class to ext before the pending IMMED reenq scan runs. If the
+	 * first task was parked by that proxy put and its cookie can run here,
+	 * leave it local for the next pick instead of sending it back to BPF.
+	 */
+	if (rq->donor == rq->idle) {
+		struct task_struct *first;
+
+		first = list_first_entry_or_null(&rq->scx.local_dsq.list,
+						 struct task_struct, scx.dsq_list.node);
+		if (first && (first->scx.dsq_flags & SCX_TASK_DSQ_PROXY_IDLE) &&
+		    sched_cpu_cookie_match(rq, first))
+			return true;
+	}
+
 	/*
 	 * @rq is either picking, in transition to, or running an SCX task.
 	 */
@@ -1608,6 +1625,7 @@ static void dsq_dec_nr(struct scx_dispatch_q *dsq, struct task_struct *p)
 {
 	/* see dsq_inc_nr() */
 	WRITE_ONCE(dsq->nr, dsq->nr - 1);
+	p->scx.dsq_flags &= ~SCX_TASK_DSQ_PROXY_IDLE;
 
 	if (p->scx.flags & SCX_TASK_IMMED) {
 		struct rq *rq = container_of(dsq, struct rq, scx.local_dsq);
@@ -3521,12 +3539,18 @@ static void switch_class(struct rq *rq, struct task_struct *next)
 	}
 }
 
+/*
+ * A proxy retry can bypass pick_task_scx() and leave PROXY_PENDING set.
+ * A cookie mismatch means core scheduling forced this CPU idle, not that
+ * proxy resolution put a selected donor through idle.
+ */
 static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 			      struct task_struct *next)
 {
 	struct scx_sched *sch = scx_task_sched(p);
 	bool proxy_put = p->is_blocked && next == rq->idle &&
-			 (rq->scx.flags & SCX_RQ_PROXY_PENDING);
+			 (rq->scx.flags & SCX_RQ_PROXY_PENDING) &&
+			 sched_cpu_cookie_match(rq, p);
 	bool proxy_block = rq->scx.flags & SCX_RQ_PROXY_BLOCKING;
 	bool rescue_keep = false;
 
@@ -3612,6 +3636,8 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 						enq_flags |= SCX_ENQ_IMMED;
 				}
 
+				if (proxy_put)
+					p->scx.dsq_flags |= SCX_TASK_DSQ_PROXY_IDLE;
 				scx_dispatch_enqueue(sch, rq, &rq->scx.local_dsq, p, 0, 0,
 						     enq_flags);
 			}
-- 
2.56.0