Re: [PATCH V2 1/1] sched_ext: Fix spurious WARN on stale ops_state in ops_dequeue()

[Andrea Righi <arighi@nvidia.com>]

Hi Samuele,

On Thu, May 21, 2026 at 12:59:11PM +0200, Samuele Mariotti wrote:
> ops_dequeue() can race with finish_dispatch() and spuriously trigger the
> "queued task must be in BPF scheduler's custody" warning.
> 
> ops_dequeue() snapshots p->scx.ops_state via atomic_long_read_acquire()
> and then, in the SCX_OPSS_QUEUED arm, asserts that SCX_TASK_IN_CUSTODY
> is set. The two reads are not atomic w.r.t. a concurrent
> finish_dispatch() running on another CPU:
> 
> CPU 1                                    CPU 2
> =====                                    =====
>                                          dequeue_task_scx()
>                                            ops_dequeue()
>                                              opss = read_acquire(ops_state)
>                                                   = SCX_OPSS_QUEUED
> finish_dispatch()
>   cmpxchg ops_state:
>     SCX_OPSS_QUEUED -> SCX_OPSS_DISPATCHING  [succeeds]
>   dispatch_enqueue(SCX_DSQ_GLOBAL,
>                    SCX_ENQ_CLEAR_OPSS)
>     call_task_dequeue()
>       p->scx.flags &= ~SCX_TASK_IN_CUSTODY
>                                              WARN_ON_ONCE(!(p->scx.flags &
>                                                      SCX_TASK_IN_CUSTODY))
>                                             /* opss is stale: QUEUED,
>                                              * but task already claimed */
>     set_release(ops_state, SCX_OPSS_NONE)
> 
> The race has been observed via two distinct call chains: the most common
> goes through sched_setaffinity(), a rarer variant through
> sched_change_begin().
> 
> For SCX_DSQ_GLOBAL / SCX_DSQ_BYPASS, dispatch_enqueue() clears
> SCX_TASK_IN_CUSTODY before clearing ops_state to SCX_OPSS_NONE
> (intentional, to avoid concurrent non-atomic RMW of p->scx.flags against
> ops_dequeue()). The window between those two writes is exactly what
> ops_dequeue() observes as "QUEUED without custody".
> 
> The observed state is not actually inconsistent, it just means CPU 1 has
> already claimed the task and the QUEUED value held by CPU 2 is stale.
> Re-read ops_state in that case; the next read is guaranteed to return
> SCX_OPSS_DISPATCHING or SCX_OPSS_NONE, both of which exit the switch
> cleanly. The retry is bounded: once IN_CUSTODY is cleared, ops_state has
> already advanced past QUEUED for this dispatch cycle, and a fresh QUEUED
> would require re-enqueue under p's rq lock, which CPU 2 holds.
> 
> Changes in v2:
> - Use READ_ONCE() for p->scx.flags to ensure fresh reads and prevent
>   compiler reordering in the lockless path
> - Add cpu_relax() to reduce power consumption and improve performance
>   during the spin-wait
> - Use unlikely() to optimize branch prediction for the common case
> - Expand the in-code comment to document the race condition and
>   bounded retry guarantee
> 
> Fixes: ebf1ccff79c4 ("sched_ext: Fix ops.dequeue() semantics")
> Suggested-by: Andrea Righi <arighi@nvidia.com>
> Signed-off-by: Samuele Mariotti <smariotti@disroot.org>
> Signed-off-by: Paolo Valente <paolo.valente@unimore.it>

Looks good to me.

Reviewed-by: Andrea Righi <arighi@nvidia.com>

Thanks,
-Andrea

> ---
>  kernel/sched/ext.c | 17 +++++++++++++++--
>  1 file changed, 15 insertions(+), 2 deletions(-)
> 
> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> index 547ca398f646..c1762420cc35 100644
> --- a/kernel/sched/ext.c
> +++ b/kernel/sched/ext.c
> @@ -2078,6 +2078,7 @@ static void ops_dequeue(struct rq *rq, struct task_struct *p, u64 deq_flags)
>  	/* dequeue is always temporary, don't reset runnable_at */
>  	clr_task_runnable(p, false);
>  
> +retry:
>  	/* acquire ensures that we see the preceding updates on QUEUED */
>  	opss = atomic_long_read_acquire(&p->scx.ops_state);
>  
> @@ -2091,8 +2092,20 @@ static void ops_dequeue(struct rq *rq, struct task_struct *p, u64 deq_flags)
>  		 */
>  		BUG();
>  	case SCX_OPSS_QUEUED:
> -		/* A queued task must always be in BPF scheduler's custody */
> -		WARN_ON_ONCE(!(p->scx.flags & SCX_TASK_IN_CUSTODY));
> +		/*
> +		 * A queued task must always be in BPF scheduler's custody. If
> +		 * SCX_TASK_IN_CUSTODY is clear, finish_dispatch() on another
> +		 * CPU has already passed call_task_dequeue() (which clears the
> +		 * flag), but has not yet written SCX_OPSS_NONE. That final
> +		 * store does not require this rq's lock, so retrying with
> +		 * cpu_relax() is bounded: we will observe NONE (or DISPATCHING,
> +		 * handled by the fallthrough) on a subsequent iteration.
> +		 */
> +		if (unlikely(!(READ_ONCE(p->scx.flags) & SCX_TASK_IN_CUSTODY))) {
> +			cpu_relax();
> +			goto retry;
> +		}
> +
>  		if (atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
>  					    SCX_OPSS_NONE))
>  			break;
> -- 
> 2.54.0
>