Re: [PATCH 12/17] sched_ext: Handle proxy-exec races in remote DSQ transfers

[Andrea Righi <arighi@nvidia.com>]

On Mon, Aug 17, 2026 at 07:17:23AM -1000, Tejun Heo wrote:
> Hello,
> 
> On Mon, Aug 17, 2026 at 09:15:22AM +0200, Andrea Righi wrote:
> > > > @@ -1530,12 +1539,17 @@ static void rq_owned_post_enq(struct scx_sched *sch, struct rq *rq,
> > > >  	call_task_dequeue(sch, rq, p, 0);
> > > >  
> > > >  	/*
> > > > -	 * Only local inserts get the wakeup treatment below. Rejects kick the
> > > > -	 * deferred reenq and rescue parks are paced by the rescue timer.
> > > > +	 * Only local inserts get the wakeup treatment below. Proxy-active tasks
> > > > +	 * and rescuees remain parked until their respective resolution paths.
> > > > +	 * Other rejects can be reenqueued immediately.
> > > >  	 */
> > > >  	if (unlikely(dsq->id != SCX_DSQ_LOCAL)) {
> > > > -		if (dsq->id == SCX_DSQ_REJECT)
> > > > +		if (dsq->id == SCX_DSQ_REJECT) {
> > > > +			if ((p->scx.flags & SCX_TASK_REENQ_REASON_MASK) ==
> > > > +			    SCX_TASK_REENQ_PROXY)
> > > > +				rq->scx.flags |= SCX_RQ_PROXY_REENQ;
> > > 
> > > and if my reading above is correct, this wouldn't be necessary, right?
> > 
> > The flag isn't needed to provide the post-switch ordering, but it is needed to
> > record that reject_dsq still needs another drain.
> 
> But wouldn't it be able to use the same schedule_deferred_locked() call like
> SCX_DSQ_REJECT case so that the PROXY_REENQ flag is only used for retry
> cases (and maybe renamed accordingly)?

Yes. The initial insertion into reject_dsq can use schedule_deferred_locked()
without setting a proxy-specific flag. Then we can set the flag only when
scx_reenq_reject() actually skips a task because it is still running or
donating. How about renaming to SCX_RQ_PROXY_RETRY to refelct the new meaning?

> 
> ...
> > It also preserves the outstanding drain if proxy re-pick calls
> > zap_balance_callbacks() before the initially queued callback runs. In
> > that case, the task is already parked on reject_dsq, but the flag lets
> > scx_proxy_resolved() queue replacement deferred work.
> 
> I think this is a generic problem with core scheduling. SCX assumes that
> deferred scheduilng always runs but core-sched can zap them. We probably
> need to address this directly using a similar but generic deferred work
> pending flag.

Agreed. This isn't specific to proxy-rejected tasks, any sched_ext deferred work
queued through a balance callback can be lost if core scheduling zaps the
callback.

Maybe we can introduce a new generic sched_ext deferred work pending flag, kept
set until the deferred work actually runs and we can use it to rearm it after
callback cancellation?

Probably something to do for another patch series...

> 
> > > > +{
> > > > +	lockdep_assert_rq_held(rq);
> > > > +	WARN_ON_ONCE((p->scx.flags & SCX_TASK_REENQ_REASON_MASK) &&
> > > > +		     !(enq_flags & SCX_ENQ_REENQ));
> > > 
> > > In the previous patch, is it possible to clear reason before reenqueueing
> > > it and get rid of overwrite cases or does that quite not work out?
> > 
> > Yes, I think we can get rid of the overwrite cases. There is one constraint,
> > though: the reason must remain set while ops.enqueue() runs, because BPF reads
> > it from p->scx.flags.
> > 
> > I'll move the clearing into scx_do_enqueue_task() so that it happens immediately
> > after ops.enqueue() returns and before resolving any direct dispatch requested
> > by the callback. A rejection caused by that new placement will then see a clear
> > reason field and can set its own reason without overwriting the previous one.
> 
> I haven't really thought through it so please do whatever that makes the
> code least ugly.

Ok, probably moving the clearing into scx_do_enqueue_task(), after ops.enqueue()
returns but before resolving the resulting placement is the least ugly way...
This keeps the reason visible to BPF while ensuring that a subsequent rejection
starts with a clear reason field.

> 
> > > > @@ -2633,6 +2719,19 @@ static struct rq *move_task_between_dsqs(struct scx_sched *sch,
> > > >  
> > > >  	if (dst_dsq->id == SCX_DSQ_LOCAL) {
> > > >  		dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
> > > > +		/*
> > > > +		 * Unlike the rq-lock handoff paths, @src_rq has been locked
> > > > +		 * throughout this operation. Only active proxy state can race the
> > > > +		 * move here; let the enforcing check below diagnose an ordinary
> > > > +		 * migration-disabled task.
> > > > +		 */
> > > > +		proxy_raced = src_rq != dst_rq && task_proxy_move_active(p);
> > > > +		if (unlikely(proxy_raced)) {
> > > 
> > > I don't know what bouncing through the local var buys. Out of curiosity, if
> > > task_move_proxy_raced() is used here, does something break or is it just to
> > > avoid unnecessary tests?
> > 
> > Yeah, the local variable doesn't buy anything, I'll remove it.
> > 
> > Using task_move_proxy_raced() (aka task_proxy_unsafe_to_move() after the rename)
> > would change the behavior for an ordinary migration-disabled task. This path
> > holds src_rq, so that state isn't a lock-handoff race and should still reach
> > task_can_run_on_remote_rq(..., true), which diagnoses the invalid BPF-directed
> > migration. Treating it as a proxy rejection instead could hide that error and
> > cause repeated reenqueues.
> 
> That sounds like something which is worth noting in the comment. It's subtle
> that this site needs a different set of conditions.

Agreed. I'll remove the local variable and expand the comment.

Thanks for all the help! Really appreciated!
-Andrea