Re: [PATCH v2 sched_ext/for-7.3-fixes] sched_ext: Don't leave balance callbacks queued across a pick retry

[Andrea Righi <arighi@nvidia.com>]

On Sat, Oct 17, 2026 at 03:40:58PM +0000, sashiko-bot@kernel.org wrote:
> Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
> - [High] sched_ext: balance callbacks queued before core-sched pick is final
> --
> 
> commit 859b8e499322d8a16b0b359431e360cdd5ae7494
> Author: Andrea Righi <arighi@nvidia.com>
> 
> sched_ext: Don't leave balance callbacks queued across a pick retry
> 
> This commit moves the queueing of rq->scx.deferred_bal_cb and
> rq->scx.kick_sync_bal_cb to the exit path of do_pick_task_scx() so that
> balance callbacks are only queued when do_pick_task_scx() returns a final
> task rather than RETRY_TASK or a NULL dl_server pick.
> 
> > diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
> > index 0403c46342813..e19961164cb01 100644
> > --- a/kernel/sched/ext/ext.c
> > +++ b/kernel/sched/ext/ext.c
> 
> [ ... ]
> 
> > @@ -3342,16 +3331,8 @@ static enum scx_dsp_verdict dispatch_core_pick(struct rq *rq, struct rq_flags *r
> >  
> >  	verdict = dispatch_one(rq, prev);
> >  
> > -	if (cpu_of(rq) == smp_processor_id()) {
> > -		maybe_queue_balance_callback(rq);
> > -
> > -		/* see dispatch_pick() */
> > -		if (unlikely(rq->scx.kick_sync_pending)) {
> > -			rq->scx.kick_sync_pending = false;
> > -			queue_balance_callback(rq, &rq->scx.kick_sync_bal_cb,
> > -					       kick_sync_wait_bal_cb);
> > -		}
> > -	} else if (unlikely(rq->scx.flags & SCX_RQ_BAL_CB_PENDING)) {
> > +	if (cpu_of(rq) != smp_processor_id() &&
> > +	    unlikely(rq->scx.flags & SCX_RQ_BAL_CB_PENDING)) {
> >  		/*
> >  		 * Balance callbacks must run in the context that queued them,
> >  		 * so they can't be queued on another CPU's rq. Run the deferred
> > @@ -3394,8 +3375,10 @@ do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
> >  	else
> >  		verdict = dispatch_pick(rq, rf, prev);
> >  
> > -	if (verdict == SCX_DSP_RETRY)
> > -		return RETRY_TASK;
> > +	if (verdict == SCX_DSP_RETRY) {
> > +		p = RETRY_TASK;
> > +		goto out;
> > +	}
> 
> [ ... ]
> 
> > @@ -3437,6 +3422,38 @@ do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
> >  		}
> >  	}
> >  
> > +out:
> > +	if (cpu_of(rq) != smp_processor_id())
> > +		return p;
> > +
> > +	/*
> > +	 * Balance callbacks must be consumed before the rq lock is released, so
> > +	 * queue them only when the pick is final. On RETRY_TASK, or after a
> > +	 * NULL pick from the dl_server, the core continues picking and can
> > +	 * release the rq lock (e.g., prev_balance() or another dispatch). In
> > +	 * that case, run the deferred work from irq_work, which doesn't depend
> > +	 * on the current rq lock section, and leave the kick sync wait pending
> > +	 * for the next pick.
> > +	 */
> > +	if (p != RETRY_TASK && (p || !force_scx)) {
> > +		maybe_queue_balance_callback(rq);
> > +
> > +		/*
> > +		 * Defer to a balance callback which can drop rq lock and enable
> > +		 * IRQs. Waiting directly in the pick path would deadlock against
> > +		 * CPUs sending us IPIs (e.g., TLB flushes) while we wait for
> > +		 * them.
> > +		 */
> > +		if (unlikely(rq->scx.kick_sync_pending)) {
> > +			rq->scx.kick_sync_pending = false;
> > +			queue_balance_callback(rq, &rq->scx.kick_sync_bal_cb,
> > +					       kick_sync_wait_bal_cb);
> > +		}
> > +	} else if (rq->scx.flags & SCX_RQ_BAL_CB_PENDING) {
> > +		rq->scx.flags &= ~SCX_RQ_BAL_CB_PENDING;
> > +		schedule_deferred(rq);
> > +	}
> > +
> >  	return p;
> >  }
> 
> [Severity: High]
> Under CONFIG_SCHED_CORE, can a non-RETRY_TASK return from
> do_pick_task_scx() on the local rq still be followed by further picks that
> release the shared core lock before the core-wide pick is final?
> 
> In kernel/sched/core.c:pick_next_task(), a single shared core lock is held
> across all SMT siblings, and pick_task() can be called multiple times after
> the local rq pick succeeds:
> 
> 1. In the !need_sync fastpath, pick_task(rq, rf) runs on the local rq:
> 
> kernel/sched/core.c:pick_next_task() {
> 	if (!need_sync) {
> 		next = pick_task(rq, rf);
> 		...
> 		if (!next->core_cookie) {
> 			...
> 			goto out_set_next;
> 		}
> 	}
> }
> 
> If do_pick_task_scx() returns a cookied SCX task (next->core_cookie != 0),
> the out: block sees p != RETRY_TASK && (p || !force_scx) and queues
> rq->scx.deferred_bal_cb and/or rq->scx.kick_sync_bal_cb on
> rq->balance_callback. Because next->core_cookie is non-zero,
> pick_next_task() then falls through into the for_each_cpu_wrap(i, smt_mask,
> cpu) loop to re-pick across all siblings.
> 
> 2. In the for_each_cpu_wrap(i, smt_mask, cpu) loop in pick_next_task(),
> iteration starts at i = cpu (where cpu_of(rq_i) == smp_processor_id()):
> 
> kernel/sched/core.c:pick_next_task() {
> 	for_each_cpu_wrap(i, smt_mask, cpu) {
> 		rq_i = cpu_rq(i);
> 		...
> 		p = pick_task(rq_i, rf);
> 		if (unlikely(p == RETRY_TASK)) {
> 			...
> 			goto restart;
> 		}
> 	}
> }
> 
> Once do_pick_task_scx() on the local rq succeeds and queues the balance
> callbacks on rq->balance_callback, the loop proceeds to call
> pick_task(rq_i, rf) on the remaining SMT siblings (i != cpu). When
> dispatch_core_pick(rq_i, rf, prev) runs for a sibling rq_i,
> dispatch_one(rq_i, prev) or run_deferred(rq_i) can drop the shared core
> lock and return SCX_DSP_RETRY, causing pick_next_task() to jump to restart
> with rq->balance_callback already queued on the local rq. Any CPU acquiring
> the core lock during that window will trigger the warning in
> assert_balance_callbacks_empty() via rq_pin_lock() or flush the callback
> early.

I'm planning to fix this with a follow-up fix, it was also pointed out by Tejun
in his review of v1.

> 
> 3. Also in dispatch_core_pick(), if an SMT sibling already has an in-flight
> dispatch on the local rq:
> 
> kernel/sched/ext/ext.c:dispatch_core_pick() {
> 	/* another dispatch is in flight on @rq, let that handle it */
> 	if (rq->scx.flags & SCX_RQ_IN_DISPATCH)
> 		return SCX_DSP_NONE;
> }
> 
> Previously, returning SCX_DSP_NONE here bypassed callback queueing in
> dispatch_core_pick(). With the callback handling moved to the out: label in
> do_pick_task_scx(), won't this early return now fall through to out: and
> consume SCX_RQ_BAL_CB_PENDING and rq->scx.kick_sync_pending while the
> sibling's dispatch_one() on rq is still in flight?

Hm... yeah this looks legit for SCX_RQ_BAL_CB_PENDING, I'll fix this in v3.

-Andrea