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

[sashiko-bot@kernel.org]

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.

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?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261017152226.3690249-1-arighi@nvidia.com?part=1