[PATCH 1/2] sched_ext: Initialize idle masks before ops.init()

[Andrea Righi <arighi@nvidia.com>]

The built-in idle masks are reset with all online CPUs marked idle, but
idle state tracking starts only after the scheduler is fully enabled.
As a result, ops.init() can observe busy CPUs as idle, and those CPUs
remain incorrectly advertised until their next idle transition.

Enable idle tracking before ops.init() and refresh every online CPU under
its rq lock. Once a CPU is refreshed, later transitions keep its state
accurate. Keep ops.update_idle() notifications disabled until the
scheduler is fully enabled.

Suggested-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/ext.h  |  4 +++-
 kernel/sched/ext/idle.c | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext/ext.h b/kernel/sched/ext/ext.h
index 0b7fc46aee08c..495e7023c9ee1 100644
--- a/kernel/sched/ext/ext.h
+++ b/kernel/sched/ext/ext.h
@@ -59,11 +59,13 @@ static inline void init_sched_ext_class(void) {}
 #endif	/* CONFIG_SCHED_CLASS_EXT */
 
 #ifdef CONFIG_SCHED_CLASS_EXT
+DECLARE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
+
 void __scx_update_idle(struct rq *rq, bool idle, bool do_notify);
 
 static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 {
-	if (scx_enabled())
+	if (static_branch_unlikely(&scx_idle_tracking_enabled))
 		__scx_update_idle(rq, idle, do_notify);
 }
 #else
diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 3e9d6a44bf431..6d81bb7c43966 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -14,6 +14,9 @@
 #include "idle.h"
 #include "sub.h"
 
+/* Enable/disable idle state tracking */
+DEFINE_STATIC_KEY_FALSE(scx_idle_tracking_enabled);
+
 /* Enable/disable built-in idle CPU selection policy */
 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
 
@@ -810,6 +813,15 @@ void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
 	if (static_branch_likely(&scx_builtin_idle_enabled))
 		update_builtin_idle(cpu, idle);
 
+	/*
+	 * Idle tracking starts before the scheduler is enabled so that the
+	 * built-in idle masks are accurate when ops.init() runs. Suppress
+	 * ops.update_idle() notifications until the scheduler is fully
+	 * enabled.
+	 */
+	if (!scx_enabled())
+		return;
+
 	/*
 	 * ops.update_idle() fires on real idle transitions, indicated by
 	 * @do_notify and managed by put_prev_task_idle()/set_next_task_idle().
@@ -838,8 +850,8 @@ static void reset_idle_masks(struct sched_ext_ops *ops)
 	int node;
 
 	/*
-	 * Consider all online cpus idle. Should converge to the actual state
-	 * quickly.
+	 * Seed all online CPUs as idle. refresh_idle_masks() below corrects
+	 * their state before ops.init() runs.
 	 */
 	if (!(ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE)) {
 		cpumask_copy(idle_cpumask(NUMA_NO_NODE)->cpu, cpu_online_mask);
@@ -855,6 +867,23 @@ static void reset_idle_masks(struct sched_ext_ops *ops)
 	}
 }
 
+static void refresh_idle_masks(void)
+{
+	int cpu;
+
+	/*
+	 * Idle tracking is already enabled and the online CPU set is stable.
+	 * Once a CPU is refreshed under its rq lock, subsequent transitions
+	 * keep its state up to date.
+	 */
+	for_each_online_cpu(cpu) {
+		struct rq *rq = cpu_rq(cpu);
+
+		scoped_guard(rq_lock_irqsave, rq)
+			update_builtin_idle(cpu, rq->curr == rq->idle);
+	}
+}
+
 void scx_idle_enable(struct sched_ext_ops *ops)
 {
 	if (!ops->update_idle || (ops->flags & SCX_OPS_KEEP_BUILTIN_IDLE))
@@ -868,10 +897,15 @@ void scx_idle_enable(struct sched_ext_ops *ops)
 		static_branch_disable_cpuslocked(&scx_builtin_idle_per_node);
 
 	reset_idle_masks(ops);
+	static_branch_enable_cpuslocked(&scx_idle_tracking_enabled);
+
+	if (static_branch_likely(&scx_builtin_idle_enabled))
+		refresh_idle_masks();
 }
 
 void scx_idle_disable(void)
 {
+	static_branch_disable(&scx_idle_tracking_enabled);
 	static_branch_disable(&scx_builtin_idle_enabled);
 	static_branch_disable(&scx_builtin_idle_per_node);
 }
-- 
2.55.0