[PATCH 12/16] sched_ext: Track proxy execution for NOHZ_FULL

[Andrea Righi <arighi@nvidia.com>]

scx_can_stop_tick() currently treats any blocked rq->donor as an active
proxy session. However, rq->donor can temporarily still identify an
outgoing blocked task while the scheduler selects an ordinary execution
context. A dependency update in that window can keep TICK_DEP_BIT_SCHED
set after the context switch.

Track whether proxy resolution actually selected a different execution
context. Keep the tick active while that state is set instead of
deriving it from a potentially stale donor.

When proxy execution ends, clear the state and queue a balance callback
to reevaluate the tick dependency after context_switch() updates
rq->curr. This lets sched_can_stop_tick() observe the complete donor and
execution selection without adding work outside the existing
proxy-execution branch.

Conservatively re-enable the periodic scheduler tick for the duration of
proxy execution. Allowing proxy execution itself to run tickless would
require making the remote NOHZ scheduler tick aware of the split between
rq->curr and rq->donor, which is left as a future improvement.

Perform the tick update from scx_proxy_reenqueue_retry(), which already
runs after proxy resolution, so scheduler core needs no additional hook.
Compile the additional runqueue state and callback out when
CONFIG_NO_HZ_FULL is disabled.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.c | 44 +++++++++++++++++++++++++++++++++++++++---
 kernel/sched/sched.h   |  4 ++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 1f4016c04d527..247da76021044 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1130,13 +1130,51 @@ static void schedule_deferred_locked(struct rq *rq)
 	schedule_deferred(rq);
 }
 
+#ifdef CONFIG_NO_HZ_FULL
+static void scx_proxy_tick_bal_cb(struct rq *rq)
+{
+	sched_update_tick_dependency(rq);
+}
+
+static void scx_proxy_update_tick(struct rq *rq, struct task_struct *next)
+{
+	bool proxy = next != rq->donor;
+	bool was_proxy = rq->scx.flags & SCX_RQ_PROXY_TICK;
+
+	if (likely(proxy == was_proxy))
+		return;
+
+	if (proxy) {
+		/* Keep proxy execution tick-driven for now. */
+		rq->scx.flags |= SCX_RQ_PROXY_TICK;
+		tick_nohz_dep_set_cpu(cpu_of(rq), TICK_DEP_BIT_SCHED);
+	} else {
+		rq->scx.flags &= ~SCX_RQ_PROXY_TICK;
+		/*
+		 * The selected donor is already visible, but rq->curr still
+		 * identifies the outgoing execution context. Reevaluate after
+		 * context_switch() updates rq->curr so sched_can_stop_tick() sees
+		 * the complete selection.
+		 */
+		queue_balance_callback(rq, &rq->scx.proxy_tick_bal_cb,
+				       scx_proxy_tick_bal_cb);
+	}
+}
+#endif
+
 /*
- * Retry proxy-rejected tasks which couldn't be reenqueued by an earlier drain.
+ * Complete sched_ext bookkeeping after proxy resolution and retry tasks which
+ * couldn't be reenqueued by an earlier reject DSQ drain.
  */
 void scx_proxy_reenqueue_retry(struct rq *rq, struct task_struct *next)
 {
 	lockdep_assert_rq_held(rq);
 
+#ifdef CONFIG_NO_HZ_FULL
+	if (scx_enabled() && tick_nohz_full_cpu(cpu_of(rq)))
+		scx_proxy_update_tick(rq, next);
+#endif
+
 	if (rq->scx.flags & SCX_RQ_PROXY_RETRY) {
 		rq->scx.flags &= ~SCX_RQ_PROXY_RETRY;
 		schedule_deferred_locked(rq);
@@ -4966,8 +5004,8 @@ bool scx_can_stop_tick(struct rq *rq)
 	struct task_struct *p = rq->donor;
 	struct scx_sched *sch = scx_task_sched(p);
 
-	/* Keep the tick running while a blocked proxy donor is selected. */
-	if (p->is_blocked)
+	/* Proxy execution is conservatively tick-driven for now. */
+	if (rq->scx.flags & SCX_RQ_PROXY_TICK)
 		return false;
 
 	if (p->sched_class != &ext_sched_class)
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 52c60a884994f..083074edd0bab 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -789,6 +789,7 @@ enum scx_rq_flags {
 	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_RETRY	= 1 << 9, /* proxy-rejected tasks need retry */
+	SCX_RQ_PROXY_TICK	= 1 << 10, /* proxy execution requires the tick */
 
 	SCX_RQ_IN_WAKEUP	= 1 << 16,
 	SCX_RQ_IN_DISPATCH	= 1 << 17,
@@ -843,6 +844,9 @@ struct scx_rq {
 	struct list_head	deferred_reenq_users;	/* user DSQs requesting reenq */
 	struct balance_callback	deferred_bal_cb;
 	struct balance_callback	kick_sync_bal_cb;
+#ifdef CONFIG_NO_HZ_FULL
+	struct balance_callback	proxy_tick_bal_cb;
+#endif
 	struct irq_work		deferred_irq_work;
 	struct irq_work		kick_cpus_irq_work;
 };
-- 
2.55.0