Assume task RT1 and RT2 have RT prio, one cpu has scheduled task RT1,
task idle, and task RT2 in order, and rq->scx.cpu_released was true while
task RT1 was running. then rq->scx.cpu_released was changed from true to
false while task RT1 was scheduled out and idle was scheduled in.
But rq->scx.cpu_released could not be changed to true while task RT2 was
scheduled in later.

The sched_class of next task can be observed while update_idle and the
state of rq->scx.cpu_released can be changed properly.

Signed-off-by: liuwenfang liuwenfang@honor.com 
---
 kernel/sched/ext.c      |  2 +-
 kernel/sched/ext.h      | 11 +++++++----
 kernel/sched/ext_idle.c |  6 +++++-
 kernel/sched/ext_idle.h |  1 +
 kernel/sched/idle.c     |  6 +++---
 5 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index f5133249f..6bbea0ea1 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -3187,7 +3187,7 @@ preempt_reason_from_class(const struct sched_class *class)
 	return SCX_CPU_PREEMPT_UNKNOWN;
 }
 
-static void switch_class(struct rq *rq, struct task_struct *next)
+void switch_class(struct rq *rq, struct task_struct *next)
 {
 	const struct sched_class *next_class = next->sched_class;
 
diff --git a/kernel/sched/ext.h b/kernel/sched/ext.h
index 1bda96b19..eddfcc98e 100644
--- a/kernel/sched/ext.h
+++ b/kernel/sched/ext.h
@@ -67,15 +67,18 @@ static inline void init_sched_ext_class(void) {}
 #endif	/* CONFIG_SCHED_CLASS_EXT */
 
 #if defined(CONFIG_SCHED_CLASS_EXT) && defined(CONFIG_SMP)
-void __scx_update_idle(struct rq *rq, bool idle, bool do_notify);
+void __scx_update_idle(struct rq *rq, struct task_struct *next,
+		       bool idle, bool do_notify);
 
-static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify)
+static inline void scx_update_idle(struct rq *rq, struct task_struct *next,
+				   bool idle, bool do_notify)
 {
 	if (scx_enabled())
-		__scx_update_idle(rq, idle, do_notify);
+		__scx_update_idle(rq, next, idle, do_notify);
 }
 #else
-static inline void scx_update_idle(struct rq *rq, bool idle, bool do_notify) {}
+static inline void scx_update_idle(struct rq *rq, struct task_struct *next,
+				   bool idle, bool do_notify) {}
 #endif
 
 #ifdef CONFIG_CGROUP_SCHED
diff --git a/kernel/sched/ext_idle.c b/kernel/sched/ext_idle.c
index e67a19a07..9735f1fb1 100644
--- a/kernel/sched/ext_idle.c
+++ b/kernel/sched/ext_idle.c
@@ -660,12 +660,16 @@ static void update_builtin_idle(int cpu, bool idle)
  * while avoiding unnecessary updates and maintaining balanced state
  * transitions.
  */
-void __scx_update_idle(struct rq *rq, bool idle, bool do_notify)
+void __scx_update_idle(struct rq *rq, struct task_struct *next,
+		       bool idle, bool do_notify)
 {
 	int cpu = cpu_of(rq);
 
 	lockdep_assert_rq_held(rq);
 
+	if (!idle && !rq->scx.cpu_released && next)
+		switch_class(rq, next);
+
 	/*
 	 * Trigger ops.update_idle() only when transitioning from a task to
 	 * the idle thread and vice versa.
diff --git a/kernel/sched/ext_idle.h b/kernel/sched/ext_idle.h
index 511cc2221..83d74ea37 100644
--- a/kernel/sched/ext_idle.h
+++ b/kernel/sched/ext_idle.h
@@ -31,5 +31,6 @@ s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags, u64
 void scx_idle_enable(struct sched_ext_ops *ops);
 void scx_idle_disable(void);
 int scx_idle_init(void);
+void switch_class(struct rq *rq, struct task_struct *next);
 
 #endif /* _KERNEL_SCHED_EXT_IDLE_H */
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index 2c85c86b4..1b77b56bc 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -452,20 +452,20 @@ static void wakeup_preempt_idle(struct rq *rq, struct task_struct *p, int flags)
 static void put_prev_task_idle(struct rq *rq, struct task_struct *prev, struct task_struct *next)
 {
 	dl_server_update_idle_time(rq, prev);
-	scx_update_idle(rq, false, true);
+	scx_update_idle(rq, next, false, true);
 }
 
 static void set_next_task_idle(struct rq *rq, struct task_struct *next, bool first)
 {
 	update_idle_core(rq);
-	scx_update_idle(rq, true, true);
+	scx_update_idle(rq, next, true, true);
 	schedstat_inc(rq->sched_goidle);
 	next->se.exec_start = rq_clock_task(rq);
 }
 
 struct task_struct *pick_task_idle(struct rq *rq)
 {
-	scx_update_idle(rq, true, false);
+	scx_update_idle(rq, NULL, true, false);
 	return rq->idle;
 }
 
-- 
2.17.1