Re: [PATCH v3 2/2] selftests/sched_ext: Add dispatch_cookie test

[Andrea Righi <arighi@nvidia.com>]

Hi Cheng-Yang,

On Sun, May 10, 2026 at 03:11:57AM +0800, Cheng-Yang Chou wrote:
> Test scx_bpf_dsq_insert_begin() and scx_bpf_dsq_insert_commit().
> 
> The BPF scheduler enqueues tasks into a BPF queue map in ops.enqueue()
> and dispatches them via the begin/commit transaction API in ops.dispatch().
> After a successful dispatch, the token is stored. On the task's next
> dispatch (after dequeue/re-enqueue increments qseq), the stored token is
> stale and finish_dispatch() silently drops the buffered entry.
> 
> Userspace forks spinning children and repeatedly flips their CPU affinity
> to trigger dequeue/re-enqueue cycles. The test verifies both nr_tx_dispatched
> and nr_tx_stale are positive, exercising both the happy path and the
> race-detection window.
> 
> Suggested-by: Tejun Heo <tj@kernel.org>
> Suggested-by: Kuba Piecuch <jpiecuch@google.com>
> Signed-off-by: Cheng-Yang Chou <yphbchou0911@gmail.com>
> ---
...
> +void BPF_STRUCT_OPS(dispatch_cookie_dispatch, s32 cpu,
> +		    struct task_struct *prev)
> +{
> +	s32 pid;
> +	struct task_struct *p;
> +	u64 *stored, token;
> +
> +	if (bpf_map_pop_elem(&queue, &pid))
> +		return;
> +
> +	p = bpf_task_from_pid(pid);
> +	if (!p)
> +		return;
> +
> +	/*
> +	 * After a successful fresh dispatch, store the token. On the task's
> +	 * next dispatch (after re-enqueue increments qseq), the stored token
> +	 * exercises the race-detection window in finish_dispatch().
> +	 *
> +	 * scx_bpf_dsq_insert_commit() always returns %true when the preamble
> +	 * passes. Stale detection fires asynchronously in finish_dispatch()
> +	 * with no BPF-observable signal. When using a stored token, always
> +	 * pair the commit() call with a fallback scx_bpf_dsq_insert(): if
> +	 * the token is stale, finish_dispatch() drops the first entry and
> +	 * the fallback dispatches the task. If the token is still fresh,
> +	 * finish_dispatch() dispatches it and the fallback's CAS is a no-op.
> +	 */
> +	stored = bpf_map_lookup_elem(&last_token, &pid);
> +	if (stored) {
> +		token = *stored;
> +		bpf_map_delete_elem(&last_token, &pid);
> +		scx_bpf_dsq_insert_commit(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, token);

I think we should pass 0 instead of SCX_SLICE_DFL here:

 scx_bpf_dsq_insert_commit(p, dsq_id, enq_flags, token)

> +		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
> +		__sync_fetch_and_add(&nr_tx_stale, 1);
> +	} else {
> +		token = scx_bpf_dsq_insert_begin(p);
> +		if (scx_bpf_dsq_insert_commit(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, token)) {

Ditto.

> +			__sync_fetch_and_add(&nr_tx_dispatched, 1);

Also, stale detection happens asynchronously later in finish_dispatch() with no
BPF-observable signal, so commit can return true here and write the buffered
entry with stale qseq=N, but the counter is incremented before finish_dispatch()
runs. Later, when finish_dispatch() processes the entry, it detects the qseq
mismatch (task now has N+1) and silently drops the dispatch. So nr_tx_stale
counts "attempted stale commits", not "detected and rejected stale tokens".

> +			bpf_map_update_elem(&last_token, &pid, &token, BPF_ANY);
> +		} else {
> +			scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
> +		}
> +	}
> +
> +	bpf_task_release(p);
> +}
> +
> +void BPF_STRUCT_OPS(dispatch_cookie_exit, struct scx_exit_info *ei)
> +{
> +	UEI_RECORD(uei, ei);
> +}
> +
> +SEC(".struct_ops.link")
> +struct sched_ext_ops dispatch_cookie_ops = {
> +	.enqueue		= (void *) dispatch_cookie_enqueue,
> +	.dispatch		= (void *) dispatch_cookie_dispatch,
> +	.exit			= (void *) dispatch_cookie_exit,
> +	.name			= "dispatch_cookie",
> +	.timeout_ms		= 5000U,
> +};

...

> +static enum scx_test_status run(void *ctx)
> +{
> +	struct dispatch_cookie_ctx *tctx = ctx;
> +	cpu_set_t cpuset_one, cpuset_all;
> +	pid_t pids[NUM_CHILDREN];
> +	int i, j, nforked = 0, status;
> +
> +	tctx->link = bpf_map__attach_struct_ops(tctx->skel->maps.dispatch_cookie_ops);
> +	SCX_FAIL_IF(!tctx->link, "Failed to attach scheduler");
> +
> +	CPU_ZERO(&cpuset_one);
> +	CPU_SET(0, &cpuset_one);

This is probably fine as it is, but it's worth noticing that if for any reason
CPU0 is offline the test is probably failing.

> +	SCX_FAIL_IF(sched_getaffinity(0, sizeof(cpuset_all), &cpuset_all),
> +		    "Failed to get CPU affinity (%d)", errno);
> +
> +	for (i = 0; i < NUM_CHILDREN; i++) {
> +		pids[i] = fork();
> +		if (pids[i] == 0) {
> +			while (1)
> +				sched_yield();
> +		}
> +		if (pids[i] > 0)
> +			nforked++;
> +	}

Thanks,
-Andrea