Re: [PATCH sched_ext/for-6.12-fixes] sched_ext: Call __balance_callbacks() from __scx_task_iter_rq_unlock()

[Peter Zijlstra <peterz@infradead.org>]

On Fri, Nov 01, 2024 at 12:36:34AM +0100, Peter Zijlstra wrote:
> On Wed, Oct 30, 2024 at 11:41:39AM -1000, Tejun Heo wrote:
> 
> > --- a/kernel/sched/ext.c
> > +++ b/kernel/sched/ext.c
> > @@ -1315,6 +1315,8 @@ static void scx_task_iter_start(struct s
> >  static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
> >  {
> >  	if (iter->locked) {
> > +		/* ->switched_from() may have scheduled balance callbacks */
> > +		__balance_callbacks(iter->rq);
> >  		task_rq_unlock(iter->rq, iter->locked, &iter->rf);
> >  		iter->locked = NULL;
> >  	}
> 
> I think you need to unpin/repin around it. The balance callbacks like to
> drop rq->lock at times.

Maybe something like so.. I'm not sure it's an improvement.

---
 kernel/sched/core.c  | 12 +++++-------
 kernel/sched/ext.c   |  1 +
 kernel/sched/sched.h |  9 +++++++++
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 5de31c312189..c826745eedef 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5039,7 +5039,7 @@ struct balance_callback *splice_balance_callbacks(struct rq *rq)
 	return __splice_balance_callbacks(rq, true);
 }
 
-static void __balance_callbacks(struct rq *rq)
+void __balance_callbacks(struct rq *rq)
 {
 	do_balance_callbacks(rq, __splice_balance_callbacks(rq, false));
 }
@@ -6721,9 +6721,8 @@ static void __sched notrace __schedule(int sched_mode)
 		/* Also unlocks the rq: */
 		rq = context_switch(rq, prev, next, &rf);
 	} else {
-		rq_unpin_lock(rq, &rf);
-		__balance_callbacks(rq);
-		raw_spin_rq_unlock_irq(rq);
+		rf.balance = true;
+		rq_unlock_irq(rq, &rf);
 	}
 }
 
@@ -7217,9 +7216,8 @@ void rt_mutex_setprio(struct task_struct *p, struct task_struct *pi_task)
 	/* Avoid rq from going away on us: */
 	preempt_disable();
 
-	rq_unpin_lock(rq, &rf);
-	__balance_callbacks(rq);
-	raw_spin_rq_unlock(rq);
+	rf.balance = true;
+	rq_unlock(rq, &rf);
 
 	preempt_enable();
 }
diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 7bc4fb7f3926..8e5b31983f37 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1314,6 +1314,7 @@ static void scx_task_iter_start(struct scx_task_iter *iter)
 static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
 {
 	if (iter->locked) {
+		iter->rf.balance = true;
 		task_rq_unlock(iter->rq, iter->locked, &iter->rf);
 		iter->locked = NULL;
 	}
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 0aa6da2633aa..5a39ad0fa574 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1704,6 +1704,7 @@ static inline void rq_clock_stop_loop_update(struct rq *rq)
 struct rq_flags {
 	unsigned long flags;
 	struct pin_cookie cookie;
+	bool balance;
 #ifdef CONFIG_SCHED_DEBUG
 	/*
 	 * A copy of (rq::clock_update_flags & RQCF_UPDATED) for the
@@ -1739,6 +1740,8 @@ static inline void rq_pin_lock(struct rq *rq, struct rq_flags *rf)
 #endif
 }
 
+extern void __balance_callbacks(struct rq *rq);
+
 static inline void rq_unpin_lock(struct rq *rq, struct rq_flags *rf)
 {
 #ifdef CONFIG_SCHED_DEBUG
@@ -1747,6 +1750,12 @@ static inline void rq_unpin_lock(struct rq *rq, struct rq_flags *rf)
 #endif
 
 	lockdep_unpin_lock(__rq_lockp(rq), rf->cookie);
+#ifdef CONFIG_SMP
+	if (rf->balance) {
+		rf->balance = false;
+		__balance_callbacks(rq);
+	}
+#endif
 }
 
 static inline void rq_repin_lock(struct rq *rq, struct rq_flags *rf)