Re: [PATCH] sched_ext: Add cookie API for early qseq capture

[sashiko-bot@kernel.org]

Thank you for your contribution! Sashiko AI review found 3 potential issue(s) to consider:
- [High] Unprotected task state mutation causing data race on slice
- [High] Missing acquire memory barrier in scx_bpf_task_get_cookie() defeats sequence-lock protection
- [Medium] Missing API variant for PRIQ (vtime-ordered) DSQ insertion
--

commit c9d30798b70b7101ea880a0f906790c5875384ac
Author: Cheng-Yang Chou <yphbchou0911@gmail.com>

sched_ext: Add cookie API for early qseq capture

This commit introduces two new kfuncs, scx_bpf_task_get_cookie() and
scx_bpf_dsq_insert_with_cookie(), to allow BPF schedulers to capture a
task's qseq before performing pre-dispatch checks and safely insert it
if the sequence hasn't changed.

Link: https://lore.kernel.org/r/20260203230639.1259869-1-arighi@nvidia.com/

> diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
> index 5600618dcef8e..5ee74c964a971 100644
> --- a/kernel/sched/ext.c
> +++ b/kernel/sched/ext.c

[ ... ]

> +/**
> + * scx_bpf_task_get_cookie - Get an opaque dispatch cookie for a task
> + * @p: task_struct to read cookie from
> + *
> + * Returns an opaque u64 cookie encoding @p's current qseq. Call this
> + * before pre-dispatch validity checks and pass the result to
> + * scx_bpf_dsq_insert_with_cookie() to extend the qseq protection window.
> + *
> + * For schedulers that do not implement properly synchronized dequeue only.
> + */
> +__bpf_kfunc u64 scx_bpf_task_get_cookie(struct task_struct *p)
> +{
> +	return atomic_long_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK;

Should this use atomic_long_read_acquire() instead of a relaxed atomic read?

Since this conceptually starts a sequence-lock read-side critical section,
a relaxed read doesn't emit hardware memory barriers. On weakly-ordered
architectures, the CPU could reorder the subsequent BPF memory loads (the
validity checks) to execute before the ops_state read.

If a concurrent dequeue/re-enqueue race occurs, the BPF checks could be
evaluated using stale task data, but the delayed ops_state read captures
the new qseq. The protection window is defeated, and finish_dispatch()
would accept the dispatch despite violated constraints.

> +}
> +
> +/**
> + * scx_bpf_dsq_insert_with_cookie - Insert a task using an early-captured cookie
> + * @p: task_struct to insert
> + * @dsq_id: DSQ to insert into
> + * @enq_flags: SCX_ENQ_*
> + * @cookie: cookie from scx_bpf_task_get_cookie()
> + * @aux: implicit BPF argument
> + *
> + * Like scx_bpf_dsq_insert() with slice=0, but uses @cookie's qseq instead
> + * of re-reading ops_state at insert time. A stale cookie causes
> + * finish_dispatch() to silently discard the dispatch. Use
> + * scx_bpf_task_set_slice() to set a non-default slice.
> + *
> + * Returns %true on success, %false on failure.
> + */
> +__bpf_kfunc bool scx_bpf_dsq_insert_with_cookie(struct task_struct *p,
> +						 u64 dsq_id, u64 enq_flags,
> +						 u64 cookie,
> +						 const struct bpf_prog_aux *aux)
> +{
> +	struct scx_sched *sch;
> +
> +	guard(rcu)();
> +	sch = scx_prog_sched(aux);
> +	if (unlikely(!sch))
> +		return false;
> +
> +	if (!scx_dsq_insert_preamble(sch, p, dsq_id, &enq_flags))
> +		return false;
> +
> +	p->scx.slice = p->scx.slice ?: 1;

Does modifying the slice here create a data race if the task is currently
running on another CPU?

Because this kfunc is meant for scenarios where the task reference might be
stale due to a concurrent dequeue/re-enqueue race, the task might be in the
TASK_RUNNING state elsewhere.

Modifying p->scx.slice without holding the task's current rq->lock can race
with update_curr_scx() on the other CPU, potentially overwriting the slice
decrement from the timer tick.

> +
> +	scx_dsq_insert_commit(sch, p, dsq_id, enq_flags, (unsigned long)cookie);
> +
> +	return true;
> +}

Is there a missing API variant for inserting into PRIQ (vtime-ordered) DSQs
with a cookie?

To insert a task into a PRIQ DSQ, SCX_ENQ_DSQ_PRIQ must be set. However,
BPF schedulers cannot manually pass this flag in enq_flags because
scx_dsq_insert_preamble() rejects it as part of __SCX_ENQ_INTERNAL_MASK.
This prevents BPF schedulers from using the new cookie mechanism for
vtime-ordered priority queues.

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260506075925.371138-1-yphbchou0911@gmail.com?part=1