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

[Andrea Righi <arighi@nvidia.com>]

Hi Tejun,

On Sun, Aug 16, 2026 at 01:54:22PM -1000, Tejun Heo wrote:
> Hello,
> 
> On Sun, Aug 16, 2026 at 07:35:10PM +0200, Andrea Righi wrote:
> > +/*
> > + * Proxy resolution happens before rq->curr is switched. Queue deferred work
> > + * on the rq so that an outgoing proxy owner has cleared on_cpu by the time
> > + * reject_dsq is drained.
> > + */
> 
> This doesn't make sense to me. Deferred operatoins can only execute after
> the scheduling operations are complete no matter when they were queued. I
> can see why you'd need this to kick rejected tasks that can't be queued yet.

You're right, this comment comes from an earlier version where proxy-active
rejects did not schedule an initial deferred drain and the hook was used to
arrange the post-switch drain.

The rejection path has changed to schedule deferred work for every reject, so
queueing it from this hook is not what provides the post-switch ordering.

But the hook is still needed as a retry notification. When the initial drain
runs, the rejected task may still be on_cpu or remain the active donor if the
scheduling operation retained or selected the same proxy execution context. In
that case, scx_reenq_reject() must leave the task parked and record that another
attempt is needed. A later proxy resolution then queues that attempt.

The flag also records the outstanding drain independently of the balance
callback used to schedule it. If proxy re-pick calls zap_balance_callbacks()
before that callback runs, scx_proxy_resolved() uses the flag to queue
replacement deferred work.

I'll rewrite the comment to describe this retry role.

> 
> >  void scx_proxy_resolved(struct rq *rq)
> >  {
> > +	lockdep_assert_rq_held(rq);
> > +
> > +	if (rq->scx.flags & SCX_RQ_PROXY_REENQ)
> > +		schedule_deferred_locked(rq);
> >  }
> >  
> >  void schedule_dsq_reenq(struct scx_sched *sch, struct scx_dispatch_q *dsq,
> > @@ -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.

scx_reenq_reject() sets it when a proxy-rejected task remains on_cpu or
continues to be the active donor, scx_proxy_resolved() then can use it to queue
another attempt.

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.

> 
> >  			schedule_deferred_locked(rq);
> > +		}
> >  		return;
> >  	}
> >  
> ...
> > +static bool task_proxy_move_active(struct task_struct *p)
> 
> How about task_proxy_running_or_donating()

Ack.

> 
> > +static bool task_move_proxy_raced(struct task_struct *p)
> 
> and task_proxy_unsafe_to_move() instead?

Ack.

> 
> > +/*
> > + * Park a task whose remote transfer raced with proxy execution. Reenqueueing
> > + * from the source rq makes the task's owning scheduler choose its placement
> > + * again and preserves sub-scheduler containment.
> > + */
> > +static void scx_reject_task(struct scx_sched *sch, struct rq *rq,
> > +			    struct task_struct *p, u64 enq_flags)
> 
> and scx_proxy_reject_task().

Ack.

> 
> > +{
> > +	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.

For paths which bypass ops.enqueue(), the reason can be cleared before the
kernel placement. With that lifetime enforced in scx_do_enqueue_task(),
scx_finish_reenqueue() and the overwrite handling in the previous patch can go
away.

> 
> > +	p->scx.flags &= ~SCX_TASK_REENQ_REASON_MASK;
> > +
> > +	p->scx.holding_cpu = -1;
> > +	p->scx.flags |= SCX_TASK_REENQ_PROXY;
> > +	scx_prepare_dsq_divert(p, &enq_flags);
> > +
> > +	scx_dispatch_enqueue(sch, rq, &rq->scx.reject_dsq, p, 0, 0, enq_flags);
> > +}
> > +
> >  /**
> >   * unlink_dsq_and_switch_rq_lock() - Unlink task and switch to its rq lock
> >   * @p: target task
> > @@ -2596,6 +2667,20 @@ static bool consume_remote_task(struct scx_sched *sch, struct rq *this_rq,
> >  				struct scx_dispatch_q *dsq, struct rq *src_rq)
> >  {
> >  	if (unlink_dsq_and_switch_rq_lock(p, dsq, this_rq, src_rq)) {
> > +		/*
> > +		 * Proxy execution may have changed @p's running or
> > +		 * migration-disabled state while switching rq locks without
> > +		 * clearing holding_cpu. Park it on the source rq and let its
> > +		 * owning scheduler choose its placement again.
> > +		 */
> > +		if (unlikely(task_move_proxy_raced(p))) {
> > +			p->scx.dsq = NULL;
> > +			scx_reject_task(sch, src_rq, p,
> > +					enq_flags | SCX_ENQ_CLEAR_OPSS);
> 
> No need for line break.

Ack.

> 
> > +			switch_rq_lock(src_rq, this_rq);
> > +			return false;
> > +		}
> > +
> >  		move_remote_task_to_local_dsq(sch, p, enq_flags, src_rq, this_rq);
> >  		return true;
> >  	} else {
> > @@ -2626,6 +2711,7 @@ static struct rq *move_task_between_dsqs(struct scx_sched *sch,
> >  					 struct scx_dispatch_q *dst_dsq)
> >  {
> >  	struct rq *src_rq = task_rq(p), *dst_rq;
> > +	bool proxy_raced;
> >  
> >  	BUG_ON(src_dsq->id == SCX_DSQ_LOCAL);
> >  	lockdep_assert_held(&src_dsq->lock);
> > @@ -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.

> 
> > +			dispatch_dequeue_locked(p, src_dsq);
> > +			raw_spin_unlock(&src_dsq->lock);
> > +			scx_reject_task(sch, src_rq, p, enq_flags);
> > +			return src_rq;
> > +		}
> >  		if (src_rq != dst_rq &&
> >  		    unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
> >  			dst_dsq = find_global_dsq(sch, task_cpu(p));
> > @@ -2788,7 +2887,9 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
> >  	/* task_rq couldn't have changed if we're still the holding cpu */
> >  	if (likely(p->scx.holding_cpu == raw_smp_processor_id()) &&
> >  	    !WARN_ON_ONCE(src_rq != task_rq(p))) {
> > +		bool proxy_raced = src_rq != dst_rq && task_move_proxy_raced(p);
> >  		bool fallback = false;
> > +
> >  		/*
> >  		 * If @p is staying on the same rq, there's no need to go
> >  		 * through the full deactivate/activate cycle. Optimize by
> > @@ -2798,9 +2899,13 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
> >  			p->scx.holding_cpu = -1;
> >  			scx_dispatch_enqueue(sch, dst_rq, &dst_rq->scx.local_dsq, p,
> >  					     slice, vtime, enq_flags | SCX_ENQ_APPLY_SLICE);
> > -		} else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
> > -			p->scx.holding_cpu = -1;
> > +		} else if (unlikely(proxy_raced)) {
> 
> Ditto, I don't know what the bouncing through proxy_raced buysk. Alos, in
> this else clause src_rq != dst_rq is already established.

Ack.

> 
> >  			fallback = true;
> > +			scx_reject_task(sch, src_rq, p, enq_flags);
> > +		} else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq,
> > +							       true))) {
> 
> No need for line break.

Ack.

> 
> > +			fallback = true;
> > +			p->scx.holding_cpu = -1;
> 
> fallback = true and p->scx.holding_cpu = -1 lines are swapped compared to
> the original code making the diff look a bit confusing.

Ack.

> 
> >  static void scx_reenq_reject(struct rq *rq)
> >  {
> >  	LIST_HEAD(tasks);
> >  	struct task_struct *p, *n;
> > +	bool proxy_pending = false;
> >  
> >  	lockdep_assert_rq_held(rq);
> >  
> > -	if (list_empty(&rq->scx.reject_dsq.list))
> > +	if (list_empty(&rq->scx.reject_dsq.list)) {
> > +		rq->scx.flags &= ~SCX_RQ_PROXY_REENQ;
> >  		return;
> > +	}
> >  
> >  	/*
> > -	 * Move tasks to a private list so a task re-rejected by
> > +	 * Move ready tasks to a private list so a task re-rejected by
> >  	 * scx_do_enqueue_task() below isn't revisited this round.
> >  	 */
> >  	list_for_each_entry_safe(p, n, &rq->scx.reject_dsq.list, scx.dsq_list.node) {
> >  		u32 reason = p->scx.flags & SCX_TASK_REENQ_REASON_MASK;
> >  
> > -		/* migration_pending tasks should have bypassed to local DSQ */
> > -		WARN_ON_ONCE(p->migration_pending);
> >  		WARN_ON_ONCE(!reason);
> >  
> > +		/*
> > +		 * The affinity machinery owns placement while a migration is
> > +		 * pending and will dequeue and reactivate @p as necessary. Don't
> > +		 * return it to BPF in the meantime. This isn't a proxy-resolution
> > +		 * state and thus doesn't contribute to @proxy_pending.
> > +		 */
> > +		if (p->migration_pending) {
> > +			WARN_ON_ONCE(reason != SCX_TASK_REENQ_PROXY);
> > +			continue;
> > +		}
> > +
> > +		if (reason == SCX_TASK_REENQ_PROXY &&
> > +		    (task_on_cpu(rq, p) || task_current_donor(rq, p))) {
> > +			proxy_pending = true;
> > +			continue;
> > +		}
> 
> Can you structure the code so that it's:
> 
>         if (proxy_exec_enabled() && reason == SCX_TASK_REENQ_PROXY) {
>                 ...
>         else {
>                 WARN_ON_ONCE(p->migration_pending);
>                 ...
>         }
>

Ack.

> >  		scx_dispatch_dequeue(rq, p);
> >  		p->scx.flags |= reason;
> >  
> >  		list_add_tail(&p->scx.dsq_list.node, &tasks);
> >  	}
> >  
> > +	if (proxy_pending)
> > +		rq->scx.flags |= SCX_RQ_PROXY_REENQ;
> > +	else
> > +		rq->scx.flags &= ~SCX_RQ_PROXY_REENQ;
> 
> Hmm... this looks unconventional. Shouldn't the cosumer of the flag -
> scx_proxy_resolved() - clear it?

Agreed. I'll make scx_proxy_resolved() clear the flag before scheduling the
deferred work.

> 
> > diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
> > index 2b2dcde923600..8b0be25cda7d0 100644
> > --- a/kernel/sched/ext/internal.h
> > +++ b/kernel/sched/ext/internal.h
> > @@ -1758,6 +1758,14 @@ enum scx_enq_flags {
> >  	SCX_ENQ_SLICE_DFL	= 1LLU << 62,	/* carried slice is a default refill */
> >  };
> >  
> > +/* Strip priority and carried slice state when diverting from a local DSQ. */
> > +static inline void scx_prepare_dsq_divert(struct task_struct *p, u64 *enq_flags)
> 
> How about scx_divert_strip_flags()?

Ack.

Thanks tons for the review!
-Andrea