[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 built-in 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  |  5 ++++-
 kernel/sched/ext/idle.c | 35 ++++++++++++++++++++++++++++++++---
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext/ext.h b/kernel/sched/ext/ext.h
index 0b7fc46aee08c..6d0dab822711f 100644
--- a/kernel/sched/ext/ext.h
+++ b/kernel/sched/ext/ext.h
@@ -59,11 +59,14 @@ static inline void init_sched_ext_class(void) {}
 #endif	/* CONFIG_SCHED_CLASS_EXT */
 
 #ifdef CONFIG_SCHED_CLASS_EXT
+DECLARE_STATIC_KEY_FALSE(scx_builtin_idle_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 (scx_enabled() ||
+	    static_branch_unlikely(&scx_builtin_idle_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..edbfc80a04c74 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -15,7 +15,7 @@
 #include "sub.h"
 
 /* Enable/disable built-in idle CPU selection policy */
-static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
+DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_enabled);
 
 /* Enable/disable per-node idle cpumasks */
 static DEFINE_STATIC_KEY_FALSE(scx_builtin_idle_per_node);
@@ -810,6 +810,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 +847,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 +864,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,6 +894,9 @@ void scx_idle_enable(struct sched_ext_ops *ops)
 		static_branch_disable_cpuslocked(&scx_builtin_idle_per_node);
 
 	reset_idle_masks(ops);
+
+	if (static_branch_likely(&scx_builtin_idle_enabled))
+		refresh_idle_masks();
 }
 
 void scx_idle_disable(void)
-- 
2.55.0