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

[Andrea Righi <arighi@nvidia.com>]

Hi Tejun,

On Mon, Aug 03, 2026 at 11:36:39AM -1000, Tejun Heo wrote:
> Hello,
> 
> On Tue, Jul 28, 2026 at 05:43:28PM +0200, Andrea Righi wrote:
> ...
> > @@ -135,6 +135,9 @@ enum scx_ent_flags {
> > @@ -145,6 +148,8 @@ enum scx_ent_flags {
> >  	SCX_TASK_REENQ_IMMED	= 2 << SCX_TASK_REENQ_REASON_SHIFT,
> >  	SCX_TASK_REENQ_PREEMPTED = 3 << SCX_TASK_REENQ_REASON_SHIFT,
> >  	SCX_TASK_REENQ_CAP	= 4 << SCX_TASK_REENQ_REASON_SHIFT,
> > +	SCX_TASK_REENQ_MIGRATION_DISABLED = 5 << SCX_TASK_REENQ_REASON_SHIFT,
> > +	SCX_TASK_REENQ_PROXY	= 6 << SCX_TASK_REENQ_REASON_SHIFT,
> 
> Given that MIGRATION_DISABLED can only happen with proxy execution, it may
> be better to name it accordingly. Maybe that's too long. Do we have to
> distinguish between PROXY and MIGRATION_DISABLED? We can count both as
> PROXY, no?

Ack, both conditions have the same outcome: abort the BPF-directed remote
transfer, park the task on the source rq and return it to its owning scheduler.

I'll drop SCX_TASK_REENQ_MIGRATION_DISABLED and use SCX_TASK_REENQ_PROXY for all
of these cases.

> 
> > +/*
> > + * Proxy execution can change @p's execution and migration-disabled state
> > + * without touching its DSQ entry or clearing holding_cpu. Check those states
> > + * with @p's rq locked. Without proxy execution, the holding_cpu handshake is
> > + * sufficient and this must not affect the existing migration path.
> > + */
> > +static u32 task_move_reject_reason(struct task_struct *p)
> > +{
> > +	struct rq *src_rq = task_rq(p);
> > +
> > +	lockdep_assert_rq_held(src_rq);
> > +
> > +	if (!sched_proxy_exec())
> > +		return SCX_TASK_REENQ_NONE;
> > +
> > +	/* @p may be rq->curr under another task's scheduling context. */
> > +	if (task_on_cpu(src_rq, p))
> > +		return SCX_TASK_REENQ_PROXY;
> > +
> > +	/*
> > +	 * Reject only BPF-directed migration. proxy_migrate_task() may still
> > +	 * move a blocked donor's scheduling context to its lock owner's CPU.
> > +	 */
> 
> I have a hard time understanding this comment. Can you explain a bit more?

The distinction here is between normal task placement and proxy donor migration.

A BPF-directed transfer to a remote local DSQ performs a normal task migration
and, as we know, migration-disabled tasks can't migrate.

With proxy-exec, instead, donors can always be migrated towards the mutex
owner's CPU, because they don't actually execute there; it's the owner that runs
using the donated scheduling context; wake_cpu is preserved so that the donor
returns to its original CPU when the owner releases the mutex. For that reason,
proxy_migrate_task() should be allowed to migrate even when the donor is
migration-disabled or when the target CPU outside its affinity mask.

I'll expand the comment to make this more clear.

> 
> > +	if (is_migration_disabled(p))
> > +		return SCX_TASK_REENQ_MIGRATION_DISABLED;
> > +
> > +	/* Don't move an active scheduling context off its source rq. */
> > +	if (task_current_donor(src_rq, p))
> > +		return SCX_TASK_REENQ_PROXY;
> > +
> > +	return SCX_TASK_REENQ_NONE;
> > +}
> > +
> > +/*
> > + * 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, u32 reason)
> > +{
> > +	lockdep_assert_rq_held(rq);
> > +	WARN_ON_ONCE(reason != SCX_TASK_REENQ_MIGRATION_DISABLED &&
> > +		     reason != SCX_TASK_REENQ_PROXY);
> > +	WARN_ON_ONCE(p->scx.reject_reason);
> > +
> > +	p->scx.holding_cpu = -1;
> > +	p->scx.reject_reason = reason;
> > +	p->scx.flags &= ~SCX_TASK_IMMED;
> > +	enq_flags &= ~(SCX_ENQ_IMMED | SCX_ENQ_PREEMPT);
> 
> SCX_ENQ_HEAD likely needs clearing too. After applying the rescue patchset,
> scx_resolve_local_dsq() has:

Right, we're missing SCX_ENQ_HEAD here.

> 
> 	/*
> 	 * Diverting to rescue or reject, neither of which honors IMMED, PREEMPT
> 	 * or HEAD - a diversion has no priority and IMMED is not allowed on
> 	 * non-local DSQs. Strip the enq and task flags along with the slice.
> 	 */
> 	*enq_flags &= ~(SCX_ENQ_IMMED | SCX_ENQ_PREEMPT | SCX_ENQ_HEAD |
> 			SCX_ENQ_APPLY_SLICE | SCX_ENQ_SLICE_DFL);
> 	p->scx.flags &= ~SCX_TASK_IMMED;
> 
> We should probably factor that out and use that whenever we're diverting.

I'll factor this into a common helper and use it from both
scx_resolve_local_dsq() and the proxy rejection path.

> 
> > @@ -2533,6 +2617,15 @@ 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);
> > +		reject_reason = src_rq != dst_rq ?
> > +			task_move_reject_reason(p) : SCX_TASK_REENQ_NONE;
> > +		if (unlikely(reject_reason)) {
> > +			dispatch_dequeue_locked(p, src_dsq);
> > +			raw_spin_unlock(&src_dsq->lock);
> > +			scx_reject_task(sch, src_rq, p, enq_flags,
> > +					reject_reason);
> > +			return src_rq;
> > +		}
> 
> I wonder whether it'd be cleaner if we just mark the task for rejection and
> then make scx_resolve_local_dsq() resolve that to reject dsq instead of
> directly inserting from each site. I guess the problem is that the rejection
> has to be on the source rq, not the destination one. I hope there's a neater
> way to do this.

Right, the source rq requirement is the complication. By the time
scx_resolve_local_dsq() is normally called, it is resolving an insertion
relative to the requested destination rq. In this case the transfer must be
aborted before changing task_rq() and the task must be parked on the source rq,
so that its owning scheduler/sub-scheduler containment are preserved.

I can simplify this by making the race predicate boolean, making the PROXY
reason implicit in the rejection helper, and using the common helper mentioned
above. The individual transfer paths will still need a small branch, because
they reach this point with different DSQ and rq locking states.

I don't see a better way to do this...

> 
> >  /*
> > - * Drain @rq->scx.reject_dsq and reenqueue each task so that its owning BPF
> > - * scheduler chooses placement again.
> > + * Drain ready tasks from @rq->scx.reject_dsq and reenqueue them so that their
> > + * owning BPF schedulers choose placement again. Proxy-active tasks remain
> > + * parked until proxy resolution schedules another drain after switch-out.
> >   *
> >   * A task can be re-rejected repeatedly. Reenqueues are bounded per task by
> >   * SCX_REENQ_MAX_REPEAT in scx_do_enqueue_task(), which ejects the owning
> > - * scheduler. The private list below prevents a task from being revisited in
> > - * the same round.
> > + * scheduler.
> >   */
> >  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.reject_reason;
> >  
> >  		/* migration_pending tasks should have bypassed to local DSQ */
> > -		if (WARN_ON_ONCE(p->migration_pending))
> > -			continue;
> > +		WARN_ON_ONCE(p->migration_pending);
> 
> Why this change? Also, wouldn't this now be allowed to happen? During task
> move, if it hits migration_pending due to proxy exec, the task would be put
> on reject DSQ, right? Reenq can trigger from other sources before
> migration_pending is cleared and then would see the migration_pending set.
> Shouldn't it just continue?

Yes, it should continue. Dropping the continue is wrong.

migration_pending is separate from proxy donor migration. It belongs to the
generic affinity machinery, whereas proxy_migrate_task() moves only the donor's
scheduling context and does not use migration_pending.

If a task on the proxy reject DSQ has migration_pending set, the affinity
machinery owns its placement and will dequeue/reactivate it as necessary.

> 
> >  		if (WARN_ON_ONCE(!reason))
> >  			continue;
> >  
> > +		if (reason == SCX_TASK_REENQ_PROXY &&
> > +		    (task_on_cpu(rq, p) || task_current_donor(rq, p))) {
> > +			proxy_pending = true;
> > +			continue;
> > +		}
> 
> ie. Shouldn't migration_pending() be one of the || conditions above? And
> then if it's still migraiton_pending, that should trigger a warning?

I think it should be handled by a separate early continue rather than included
in this condition.

task_on_cpu() and task_current_donor() are resolved by the proxy switch-out
path, so they set proxy_pending and arrange another drain from
scx_proxy_resolved().

migration_pending is resolved by the affinity migration machinery, not by
scx_proxy_resolved(). Treating it as proxy_pending would associate it with the
wrong completion event.

Maybe we can structure it as following?

 	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;
  	}

This keeps the task parked while the affinity request is pending without
restricting proxy_migrate_task(), which remains free to move blocked donors
towards their mutex owners.

Thanks,
-Andrea