[PATCH v2 00/12] sched: core-sched fixes and balancing

[PATCH v2 00/12] sched: core-sched fixes and balancing

[Peter Zijlstra]

Hi!

Previous version at:

  https://lore.kernel.org/r/20260828101659.812011872@infradead.org

The main difference is reworking of rq_clock due to TJ noticing the previous
one got it all wrong. The new patches might seem a little aimless at times;
this is because my initial idea for fixing the problem didn't pan out, but I
kept some of the groundwork for it anyway, since it resulted in simpler (I
hope) code.

Code at:

  git://git.kernel.org/pub/scm/linux/kernel/git/peterz/queue.git sched/hackery

[PATCH v2 01/12] sched/core: Fix pick_next_task() self recursion

[Peter Zijlstra]

It is possible for another sibling to end up in pick_next_task() when:

  pick_next_task()
    pick_task()
      sched_class::pick_task()

drops the core wide rq->lock. In this case they end up trampling the core wide
task selection state, possibly leading to NULL derefs. Detect this case by
keeping a local copy of core_task_seq, a value that is incremented on
{en,de}queue and schedule.

Since RETRY_TASK is only possible when a higher priority task gets enqueued
during the lock break, this must mean core_task_seq will also be incremented
and is thus completely covered by the seq number mismatch.

If there are forward progress issues; those might possibly be solved by using
core_pick_in_flight.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c |   22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6296,14 +6296,14 @@ static struct task_struct *
 pick_next_task(struct rq *rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(rq))
 {
+	bool core_clock_updated = (rq == rq->core);
 	struct task_struct *next, *p, *max;
 	const struct cpumask *smt_mask;
+	int i, cpu, seq, occ = 0;
 	bool fi_before = false;
-	bool core_clock_updated = (rq == rq->core);
+	bool need_sync = false;
 	unsigned long cookie;
-	int i, cpu, occ = 0;
 	struct rq *rq_i;
-	bool need_sync = false;
 
 	if (!sched_core_enabled(rq))
 		return __pick_next_task(rq, rf);
@@ -6377,7 +6377,7 @@ pick_next_task(struct rq *rq, struct rq_
 	 * However, preemptions can cause multiple picks on the same task set.
 	 * 'Fix' this by also increasing @task_seq for every pick.
 	 */
-	rq->core->core_task_seq++;
+	seq = ++rq->core->core_task_seq;
 
 	/*
 	 * Optimize for common case where this CPU has no cookies
@@ -6394,6 +6394,15 @@ pick_next_task(struct rq *rq, struct rq_
 		}
 
 		if (!next->core_cookie) {
+			/*
+			 * pick_task() can drop the core rq lock through
+			 * newidle balance. If a sibling established a
+			 * core-wide cookie while the lock was dropped, the
+			 * uncookied no-sync fast path is no longer valid.
+			 */
+			if (unlikely(rq->core->core_cookie))
+				goto restart;
+
 			rq->core_pick = NULL;
 			rq->core_dl_server = NULL;
 			/*
@@ -6425,7 +6434,8 @@ pick_next_task(struct rq *rq, struct rq_
 			update_rq_clock(rq_i);
 
 		p = pick_task(rq_i, rf);
-		if (unlikely(p == RETRY_TASK)) {
+		if (unlikely(seq != rq->core->core_task_seq ||
+			     WARN_ON_ONCE(p == RETRY_TASK))) {
 			/* rq lock may have been dropped, clocks invalidated */
 			core_clock_updated = false;
 			if (!(rq->clock_update_flags & RQCF_UPDATED))
@@ -6455,7 +6465,7 @@ pick_next_task(struct rq *rq, struct rq_
 			if (cookie)
 				p = sched_core_find(rq_i, cookie);
 			if (!p)
-				p = idle_sched_class.pick_task(rq_i, rf);
+				p = idle_sched_class.pick_task(rq_i, NULL);
 		}
 
 		rq_i->core_pick = p;

[PATCH v2 02/12] sched/core: Introduce rq_drop_lock()

[Peter Zijlstra]

Currently rq_unpin_lock() serves a dual purpose, it pairs with both
rq_pin_lock() and rq_repin_lock(). That is:

	rq_pin_lock()
	  rq_unpin_lock()
          rq_repin_lock()
	rq_unpin_lock()

In prepration of making the temporary drop and the final drop
distinct, this needs another function. Introduce and use
rq_drop_lock() for the temporary drop, like:

	rq_pin_lock()
	  rq_drop_lock()
          rq_repin_lock()
	rq_unpin_lock()

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c     |    8 ++++----
 kernel/sched/deadline.c |    4 ++--
 kernel/sched/ext/ext.c  |    4 ++--
 kernel/sched/fair.c     |    2 +-
 kernel/sched/rt.c       |    2 +-
 kernel/sched/sched.h    |    5 +++++
 6 files changed, 15 insertions(+), 10 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3910,7 +3910,7 @@ ttwu_do_activate(struct rq *rq, struct t
 		 * Our task @p is fully woken up and running; so it's safe to
 		 * drop the rq->lock, hereafter rq is only used for statistics.
 		 */
-		rq_unpin_lock(rq, rf);
+		rq_drop_lock(rq, rf);
 		p->sched_class->task_woken(rq, p);
 		rq_repin_lock(rq, rf);
 	}
@@ -5042,7 +5042,7 @@ void wake_up_new_task(struct task_struct
 		 * Nothing relies on rq->lock after this, so it's fine to
 		 * drop it.
 		 */
-		rq_unpin_lock(rq, &rf);
+		rq_drop_lock(rq, &rf);
 		p->sched_class->task_woken(rq, p);
 		rq_repin_lock(rq, &rf);
 	}
@@ -5260,7 +5260,7 @@ struct balance_callback *splice_balance_
 void __balance_callbacks(struct rq *rq, struct rq_flags *rf)
 {
 	if (rf)
-		rq_unpin_lock(rq, rf);
+		rq_drop_lock(rq, rf);
 	do_balance_callbacks(rq, __splice_balance_callbacks(rq, false));
 	if (rf)
 		rq_repin_lock(rq, rf);
@@ -6929,7 +6929,7 @@ static inline void proxy_release_rq_lock
 	 * to pick_again, so the callbacks will get re-established.
 	 */
 	zap_balance_callbacks(rq);
-	rq_unpin_lock(rq, rf);
+	rq_drop_lock(rq, rf);
 	raw_spin_rq_unlock(rq);
 }
 
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -1129,7 +1129,7 @@ static void __push_dl_task(struct rq *rq
 		 * Nothing relies on rq->lock after this, so its safe to drop
 		 * rq->lock.
 		 */
-		rq_unpin_lock(rq, rf);
+		rq_drop_lock(rq, rf);
 		push_dl_task(rq);
 		rq_repin_lock(rq, rf);
 	}
@@ -2720,7 +2720,7 @@ static int balance_dl(struct rq *rq, str
 		 * disabled avoiding further scheduler activity on it and we've
 		 * not yet started the picking loop.
 		 */
-		rq_unpin_lock(rq, rf);
+		rq_drop_lock(rq, rf);
 		pull_dl_task(rq);
 		rq_repin_lock(rq, rf);
 	}
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3924,7 +3924,7 @@ static enum scx_dsp_verdict dispatch_pic
 {
 	enum scx_dsp_verdict verdict;
 
-	rq_unpin_lock(rq, rf);
+	rq_drop_lock(rq, rf);
 	verdict = dispatch_one(rq, prev);
 	rq_repin_lock(rq, rf);
 
@@ -3949,7 +3949,7 @@ static enum scx_dsp_verdict dispatch_cor
 	if (rq->scx.flags & SCX_RQ_IN_DISPATCH)
 		return SCX_DSP_NONE;
 
-	rq_unpin_lock(rq, rf);
+	rq_drop_lock(rq, rf);
 
 	verdict = dispatch_one(rq, prev);
 
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -14907,7 +14907,7 @@ static int sched_balance_newidle(struct
 	 * further scheduler activity on it and we're being very careful to
 	 * re-start the picking loop.
 	 */
-	rq_unpin_lock(this_rq, rf);
+	rq_drop_lock(this_rq, rf);
 
 	sd = rcu_dereference_sched_domain(this_rq->sd);
 	if (!sd)
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1611,7 +1611,7 @@ static int balance_rt(struct rq *rq, str
 		 * disabled avoiding further scheduler activity on it and we've
 		 * not yet started the picking loop.
 		 */
-		rq_unpin_lock(rq, rf);
+		rq_drop_lock(rq, rf);
 		pull_rt_task(rq);
 		rq_repin_lock(rq, rf);
 	}
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1991,6 +1991,11 @@ static inline void rq_unpin_lock(struct
 	lockdep_unpin_lock(__rq_lockp(rq), rf->cookie);
 }
 
+static inline void rq_drop_lock(struct rq *rq, struct rq_flags *rf)
+{
+	rq_unpin_lock(rq, rf);
+}
+
 static inline void rq_repin_lock(struct rq *rq, struct rq_flags *rf)
 {
 	lockdep_repin_lock(__rq_lockp(rq), rf->cookie);

[PATCH v2 03/12] sched/core: Decompose update_rq_clock()

[Peter Zijlstra]

Adds __update_rq_clock() which lacks the double update warning. Additionally it
will now always set UPDATED, even though it might skip the actual update.
Because when we skip the update, we must be up-to-date by definition.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c |   35 +++++++++++++++++++----------------
 1 file changed, 19 insertions(+), 16 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -867,29 +867,32 @@ static void update_rq_clock_task(struct
 	update_rq_clock_pelt(rq, delta);
 }
 
-void update_rq_clock(struct rq *rq)
+static void __update_rq_clock(struct rq *rq)
 {
-	s64 delta;
-	u64 clock;
-
 	lockdep_assert_rq_held(rq);
 
-	if (rq->clock_update_flags & RQCF_ACT_SKIP)
-		return;
-
-	if (sched_feat(WARN_DOUBLE_CLOCK))
-		WARN_ON_ONCE(rq->clock_update_flags & RQCF_UPDATED);
 	rq->clock_update_flags |= RQCF_UPDATED;
 
-	clock = sched_clock_cpu(cpu_of(rq));
-	scx_rq_clock_update(rq, clock);
+	if (!(rq->clock_update_flags & RQCF_ACT_SKIP)) {
+		u64 clock = sched_clock_cpu(cpu_of(rq));
+		s64 delta;
+
+		scx_rq_clock_update(rq, clock);
+
+		delta = clock - rq->clock;
+		if (delta > 0) {
+			rq->clock += delta;
+			update_rq_clock_task(rq, delta);
+		}
+	}
+}
 
-	delta = clock - rq->clock;
-	if (delta < 0)
-		return;
-	rq->clock += delta;
+void update_rq_clock(struct rq *rq)
+{
+	if (sched_feat(WARN_DOUBLE_CLOCK))
+		WARN_ON_ONCE(rq->clock_update_flags & RQCF_UPDATED);
 
-	update_rq_clock_task(rq, delta);
+	__update_rq_clock(rq);
 }
 
 #ifdef CONFIG_SCHED_HRTICK

[PATCH v2 04/12] sched/core: Make __update_rq_clock() skip updates on UPDATED

[Peter Zijlstra]

Calling update_rq_clock() on a rq which is already up-to-date would still move
the clock (after warning when WARN_DOUBLE_CLOCK is set). This is a bit silly,
so also skip the update when already up-to-date.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c  |    7 +++++--
 kernel/sched/sched.h |    2 ++
 2 files changed, 7 insertions(+), 2 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -869,11 +869,14 @@ static void update_rq_clock_task(struct
 
 static void __update_rq_clock(struct rq *rq)
 {
+	bool update;
+
 	lockdep_assert_rq_held(rq);
 
+	update = !(rq->clock_update_flags & RQCF_NOP_MASK);
 	rq->clock_update_flags |= RQCF_UPDATED;
 
-	if (!(rq->clock_update_flags & RQCF_ACT_SKIP)) {
+	if (update) {
 		u64 clock = sched_clock_cpu(cpu_of(rq));
 		s64 delta;
 
@@ -890,7 +893,7 @@ static void __update_rq_clock(struct rq
 void update_rq_clock(struct rq *rq)
 {
 	if (sched_feat(WARN_DOUBLE_CLOCK))
-		WARN_ON_ONCE(rq->clock_update_flags & RQCF_UPDATED);
+		WARN_ON_ONCE((rq->clock_update_flags & RQCF_NOP_MASK) == RQCF_UPDATED);
 
 	__update_rq_clock(rq);
 }
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1807,6 +1807,8 @@ extern void update_rq_clock(struct rq *r
 #define RQCF_ACT_SKIP		0x02
 #define RQCF_UPDATED		0x04
 
+#define RQCF_NOP_MASK		(RQCF_ACT_SKIP | RQCF_UPDATED)
+
 static inline void assert_clock_updated(struct rq *rq)
 {
 	/*

[PATCH v2 05/12] sched/core: Abstract/fix rq_clock_start_loop_update()

[Peter Zijlstra]

The existing rq_clock_start_loop_update() / rq_lock_stop_loop_update() almost
do what we need rq_drop_lock() / rq_repin_lock() to do. They freeze time by
setting ACT_SKIP.

Rename them rq_clock_freeze() / rq_clock_thaw() and have them operate on a new
bit RQCF_FREEZE. This is such that schedule() can safely nest inside
rq_clock_freeze() as is required for rq_drop_lock().

Sanitize the rq->clock_update_flags handling while adding the new bit --
notably do away with the full shift left which makes a mess of things.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c  |    4 +--
 kernel/sched/fair.c  |   15 +++++++-----
 kernel/sched/sched.h |   63 +++++++++++++++++++++++++++++----------------------
 3 files changed, 48 insertions(+), 34 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7185,9 +7185,9 @@ static void __sched notrace __schedule(i
 	hrtick_schedule_enter(rq);
 
 	/* Promote REQ to ACT */
-	rq->clock_update_flags <<= 1;
+	rq->clock_update_flags |= (rq->clock_update_flags & RQCF_REQ_SKIP) << 1;
 	update_rq_clock(rq);
-	rq->clock_update_flags = RQCF_UPDATED;
+	rq->clock_update_flags &= (RQCF_FREEZE|RQCF_UPDATED);
 
 	switch_count = &prev->nivcsw;
 
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -5445,6 +5445,7 @@ static inline void clear_tg_load_avg(str
 static void __maybe_unused clear_tg_offline_cfs_rqs(struct rq *rq)
 {
 	struct task_group *tg;
+	unsigned int cuf;
 
 	lockdep_assert_rq_held(rq);
 
@@ -5453,7 +5454,7 @@ static void __maybe_unused clear_tg_offl
 	 * set_rq_offline(), so we should skip updating
 	 * the rq clock again in unthrottle_cfs_rq().
 	 */
-	rq_clock_start_loop_update(rq);
+	cuf = rq_clock_freeze(rq);
 
 	guard(rcu)();
 
@@ -5463,7 +5464,7 @@ static void __maybe_unused clear_tg_offl
 		clear_tg_load_avg(cfs_rq);
 	}
 
-	rq_clock_stop_loop_update(rq);
+	rq_clock_thaw(rq, cuf);
 }
 
 /*
@@ -7331,6 +7332,7 @@ static void __cfsb_csd_unthrottle(void *
 {
 	struct cfs_rq *cursor, *tmp;
 	struct rq *rq = arg;
+	unsigned int cuf;
 
 	guard(rq_lock)(rq);
 
@@ -7340,7 +7342,7 @@ static void __cfsb_csd_unthrottle(void *
 	 * Do it once and skip the potential next ones.
 	 */
 	update_rq_clock(rq);
-	rq_clock_start_loop_update(rq);
+	cuf = rq_clock_freeze(rq);
 
 	/*
 	 * Since we hold rq lock we're safe from concurrent manipulation of
@@ -7359,7 +7361,7 @@ static void __cfsb_csd_unthrottle(void *
 			unthrottle_cfs_rq(cursor);
 	}
 
-	rq_clock_stop_loop_update(rq);
+	rq_clock_thaw(rq, cuf);
 }
 
 static inline void __unthrottle_cfs_rq_async(struct cfs_rq *cfs_rq)
@@ -7865,6 +7867,7 @@ static void __maybe_unused update_runtim
 static void __maybe_unused unthrottle_offline_cfs_rqs(struct rq *rq)
 {
 	struct task_group *tg;
+	unsigned int cuf;
 
 	lockdep_assert_rq_held(rq);
 
@@ -7877,7 +7880,7 @@ static void __maybe_unused unthrottle_of
 	 * set_rq_offline(), so we should skip updating
 	 * the rq clock again in unthrottle_cfs_rq().
 	 */
-	rq_clock_start_loop_update(rq);
+	cuf = rq_clock_freeze(rq);
 
 	guard(rcu)();
 
@@ -7904,7 +7907,7 @@ static void __maybe_unused unthrottle_of
 		unthrottle_cfs_rq(cfs_rq);
 	}
 
-	rq_clock_stop_loop_update(rq);
+	rq_clock_thaw(rq, cuf);
 }
 
 bool cfs_task_bw_constrained(struct task_struct *p)
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1790,24 +1790,33 @@ extern void update_rq_clock(struct rq *r
  * %RQCF_ACT_SKIP - is set from inside of __schedule() when skipping is
  *  in effect and calls to update_rq_clock() are being ignored.
  *
+ * %RQCF_FREEZE - set from rq_clock_freeze(), much like ACT_SKIP it
+ *  makes update_rq_clock() a noop.
+ *
  * %RQCF_UPDATED - is a debug flag that indicates whether a call has been
  *  made to update_rq_clock() since the last time rq::lock was pinned.
  *
- * If inside of __schedule(), clock_update_flags will have been
- * shifted left (a left shift is a cheap operation for the fast path
- * to promote %RQCF_REQ_SKIP to %RQCF_ACT_SKIP), so you must use,
- *
- *	if (rq-clock_update_flags >= RQCF_UPDATED)
- *
- * to check if %RQCF_UPDATED is set. It'll never be shifted more than
- * one position though, because the next rq_unpin_lock() will shift it
- * back.
+ *
+ * yield():
+ *   will generally result in REQ_SKIP.
+ *
+ * update_rq_clock():
+ *   sets UPDATED, will not modify rq->clock when NOP_MASK
+ *
+ * schedule():
+ *   will promote REQ_SKIP to ACT_SKIP before its update_rq_clock() to ensure
+ *   the yield() and schedule() appear to happen instantaneous.
+ *
+ * rq_clock_freeze()/rq_clock_thaw():
+ *   set FREEZE to make 'nested' update_rq_clock() calls NOP and avoid
+ *   time being moved forward. Also used in rq_drop_lock()/rq_repin_lock().
  */
 #define RQCF_REQ_SKIP		0x01
 #define RQCF_ACT_SKIP		0x02
-#define RQCF_UPDATED		0x04
+#define RQCF_FREEZE		0x04
+#define RQCF_UPDATED		0x08
 
-#define RQCF_NOP_MASK		(RQCF_ACT_SKIP | RQCF_UPDATED)
+#define RQCF_NOP_MASK		(RQCF_ACT_SKIP | RQCF_FREEZE | RQCF_UPDATED)
 
 static inline void assert_clock_updated(struct rq *rq)
 {
@@ -1815,7 +1824,7 @@ static inline void assert_clock_updated(
 	 * The only reason for not seeing a clock update since the
 	 * last rq_pin_lock() is if we're currently skipping updates.
 	 */
-	WARN_ON_ONCE(rq->clock_update_flags < RQCF_ACT_SKIP);
+	WARN_ON_ONCE(!(rq->clock_update_flags & RQCF_UPDATED));
 }
 
 static inline u64 rq_clock(struct rq *rq)
@@ -1851,25 +1860,27 @@ static inline void rq_clock_cancel_skipu
 }
 
 /*
- * During cpu offlining and rq wide unthrottling, we can trigger
- * an update_rq_clock() for several cfs and rt runqueues (Typically
- * when using list_for_each_entry_*)
- * rq_clock_start_loop_update() can be called after updating the clock
- * once and before iterating over the list to prevent multiple update.
- * After the iterative traversal, we need to call rq_clock_stop_loop_update()
- * to clear RQCF_ACT_SKIP of rq->clock_update_flags.
+ * FREEZE rq_clock. This is used in a number of places to
+ * ensure things appears to happen at the same time.
+ *
+ * Eg. during cpu offlining and rq wide unthrottling, we can trigger an
+ * update_rq_clock() for several cfs and rt runqueues (Typically when using
+ * list_for_each_entry_*)
  */
-static inline void rq_clock_start_loop_update(struct rq *rq)
+static inline unsigned int rq_clock_freeze(struct rq *rq)
 {
+	unsigned int flags = rq->clock_update_flags;
 	lockdep_assert_rq_held(rq);
-	WARN_ON_ONCE(rq->clock_update_flags & RQCF_ACT_SKIP);
-	rq->clock_update_flags |= RQCF_ACT_SKIP;
+	assert_clock_updated(rq);
+	rq->clock_update_flags |= RQCF_FREEZE;
+	return flags;
 }
 
-static inline void rq_clock_stop_loop_update(struct rq *rq)
+static inline void rq_clock_thaw(struct rq *rq, bool flags)
 {
 	lockdep_assert_rq_held(rq);
-	rq->clock_update_flags &= ~RQCF_ACT_SKIP;
+	rq->clock_update_flags &= ~RQCF_FREEZE;
+	rq->clock_update_flags |= RQCF_UPDATED | (flags & RQCF_FREEZE);
 }
 
 struct rq_flags {
@@ -1938,14 +1949,14 @@ static inline void rq_pin_lock(struct rq
 {
 	rf->cookie = lockdep_pin_lock(__rq_lockp(rq));
 
-	rq->clock_update_flags &= (RQCF_REQ_SKIP|RQCF_ACT_SKIP);
+	rq->clock_update_flags &= ~RQCF_UPDATED;
 	rf->clock_update_flags = 0;
 	assert_balance_callbacks_empty(rq);
 }
 
 static inline void rq_unpin_lock(struct rq *rq, struct rq_flags *rf)
 {
-	if (rq->clock_update_flags > RQCF_ACT_SKIP)
+	if (rq->clock_update_flags & RQCF_UPDATED)
 		rf->clock_update_flags = RQCF_UPDATED;
 
 	scx_rq_clock_invalidate(rq);

[PATCH v2 06/12] sched/core: New rq_{pin,drop,repin,unpin}_lock() semantics

[Peter Zijlstra]

The current semantics are such that UPDATED is cleared on rq_pin_lock() but
remains set on rq_unpin_lock(). This works because rq_clock() itself asserts
rq->lock is held.

However, there are two problems with this, both related to the newidle/balance
lock drop.

The code as-is, tries to preserve UPDATED over drop()/repin(), but during the
drop an interleaving lock()+update_rq_clock() will still move the clock --
since it will clear UPDATED on lock(). Make this stronger using the freshly
minted freeze and thaw. This way the clock remains where it was, restoring the
'atomicity' of schedule().

Additionally, there is a variant with core-scheduling. In that case rq->lock is
shared among the sibling threads, but rq->clock and rq->clock_update_flags are
very much per thread. This disparity can cause another form of the above
problem, where one sibling is trying to forward the clock of the other sibling.

Sort this by making lock() clear all sibling UPDATED, but only for non-nested
(using core_pick_in_flight as nesting counter).

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/sched.h |   61 ++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 51 insertions(+), 10 deletions(-)

--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1936,6 +1936,51 @@ static inline void assert_balance_callba
 }
 
 /*
+ * We want clocks to be updated _once_ during core scheduling.
+ *
+ * rq_lock(smt0)
+ *   __rq_clock_clear(); // outer
+ *
+ * core->core_pick_in_flight++;
+ * update_rq_clock(smt0);
+ * drop()
+ *   freeze();
+ *
+ *			rq_lock(smt1)
+ *			  __rq_clock_clear(); // inner
+ *
+ *			update_rq_clock(smt1)
+ *
+ *			rq_unlock(smt1)
+ *
+ * relock()
+ *   thaw();
+ * update_rq_clock(smt1) // no second update
+ * core_pick_in_flight--;
+ *
+ * rq_unlock()
+ */
+static inline void __rq_clock_clear(struct rq *rq)
+{
+	if (!sched_core_enabled(rq)) {
+		rq->clock_update_flags &= ~RQCF_UPDATED;
+		return;
+	}
+
+#ifdef CONFIG_SCHED_CORE
+	/*
+	 * Inner rq_{,un}lock() should preserve state.
+	 */
+	if (!rq->core->core_pick_in_flight) {
+		int cpu;
+
+		for_each_cpu(cpu, cpu_smt_mask(cpu_of(rq)))
+			cpu_rq(cpu)->clock_update_flags &= ~RQCF_UPDATED;
+	}
+#endif
+}
+
+/*
  * Lockdep annotation that avoids accidental unlocks; it's like a
  * sticky/continuous lockdep_assert_held().
  *
@@ -1949,33 +1994,29 @@ static inline void rq_pin_lock(struct rq
 {
 	rf->cookie = lockdep_pin_lock(__rq_lockp(rq));
 
-	rq->clock_update_flags &= ~RQCF_UPDATED;
+	__rq_clock_clear(rq);
 	rf->clock_update_flags = 0;
 	assert_balance_callbacks_empty(rq);
 }
 
 static inline void rq_unpin_lock(struct rq *rq, struct rq_flags *rf)
 {
-	if (rq->clock_update_flags & RQCF_UPDATED)
-		rf->clock_update_flags = RQCF_UPDATED;
-
+	/* __rq_clock_clear(rq); */
 	scx_rq_clock_invalidate(rq);
 	lockdep_unpin_lock(__rq_lockp(rq), rf->cookie);
 }
 
 static inline void rq_drop_lock(struct rq *rq, struct rq_flags *rf)
 {
-	rq_unpin_lock(rq, rf);
+	rf->clock_update_flags = rq_clock_freeze(rq);
+	scx_rq_clock_invalidate(rq);
+	lockdep_unpin_lock(__rq_lockp(rq), rf->cookie);
 }
 
 static inline void rq_repin_lock(struct rq *rq, struct rq_flags *rf)
 {
 	lockdep_repin_lock(__rq_lockp(rq), rf->cookie);
-
-	/*
-	 * Restore the value we stashed in @rf for this pin context.
-	 */
-	rq->clock_update_flags |= rf->clock_update_flags;
+	rq_clock_thaw(rq, rf->clock_update_flags);
 }
 
 #define __task_rq_lock(...) __acquire_ret(___task_rq_lock(__VA_ARGS__), __rq_lockp(__ret))

[PATCH v2 07/12] sched/core: Simplify/fix time updates

[Peter Zijlstra]

Use __update_rq_clock() and the new update semantics instead of using
convoluted logic to try and divinate if a clock needs updating.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c |   25 ++++++-------------------
 1 file changed, 6 insertions(+), 19 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6302,7 +6302,6 @@ static struct task_struct *
 pick_next_task(struct rq *rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(rq))
 {
-	bool core_clock_updated = (rq == rq->core);
 	struct task_struct *next, *p, *max;
 	const struct cpumask *smt_mask;
 	int i, cpu, seq, occ = 0;
@@ -6358,10 +6357,7 @@ pick_next_task(struct rq *rq, struct rq_
 	/* reset state */
 	rq->core->core_cookie = 0UL;
 	if (rq->core->core_forceidle_count) {
-		if (!core_clock_updated) {
-			update_rq_clock(rq->core);
-			core_clock_updated = true;
-		}
+		__update_rq_clock(rq->core);
 		sched_core_account_forceidle(rq);
 		/* reset after accounting force idle */
 		rq->core->core_forceidle_start = 0;
@@ -6390,14 +6386,11 @@ pick_next_task(struct rq *rq, struct rq_
 	 * and there are no cookied tasks running on siblings.
 	 */
 	if (!need_sync) {
+		__update_rq_clock(rq);
+
 		next = pick_task(rq, rf);
-		if (unlikely(next == RETRY_TASK)) {
-			/* rq lock may have been dropped, clocks invalidated */
-			core_clock_updated = false;
-			if (!(rq->clock_update_flags & RQCF_UPDATED))
-				update_rq_clock(rq);
+		if (unlikely(next == RETRY_TASK))
 			goto restart;
-		}
 
 		if (!next->core_cookie) {
 			/*
@@ -6436,18 +6429,12 @@ pick_next_task(struct rq *rq, struct rq_
 		 * pick_next_task(). If the current cpu is not the core,
 		 * the core may also have been updated above.
 		 */
-		if (i != cpu && (rq_i != rq->core || !core_clock_updated))
-			update_rq_clock(rq_i);
+		__update_rq_clock(rq_i);
 
 		p = pick_task(rq_i, rf);
 		if (unlikely(seq != rq->core->core_task_seq ||
-			     WARN_ON_ONCE(p == RETRY_TASK))) {
-			/* rq lock may have been dropped, clocks invalidated */
-			core_clock_updated = false;
-			if (!(rq->clock_update_flags & RQCF_UPDATED))
-				update_rq_clock(rq);
+			     WARN_ON_ONCE(p == RETRY_TASK)))
 			goto restart;
-		}
 
 		rq_i->core_pick = p;
 		rq_i->core_dl_server = rq_i->dl_server;

[PATCH v2 08/12] sched/core: Allow newidle for core-sched

[Peter Zijlstra]

As reported [1], it is possible for newidle, as called from pick_task_fair(),
to migrate a task that was already selected for the SMT sibling.

That is, in case of core scheduling (pick_next_task()'s restart_multi label),
pick_task() is called for each SMT sibling. Suppose SMT0 picks task A, and SMT1
has no tasks and does newidle, it must not be possible (but currently is) to
migrate A to SMT1 and also select A.

Avoids the 'stealing', but doesn't yet enable newidle just yet.

[1] https://patch.msgid.link/20260603095108.GA1684319@bytedance.com

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c  |    3 ++-
 kernel/sched/fair.c  |    3 ++-
 kernel/sched/sched.h |   15 ++++++++++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3065,7 +3065,8 @@ static int affine_move_task(struct rq *r
 		return -EINVAL;
 	}
 
-	if (task_on_cpu(rq, p) || READ_ONCE(p->__state) == TASK_WAKING) {
+	if (task_on_cpu(rq, p) || task_on_core(rq, p) ||
+	    READ_ONCE(p->__state) == TASK_WAKING) {
 		/*
 		 * MIGRATE_ENABLE gets here because 'p == current', but for
 		 * anything else we cannot do is_migration_disabled(), punt
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -11184,7 +11184,8 @@ int can_migrate_task(struct task_struct
 	env->flags &= ~LBF_ALL_PINNED;
 
 	if (task_on_cpu(env->src_rq, p) ||
-	    task_current_donor(env->src_rq, p)) {
+	    task_current_donor(env->src_rq, p) ||
+	    task_on_core(env->src_rq, p)) {
 		schedstat_inc(p->stats.nr_failed_migrations_running);
 		return 0;
 	}
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1594,6 +1594,14 @@ static inline bool task_has_sched_core(s
 	return !!p->core_cookie;
 }
 
+static inline bool task_on_core(struct rq *rq, struct task_struct *p)
+{
+	if (sched_core_disabled())
+		return false;
+
+	return rq->core_pick == p;
+}
+
 #else /* !CONFIG_SCHED_CORE: */
 
 static inline bool sched_core_enabled(struct rq *rq)
@@ -1639,6 +1647,11 @@ static inline bool task_has_sched_core(s
 	return false;
 }
 
+static inline bool task_on_core(struct rq *rq, struct task_struct *p)
+{
+	return false;
+}
+
 #endif /* !CONFIG_SCHED_CORE */
 
 #ifdef CONFIG_RT_GROUP_SCHED
@@ -4269,7 +4282,7 @@ void move_queued_task_locked(struct rq *
 static inline
 bool task_is_pushable(struct rq *rq, struct task_struct *p, int cpu)
 {
-	if (!task_on_cpu(rq, p) &&
+	if (!task_on_cpu(rq, p) && !task_on_core(rq, p) &&
 	    cpumask_test_cpu(cpu, &p->cpus_mask))
 		return true;
 

[PATCH v2 09/12] sched/rt: Add early exit on balance path

[Peter Zijlstra]

Add/make consistent an early exit on the balance_{rt,dl}() path. When there are
no pushable tasks, there is no point in taking the rq lock and all that jazz.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/deadline.c |   10 +++-------
 kernel/sched/rt.c       |    3 +++
 2 files changed, 6 insertions(+), 7 deletions(-)

--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -3250,17 +3250,13 @@ static void pull_dl_task(struct rq *this
 				   src_rq->dl.earliest_dl.next))
 			continue;
 
+		if (!has_pushable_dl_tasks(src_rq))
+			continue;
+
 		/* Might drop this_rq->lock */
 		push_task = NULL;
 		double_lock_balance(this_rq, src_rq);
 
-		/*
-		 * If there are no more pullable tasks on the
-		 * rq, we're done with it.
-		 */
-		if (src_rq->dl.dl_nr_running <= 1)
-			goto skip;
-
 		p = pick_earliest_pushable_dl_task(src_rq, this_cpu);
 
 		/*
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -2307,6 +2307,9 @@ static void pull_rt_task(struct rq *this
 		    this_rq->rt.highest_prio.curr)
 			continue;
 
+		if (!has_pushable_tasks(src_rq))
+			continue;
+
 		/*
 		 * We can potentially drop this_rq's lock in
 		 * double_lock_balance, and another CPU could

[PATCH v2 10/12] sched/fair: Reflow pick_task_fair() / newidle

[Peter Zijlstra]

This is all histerical artifacts. Clean up the code to be less convoluted.

Note: this also re-enables balance for core-sched.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/fair.c |   51 +++++++++++++--------------------------------------
 1 file changed, 13 insertions(+), 38 deletions(-)

--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -6100,7 +6100,7 @@ static inline unsigned long cfs_rq_load_
 	return cfs_rq->avg.load_avg;
 }
 
-static int sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
+static void sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(this_rq));
 
 static inline unsigned long task_util(struct task_struct *p)
@@ -10287,11 +10287,17 @@ struct task_struct *pick_task_fair(struc
 	struct cfs_rq *cfs_rq = &rq->cfs;
 	struct sched_entity *se;
 	struct task_struct *p;
-	int new_tasks;
 
 again:
-	if (!cfs_rq->h_nr_queued)
-		goto idle;
+	if (unlikely(!cfs_rq->h_nr_queued)) {
+		rq_modified_begin(rq, &fair_sched_class);
+		sched_balance_newidle(rq, rf);
+		if (rq_modified_above(rq, &fair_sched_class))
+			return RETRY_TASK;
+
+		if (!cfs_rq->h_nr_queued)
+			return NULL;
+	}
 
 	/* Might not have done put_prev_entity() */
 	if (cfs_rq->curr && cfs_rq->curr->on_rq)
@@ -10303,17 +10309,6 @@ struct task_struct *pick_task_fair(struc
 
 	p = task_of(se);
 	return p;
-
-idle:
-	if (sched_core_enabled(rq))
-		return NULL;
-
-	new_tasks = sched_balance_newidle(rq, rf);
-	if (new_tasks < 0)
-		return RETRY_TASK;
-	if (new_tasks > 0)
-		goto again;
-	return NULL;
 }
 
 static struct task_struct *
@@ -14869,13 +14864,8 @@ static inline void nohz_newidle_balance(
 /*
  * sched_balance_newidle is called by schedule() if this_cpu is about to become
  * idle. Attempts to pull tasks from other CPUs.
- *
- * Returns:
- *   < 0 - we released the lock and there are !fair tasks present
- *     0 - failed, no new tasks
- *   > 0 - success, new (fair) tasks present
  */
-static int sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
+static void sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(this_rq))
 {
 	unsigned long next_balance = jiffies + HZ;
@@ -14892,7 +14882,7 @@ static int sched_balance_newidle(struct
 	 * Return 0; the task will be enqueued when switching to idle.
 	 */
 	if (this_rq->ttwu_pending)
-		return 0;
+		return;
 
 	/*
 	 * We must set idle_stamp _before_ calling sched_balance_rq()
@@ -14903,7 +14893,7 @@ static int sched_balance_newidle(struct
 
 	/* Do not pull tasks towards !preferred CPUs */
 	if (!cpu_preferred(this_cpu))
-		return 0;
+		return;
 
 	/*
 	 * This is OK, because current is on_cpu, which avoids it being picked
@@ -14932,7 +14922,6 @@ static int sched_balance_newidle(struct
 	t0 = sched_clock_cpu(this_cpu);
 	__sched_balance_update_blocked_averages(this_rq);
 
-	rq_modified_begin(this_rq, &fair_sched_class);
 	raw_spin_rq_unlock(this_rq);
 
 	for_each_domain(this_cpu, sd) {
@@ -14990,18 +14979,6 @@ static int sched_balance_newidle(struct
 	if (curr_cost > this_rq->max_idle_balance_cost)
 		this_rq->max_idle_balance_cost = curr_cost;
 
-	/*
-	 * While browsing the domains, we released the rq lock, a task could
-	 * have been enqueued in the meantime. Since we're not going idle,
-	 * pretend we pulled a task.
-	 */
-	if (this_rq->cfs.h_nr_queued && !pulled_task)
-		pulled_task = 1;
-
-	/* If a higher prio class was modified, restart the pick */
-	if (rq_modified_above(this_rq, &fair_sched_class))
-		pulled_task = -1;
-
 out:
 	/* Move the next balance forward */
 	if (time_after(this_rq->next_balance, next_balance))
@@ -15013,8 +14990,6 @@ static int sched_balance_newidle(struct
 		nohz_newidle_balance(this_rq);
 
 	rq_repin_lock(this_rq, rf);
-
-	return pulled_task;
 }
 
 /*

[PATCH v2 11/12] sched/fair: Push sched_balance_newidle() unlock down

[Peter Zijlstra]

By pushing down the this_rq unlock, there are more cases where the lock isn't
dropped.

XXX needs numbers

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/fair.c |   23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -13706,6 +13706,7 @@ static int sched_balance_rq(int this_cpu
 			struct sched_domain *sd, enum cpu_idle_type idle,
 			int *continue_balancing)
 {
+	bool lock_rq = false, unlock_rq = idle == CPU_NEWLY_IDLE;
 	int ld_moved, cur_ld_moved, active_balance = 0;
 	struct sched_domain *sd_parent = sd->parent;
 	struct sched_group *group;
@@ -13774,6 +13775,13 @@ static int sched_balance_rq(int this_cpu
 		 */
 		env.loop_max  = min(sysctl_sched_nr_migrate, busiest->nr_running);
 
+		if (unlock_rq) {
+			lockdep_assert_rq_held(this_rq);
+			raw_spin_rq_unlock(this_rq);
+			unlock_rq = false;
+			lock_rq = true;
+		}
+
 more_balance:
 		rq_lock_irqsave(busiest, &rf);
 		update_rq_clock(busiest);
@@ -13902,6 +13910,13 @@ static int sched_balance_rq(int this_cpu
 	if (!need_active_balance(&env))
 		goto out_unbalanced;
 
+	if (unlock_rq) {
+		lockdep_assert_rq_held(this_rq);
+		raw_spin_rq_unlock(this_rq);
+		unlock_rq = false;
+		lock_rq = true;
+	}
+
 	scoped_guard (raw_spin_rq_lock_irqsave, busiest) {
 		/*
 		 * Don't kick the active_load_balance_cpu_stop,
@@ -13995,6 +14010,10 @@ static int sched_balance_rq(int this_cpu
 out:
 	if (need_unlock)
 		atomic_set_release(&sched_balance_running, 0);
+	if (lock_rq) {
+		raw_spin_rq_lock(this_rq);
+		__release_ctx_lock(__rq_lockp(this_rq));
+	}
 
 	return ld_moved;
 }
@@ -14878,8 +14897,6 @@ static void sched_balance_newidle(struct
 	t0 = sched_clock_cpu(this_cpu);
 	__sched_balance_update_blocked_averages(this_rq);
 
-	raw_spin_rq_unlock(this_rq);
-
 	for_each_domain(this_cpu, sd) {
 		u64 domain_cost;
 
@@ -14930,8 +14947,6 @@ static void sched_balance_newidle(struct
 			break;
 	}
 
-	raw_spin_rq_lock(this_rq);
-
 	if (curr_cost > this_rq->max_idle_balance_cost)
 		this_rq->max_idle_balance_cost = curr_cost;
 

[PATCH v2 12/12] sched: Remove sched_class::balance()

[Peter Zijlstra]

Ever since commit 50653216e4ff ("sched: Add support to pick functions to take
rf"), we have the unfortunate situation that both sched_class::balance() and
sched_class::pick_task() have overlapping functionality in that they drop
rq->lock and balance tasks.

Additionally, prev_balance() is only called for a single RQ in the core-sched
case, resulting in 'missed' balance opportunities in this case.

Just like pick_task(), prev_balance() would iterate the classes and stop when
there is a runnable task found. The only difference is that pick_task() starts
from the top class, while prev_balance() would start from prev->class.

This means that if prev was an RT task, we'd get an extra visit to balance_dl()
-- and similary, if prev is fair, we'd visit balance_{dl,rt}() both. However,
this is not a problem, because both need_pull_{dl,rt}_task() will DTRT.

Therefore, move balance_{rt,dl}() into pick_task_{rt,dl}() and remove
sched_class::balance().

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c      |   23 -----------------------
 kernel/sched/deadline.c  |   10 ++++++----
 kernel/sched/idle.c      |    7 -------
 kernel/sched/rt.c        |   14 ++++++--------
 kernel/sched/sched.h     |    5 -----
 kernel/sched/stop_task.c |    7 -------
 6 files changed, 12 insertions(+), 54 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6185,25 +6185,6 @@ static inline void schedule_debug(struct
 	schedstat_inc(this_rq()->sched_count);
 }
 
-static void prev_balance(struct rq *rq, struct rq_flags *rf)
-{
-	const struct sched_class *start_class = rq->donor->sched_class;
-	const struct sched_class *class;
-
-	/*
-	 * We must do the balancing pass before put_prev_task(), such
-	 * that when we release the rq->lock the task is in the same
-	 * state as before we took rq->lock.
-	 *
-	 * We can terminate the balance pass as soon as we know there is
-	 * a runnable task of @class priority or higher.
-	 */
-	for_active_class_range(class, start_class, &idle_sched_class) {
-		if (class->balance && class->balance(rq, rf))
-			break;
-	}
-}
-
 /*
  * Pick up the highest-prio task:
  */
@@ -6241,8 +6222,6 @@ __pick_next_task(struct rq *rq, struct r
 	}
 
 restart:
-	prev_balance(rq, rf);
-
 	for_each_active_class(class) {
 		p = class->pick_task(rq, rf);
 		if (unlikely(p == RETRY_TASK))
@@ -6348,8 +6327,6 @@ pick_next_task(struct rq *rq, struct rq_
 		goto out_set_next;
 	}
 
-	prev_balance(rq, rf);
-
 	smt_mask = cpu_smt_mask(cpu);
 
 restart:
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2705,7 +2705,7 @@ static void check_preempt_equal_dl(struc
 	resched_curr(rq);
 }
 
-static int balance_dl(struct rq *rq, struct rq_flags *rf)
+static void balance_dl(struct rq *rq, struct rq_flags *rf)
 {
 	/*
 	 * Note, rq->donor may change during rq lock drops,
@@ -2724,8 +2724,6 @@ static int balance_dl(struct rq *rq, str
 		pull_dl_task(rq);
 		rq_repin_lock(rq, rf);
 	}
-
-	return sched_stop_runnable(rq) || sched_dl_runnable(rq);
 }
 
 /*
@@ -2823,6 +2821,11 @@ static struct task_struct *__pick_task_d
 	struct dl_rq *dl_rq = &rq->dl;
 	struct task_struct *p;
 
+	rq_modified_begin(rq, &dl_sched_class);
+	balance_dl(rq, rf);
+	if (rq_modified_above(rq, &dl_sched_class))
+		return RETRY_TASK;
+
 again:
 	if (!sched_dl_runnable(rq))
 		return NULL;
@@ -3654,7 +3657,6 @@ DEFINE_SCHED_CLASS(dl) = {
 	.put_prev_task		= put_prev_task_dl,
 	.set_next_task		= set_next_task_dl,
 
-	.balance		= balance_dl,
 	.select_task_rq		= select_task_rq_dl,
 	.migrate_task_rq	= migrate_task_rq_dl,
 	.set_cpus_allowed       = set_cpus_allowed_dl,
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -468,12 +468,6 @@ select_task_rq_idle(struct task_struct *
 	return task_cpu(p); /* IDLE tasks as never migrated */
 }
 
-static int
-balance_idle(struct rq *rq, struct rq_flags *rf)
-{
-	return WARN_ON_ONCE(1);
-}
-
 /*
  * Idle tasks are unconditionally rescheduled:
  */
@@ -594,7 +588,6 @@ DEFINE_SCHED_CLASS(idle) = {
 	.put_prev_task		= put_prev_task_idle,
 	.set_next_task          = set_next_task_idle,
 
-	.balance		= balance_idle,
 	.select_task_rq		= select_task_rq_idle,
 	.set_cpus_allowed	= set_cpus_allowed_common,
 
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1596,7 +1596,7 @@ static void check_preempt_equal_prio(str
 	resched_curr(rq);
 }
 
-static int balance_rt(struct rq *rq, struct rq_flags *rf)
+static void balance_rt(struct rq *rq, struct rq_flags *rf)
 {
 	/*
 	 * Note, rq->donor may change during rq lock drops,
@@ -1615,8 +1615,6 @@ static int balance_rt(struct rq *rq, str
 		pull_rt_task(rq);
 		rq_repin_lock(rq, rf);
 	}
-
-	return sched_stop_runnable(rq) || sched_dl_runnable(rq) || sched_rt_runnable(rq);
 }
 
 /*
@@ -1718,14 +1716,15 @@ static struct task_struct *_pick_next_ta
 
 static struct task_struct *pick_task_rt(struct rq *rq, struct rq_flags *rf)
 {
-	struct task_struct *p;
+	rq_modified_begin(rq, &rt_sched_class);
+	balance_rt(rq, rf);
+	if (rq_modified_above(rq, &rt_sched_class))
+		return RETRY_TASK;
 
 	if (!sched_rt_runnable(rq))
 		return NULL;
 
-	p = _pick_next_task_rt(rq);
-
-	return p;
+	return _pick_next_task_rt(rq);
 }
 
 static void put_prev_task_rt(struct rq *rq, struct task_struct *p, struct task_struct *next)
@@ -2616,7 +2615,6 @@ DEFINE_SCHED_CLASS(rt) = {
 	.put_prev_task		= put_prev_task_rt,
 	.set_next_task          = set_next_task_rt,
 
-	.balance		= balance_rt,
 	.select_task_rq		= select_task_rq_rt,
 	.set_cpus_allowed       = set_cpus_allowed_common,
 	.rq_online              = rq_online_rt,
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2785,11 +2785,6 @@ struct sched_class {
 	void (*wakeup_preempt)(struct rq *rq, struct task_struct *p, int flags);
 
 	/*
-	 * schedule/pick_next_task/prev_balance: rq->lock
-	 */
-	int (*balance)(struct rq *rq, struct rq_flags *rf);
-
-	/*
 	 * schedule/pick_next_task: rq->lock
 	 */
 	struct task_struct *(*pick_task)(struct rq *rq, struct rq_flags *rf);
--- a/kernel/sched/stop_task.c
+++ b/kernel/sched/stop_task.c
@@ -15,12 +15,6 @@ select_task_rq_stop(struct task_struct *
 	return task_cpu(p); /* stop tasks as never migrate */
 }
 
-static int
-balance_stop(struct rq *rq, struct rq_flags *rf)
-{
-	return sched_stop_runnable(rq);
-}
-
 static void
 wakeup_preempt_stop(struct rq *rq, struct task_struct *p, int flags)
 {
@@ -110,7 +104,6 @@ DEFINE_SCHED_CLASS(stop) = {
 	.put_prev_task		= put_prev_task_stop,
 	.set_next_task          = set_next_task_stop,
 
-	.balance		= balance_stop,
 	.select_task_rq		= select_task_rq_stop,
 	.set_cpus_allowed	= set_cpus_allowed_common,