Re: [PATCH sched_ext/for-6.12] sched_ext: Handle cases where pick_task_scx() is called without preceding balance_scx()

[Peter Zijlstra <peterz@infradead.org>]

On Thu, Sep 05, 2024 at 11:28:58AM +0200, Peter Zijlstra wrote:
> On Wed, Sep 04, 2024 at 02:47:03PM -1000, Tejun Heo wrote:
> > sched_ext dispatches tasks from the BPF scheduler from balance_scx() and
> > thus every pick_task_scx() call must be preceded by balance_scx(). While
> > this usually holds, there are rare cases where a higher sched class's
> > balance() returns true indicating that it has tasks to run on the CPU and
> > thus terminating balance() calls but fails to actually find the next task to
> > run when pick_task() is called. 
> 
> Oh cute. Which class in particular did you see this do?
> 
> Looking at balance_fair() / sched_balance_newidle() I suppose we could
> verify we actually have a runnable task once we've re-acquired the
> rq-lock and have pulled_task > 0.
> 
> 
> Tightening all that up would probably be better than trying to deal with
> the fallout like this, hmm?

Something like so. Haven't yet looked at the rt/dl classes.

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 11e890486c1b..7db42735d504 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -12716,6 +12716,12 @@ static int sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
 	if (this_rq->cfs.h_nr_running && !pulled_task)
 		pulled_task = 1;
 
+	/*
+	 * We pulled a task, but it got stolen before we re-acquired rq->lock.
+	 */
+	if (!this_rq->cfs.h_nr_running && pulled_task)
+		pulled_task = 0;
+
 	/* Is there a task of a high priority class? */
 	if (this_rq->nr_running != this_rq->cfs.h_nr_running)
 		pulled_task = -1;