[PATCH v5 sched_ext/for-7.3 28/33] sched_ext: Route ops.update_idle() to sub-schedulers and re-notify owed scheds

__scx_update_idle() notified only the root scheduler. A sub-scheduler that
holds a cid needs that cid's idle state to place and kick on it.

Deliver ops.update_idle() to every scheduler that holds SCX_CAP_BASE on the
transitioning cid. The root holds every cap, so a real transition always
reaches it.

Real transitions are not enough on their own. A cid that is already idle
when a sub-sched gains baseline access produces no transition, so the new
holder would never learn it is idle. The ecaps sync arms a re-notify on the
gain, and the next idle pick delivers ops.update_idle() to just that sched,
leaving holders that already track the cpu untouched. A matching loss of
baseline access drops any pending re-notify.

Bypass suppresses ops.update_idle() too, so a cpu that goes idle during a
bypass window and stays idle yields no transition to re-deliver on
un-bypass. Arm the same re-notify for every sched leaving bypass. The acute
case is a child granted cids during its own ops.sub_attach(). The grant
lands while the child is bypassed and the notify walk skips it, so on
un-bypass it holds cids it never saw go idle. The root is owed the same and
is armed through a separate per-rq flag, which keeps this working when
sub-schedulers are compiled out.

v2: Gate the idle catch-up in pick_task_idle() to avoid a double ops.update_idle(). (sashiko AI)

Signed-off-by: Tejun Heo <[email protected]>
---
 kernel/sched/ext/ext.c      | 40 ++++++++++++++++-
 kernel/sched/ext/idle.c     | 86 ++++++++++++++++++++++++++-----------
 kernel/sched/ext/internal.h |  2 +
 kernel/sched/ext/sub.c      | 24 ++++++++++-
 kernel/sched/idle.c         |  8 +++-
 kernel/sched/sched.h        |  2 +
 6 files changed, 133 insertions(+), 29 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 351240310fd8..bc6e3239de3c 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -5541,6 +5541,38 @@ void scx_disable_bypass_dsp(struct scx_sched *sch)
 	}
 }
 
+/**
+ * unbypass_renotify_idle - Arm an idle re-notify for a sched leaving bypass
+ * @rq: rq of the cpu leaving bypass
+ * @pos: scheduler that just left bypass on @rq's cpu
+ * @pcpu: @pos's per-cpu state for @rq's cpu
+ *
+ * A sched leaving bypass is owed the ops.update_idle() calls suppressed while
+ * bypassing. A cpu that goes idle during the bypass window and stays idle won't
+ * produce a notification. Arm a re-notify that scx_bypass()'s resched flushes
+ * on the next idle pick.
+ *
+ * An acute case is ops.sub_attach(). If the parent grants the child cids while
+ * attaching, when attach is complete and bypass is lifted, the child may hold
+ * idle cids it never saw go idle.
+ *
+ * The root is no exception as bypass suppresses its notifications the same way.
+ * However, the root uses a separate per-rq flag so its re-notify keeps working
+ * even when !CONFIG_EXT_SUB_SCHED.
+ */
+static void unbypass_renotify_idle(struct rq *rq, struct scx_sched *pos,
+				   struct scx_sched_pcpu *pcpu)
+{
+	if (pos == scx_root) {
+		rq->scx.flags |= SCX_RQ_ROOT_IDLE_RENOTIFY;
+		return;
+	}
+#ifdef CONFIG_EXT_SUB_SCHED
+	pcpu->idle_renotify = true;
+	rq->scx.flags |= SCX_RQ_SUB_IDLE_RENOTIFY;
+#endif
+}
+
 /**
  * scx_bypass - [Un]bypass scx_ops and guarantee forward progress
  * @sch: sched to bypass
@@ -5624,11 +5656,15 @@ void scx_bypass(struct scx_sched *sch, bool bypass)
 
 		scx_for_each_descendant_pre(pos, sch) {
 			struct scx_sched_pcpu *pcpu = per_cpu_ptr(pos->pcpu, cpu);
+			bool was_bypassing = pcpu->flags & SCX_SCHED_PCPU_BYPASSING;
 
-			if (pos->bypass_depth)
+			if (pos->bypass_depth) {
 				pcpu->flags |= SCX_SCHED_PCPU_BYPASSING;
-			else
+			} else {
 				pcpu->flags &= ~SCX_SCHED_PCPU_BYPASSING;
+				if (was_bypassing)
+					unbypass_renotify_idle(rq, pos, pcpu);
+			}
 		}
 
 		raw_spin_unlock(&scx_sched_lock);
diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 8e8c6201b7df..16ebe3ab8647 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -12,6 +12,7 @@
 #include "internal.h"
 #include "cid.h"
 #include "idle.h"
+#include "sub.h"
 
 /* Enable/disable built-in idle CPU selection policy */
 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
@@ -730,6 +731,46 @@ static void update_builtin_idle(int cpu, bool idle)
 	}
 }
 
+/*
+ * Notify schedulers of an idle transition on @cpu's cid, delivering to every
+ * sched that holds %SCX_CAP_BASE on the cid (the root holds every cap). A real
+ * transition (@do_notify) reaches all holders. A forced one (@root_renotify for
+ * the root, a sub-sched's idle_renotify marker for a sub) reaches only the owed
+ * scheds.
+ */
+static void scx_idle_notify(struct rq *rq, bool idle, bool do_notify, bool root_renotify)
+{
+	s32 cpu = cpu_of(rq);
+	s32 cid = scx_cpu_arg(cpu);
+	struct scx_sched *pos;
+
+	lockdep_assert_rq_held(rq);
+
+	pos = scx_next_descendant_pre(NULL, scx_root);
+	while (pos) {
+		bool forced = false;
+
+		if (unlikely(scx_missing_caps(pos, cpu, SCX_CAP_BASE))) {
+			pos = scx_skip_subtree_pre(pos, scx_root);
+			continue;
+		}
+
+		if (pos == scx_root) {
+			forced = root_renotify;
+		}
+#ifdef CONFIG_EXT_SUB_SCHED
+		else if (per_cpu_ptr(pos->pcpu, cpu)->idle_renotify) {
+			per_cpu_ptr(pos->pcpu, cpu)->idle_renotify = false;
+			forced = true;
+		}
+#endif
+		if ((do_notify || forced) && SCX_HAS_OP(pos, update_idle) &&
+		    !scx_bypassing(pos, cpu))
+			SCX_CALL_OP(pos, update_idle, rq, cid, idle);
+		pos = scx_next_descendant_pre(pos, scx_root);
+	}
+}
+
 /*
  * Update the idle state of a CPU to @idle.
  *
@@ -748,44 +789,39 @@ static void update_builtin_idle(int cpu, bool idle)
  */
 void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 {
-	struct scx_sched *sch = scx_root;
 	int cpu = cpu_of(rq);
 
 	lockdep_assert_rq_held(rq);
 
 	/*
-	 * Update the idle masks:
-	 * - for real idle transitions (do_notify == true)
-	 * - for idle-to-idle transitions (indicated by the previous task
-	 *   being the idle thread, managed by pick_task_idle())
-	 *
-	 * Skip updating idle masks if the previous task is not the idle
-	 * thread, since set_next_task_idle() has already handled it when
-	 * transitioning from a task to the idle thread (calling this
-	 * function with do_notify == true).
-	 *
-	 * In this way we can avoid updating the idle masks twice,
-	 * unnecessarily.
+	 * pick_task_idle() calls here only on an idle-to-idle re-pick and the
+	 * transitions call with @do_notify, so every reaching call updates the
+	 * masks.
 	 */
 	if (static_branch_likely(&scx_builtin_idle_enabled))
-		if (do_notify || is_idle_task(rq->curr))
-			update_builtin_idle(cpu, idle);
+		update_builtin_idle(cpu, idle);
 
 	/*
-	 * Trigger ops.update_idle() only when transitioning from a task to
-	 * the idle thread and vice versa.
-	 *
-	 * Idle transitions are indicated by do_notify being set to true,
-	 * managed by put_prev_task_idle()/set_next_task_idle().
+	 * ops.update_idle() fires on real idle transitions, indicated by
+	 * @do_notify and managed by put_prev_task_idle()/set_next_task_idle().
+	 * An idle pick also fires it to flush a forced notify owed to a sched
+	 * that missed transitions while bypassed or on a cid it just gained.
+	 * unbypass_renotify_idle() and scx_process_sync_ecaps() arm the per-rq
+	 * gates, and scx_idle_notify() targets the owed scheds.
 	 *
-	 * This must come after builtin idle update so that BPF schedulers can
-	 * create interlocking between ops.update_idle() and ops.enqueue() -
+	 * This must come after the builtin idle update so that BPF schedulers
+	 * can create interlocking between ops.update_idle() and ops.enqueue() -
 	 * either enqueue() sees the idle bit or update_idle() sees the task
 	 * that enqueue() queued.
 	 */
-	if (SCX_HAS_OP(sch, update_idle) && do_notify &&
-	    !scx_bypassing(sch, cpu_of(rq)))
-		SCX_CALL_OP(sch, update_idle, rq, scx_cpu_arg(cpu_of(rq)), idle);
+	if (do_notify ||
+	    (idle && (rq->scx.flags &
+		      (SCX_RQ_SUB_IDLE_RENOTIFY | SCX_RQ_ROOT_IDLE_RENOTIFY)))) {
+		bool root_renotify = rq->scx.flags & SCX_RQ_ROOT_IDLE_RENOTIFY;
+
+		rq->scx.flags &= ~(SCX_RQ_SUB_IDLE_RENOTIFY | SCX_RQ_ROOT_IDLE_RENOTIFY);
+		scx_idle_notify(rq, idle, do_notify, root_renotify);
+	}
 }
 
 static void reset_idle_masks(struct sched_ext_ops *ops)
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index abbcaab8f35a..97d7ce89dd5f 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1244,6 +1244,8 @@ struct scx_sched_pcpu {
 	 */
 	u64			ecaps;
 	struct llist_node	ecaps_to_sync_node;
+	/* owed a forced update_idle() re-notify on this cpu */
+	bool			idle_renotify;
 	/* effective caps as of the last sub_ecaps_updated() delivery */
 	u64			reported_ecaps;
 #endif
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 655632b5511f..a91480debcc2 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -483,6 +483,10 @@ static void discard_queued_syncs(struct rq *rq)
  * pshard->caps[] is the target configuration. pcpu->ecaps is the effective
  * transposed copy owned by the cid's cpu and written only here under @rq's
  * lock.
+ *
+ * A sched that newly gains baseline access here is owed an update_idle() so it
+ * learns the cid's idle state. Such a gain arms the per-rq
+ * %SCX_RQ_SUB_IDLE_RENOTIFY gate so the next idle pick delivers it.
  */
 void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev)
 {
@@ -518,7 +522,7 @@ void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev)
 		struct scx_sched_pcpu *pcpu =
 			container_of(pos, struct scx_sched_pcpu, ecaps_to_sync_node);
 		struct scx_pshard *ps = pcpu->sch->pshard[shard];
-		u64 old, ecaps, lost;
+		u64 old, ecaps, lost, gained;
 
 		init_llist_node(pos);
 
@@ -530,6 +534,7 @@ void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev)
 		WRITE_ONCE(pcpu->ecaps, ecaps);
 
 		lost = old & ~ecaps;
+		gained = ecaps & ~old;
 		lost_all |= lost;
 
 		/* tell the sched its effective caps on this cid changed */
@@ -547,6 +552,18 @@ void scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev)
 			scx_flush_dispatch_buf(pcpu->sch, rq);
 			pcpu->reported_ecaps = ecaps;
 		}
+
+		/*
+		 * Gaining baseline access owes an update_idle() so the sched
+		 * learns the cpu's idle state. Arm the per-rq gate so the next
+		 * idle pick flushes it. Losing access drops any pending notify.
+		 */
+		if (gained & SCX_CAP_BASE) {
+			pcpu->idle_renotify = true;
+			rq->scx.flags |= SCX_RQ_SUB_IDLE_RENOTIFY;
+		} else if (lost & SCX_CAP_BASE) {
+			pcpu->idle_renotify = false;
+		}
 	}
 
 	/*
@@ -1441,6 +1458,11 @@ __bpf_kfunc s32 scx_bpf_sub_grant(u64 cgroup_id, u64 caps,
 
 				caps_updated_record(cps, changed_cids, granted_caps,
 						    &to_deliver);
+				/*
+				 * The sync arms an update_idle() re-notify if
+				 * the cid gains baseline access, so the holder
+				 * learns of an already-idle cid.
+				 */
 				scx_cmask_for_each_cid(cid, changed_cids)
 					queue_sync_ecaps(child, cid);
 			}
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index 052435f4d3e3..eb73b65ce6c4 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -503,7 +503,13 @@ static void set_next_task_idle(struct rq *rq, struct task_struct *next, bool fir
 
 struct task_struct *pick_task_idle(struct rq *rq, struct rq_flags *rf)
 {
-	scx_update_idle(rq, true, false);
+	/*
+	 * Notify scx only on an idle-to-idle re-pick (the cpu was already idle).
+	 * A real task->idle transition is delivered by set_next_task_idle(), so
+	 * calling here too would duplicate it.
+	 */
+	if (scx_enabled() && is_idle_task(rq->curr))
+		scx_update_idle(rq, true, false);
 	return rq->idle;
 }
 
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index cdd21c814008..54cff94556c0 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -787,6 +787,8 @@ enum scx_rq_flags {
 	SCX_RQ_BAL_KEEP		= 1 << 3, /* balance decided to keep current */
 	SCX_RQ_CLK_VALID	= 1 << 5, /* RQ clock is fresh and valid */
 	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_IN_WAKEUP	= 1 << 16,
 	SCX_RQ_IN_BALANCE	= 1 << 17,
-- 
2.55.0