Re: [PATCH v2 2/2] selftests/sched_ext: Add a test for ops.dequeue() on remote local DSQ moves

[Andrea Righi <arighi@nvidia.com>]

Hi Kuba,

few nits below, none of them are blockers.

On Wed, Sep 30, 2026 at 11:47:23AM +0000, Kuba Piecuch wrote:
> Add a dequeue_remote test that makes moves of tasks in the BPF
> scheduler's custody to another CPU's local DSQ the common case, both via
> SCX_DSQ_LOCAL_ON dispatch and via scx_bpf_dsq_move_to_local(). The BPF
> scheduler tracks each task's custody state and triggers scx_bpf_error()
> if a custody period doesn't end with exactly one ops.dequeue() before
> the task runs, or if it ends with an SCX_DEQ_CORE_SCHED_EXEC dequeue of
> a task without a core cookie.
> 
> Without the previous patch, the test fails with:
> 
>   sched_ext: dequeue_remote: dequeue_remote.bpf.c:141: 15 (rcu_preempt): late ops.dequeue() with SCX_DEQ_CORE_SCHED_EXEC (enq_cpu=3 cpu=2 seq=1)
>      ...
>      ops_dequeue+0x114/0x170
>      set_next_task_scx+0x104/0x1e0
>      __pick_next_task+0xc7/0x180
>      __schedule+0x154/0x1870
> 
> Assisted-by: Claude:claude-opus-5.5
> Signed-off-by: Kuba Piecuch <jpiecuch@google.com>

...

> +void BPF_STRUCT_OPS(dequeue_remote_dispatch, s32 cpu, struct task_struct *prev)
> +{
> +	struct task_ctx *tctx;
> +	struct task_struct *p;
> +	s32 pid;
> +	int i;
> +
> +	if (test_use_move_to_local) {
> +		scx_bpf_dsq_move_to_local(SHARED_DSQ, 0);
> +		return;
> +	}
> +
> +	/* pop past stale entries so that they don't leave this CPU idle */
> +	bpf_for(i, 0, MAX_DISPATCH_POPS) {

If there are MAX_DISPATCH_POPS or more stale entries in front of a valid one, we
return without dispatching anything and the CPU goes idle. It's quite unlikely,
but maybe we could re-kick this CPU when the loop runs out of pops while the
queue is still non-empty to make the test a bit more robust?

> +		if (bpf_map_pop_elem(&global_queue, &pid))
> +			return;
> +
> +		p = bpf_task_from_pid(pid);
> +		if (!p)
> +			continue;
> +
> +		/*
> +		 * Entries are stale if @p left custody through a property
> +		 * change dequeue or was dispatched from a duplicate entry.
> +		 */
> +		tctx = lookup_task_ctx(p);
> +		if (!tctx || tctx->state != TASK_ENQUEUED) {
> +			bpf_task_release(p);
> +			continue;
> +		}
> +
> +		if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) {
> +			if (scx_bpf_task_cpu(p) != cpu)
> +				__sync_fetch_and_add(&remote_dispatch_cnt, 1);
> +			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu,
> +					   SCX_SLICE_DFL, 0);
> +		} else {
> +			scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
> +		}
> +
> +		bpf_task_release(p);
> +		return;
> +	}
> +}
> +
> +void BPF_STRUCT_OPS(dequeue_remote_running, struct task_struct *p)
> +{
> +	struct task_ctx *tctx;
> +
> +	tctx = lookup_task_ctx(p);
> +	if (!tctx)
> +		return;
> +
> +	/* tasks can only run from a local DSQ, i.e. after leaving custody */
> +	if (tctx->state == TASK_ENQUEUED) {
> +		__sync_fetch_and_add(&missed_dequeue_cnt, 1);
> +		scx_bpf_error("%d (%s): running without ops.dequeue() (enq_cpu=%d cpu=%d seq=%llu)",
> +			      p->pid, p->comm, tctx->enq_cpu,
> +			      scx_bpf_task_cpu(p), tctx->enqueue_seq);
> +		return;
> +	}
> +
> +	if (tctx->enq_cpu >= 0 && tctx->enq_cpu != scx_bpf_task_cpu(p))
> +		__sync_fetch_and_add(&remote_running_cnt, 1);
> +	tctx->enq_cpu = -1;
> +}
> +
> +s32 BPF_STRUCT_OPS(dequeue_remote_init_task, struct task_struct *p,
> +		   struct scx_init_task_args *args)
> +{
> +	struct task_ctx *tctx;
> +
> +	tctx = bpf_task_storage_get(&task_ctx_stor, p, 0,
> +				    BPF_LOCAL_STORAGE_GET_F_CREATE);
> +	if (!tctx)
> +		return -ENOMEM;
> +
> +	/* task storage persists across attachments, start from scratch */
> +	tctx->state = TASK_NONE;
> +	tctx->enq_cpu = -1;

Nit: maybe reset enqueue_seq as well?

> +
> +	return 0;
> +}
> +
> +s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_remote_init)
> +{
> +	return scx_bpf_create_dsq(SHARED_DSQ, -1);
> +}
> +
> +void BPF_STRUCT_OPS(dequeue_remote_exit, struct scx_exit_info *ei)
> +{
> +	UEI_RECORD(uei, ei);
> +}
> +
> +SEC(".struct_ops.link")
> +struct sched_ext_ops dequeue_remote_ops = {
> +	.select_cpu		= (void *)dequeue_remote_select_cpu,
> +	.enqueue		= (void *)dequeue_remote_enqueue,
> +	.dequeue		= (void *)dequeue_remote_dequeue,
> +	.dispatch		= (void *)dequeue_remote_dispatch,
> +	.running		= (void *)dequeue_remote_running,
> +	.init_task		= (void *)dequeue_remote_init_task,
> +	.init			= (void *)dequeue_remote_init,
> +	.exit			= (void *)dequeue_remote_exit,
> +	.flags			= SCX_OPS_ENQ_LAST,
> +	.name			= "dequeue_remote",
> +};
> diff --git a/tools/testing/selftests/sched_ext/dequeue_remote.c b/tools/testing/selftests/sched_ext/dequeue_remote.c
> new file mode 100644
> index 000000000000..0f03044a5bc6
> --- /dev/null
> +++ b/tools/testing/selftests/sched_ext/dequeue_remote.c

...

> +static enum scx_test_status run_scenario(struct dequeue_remote *skel,
> +					 bool use_move_to_local,
> +					 const char *name)
> +{
> +	enum scx_test_status ret = SCX_TEST_PASS;
> +	struct bpf_link *link;
> +	pid_t pids[MAX_WORKERS];
> +	int nr_workers, nr_forked, i;
> +
> +	nr_workers = 2 * nr_cpus;
> +	if (nr_workers < 4)
> +		nr_workers = 4;
> +	if (nr_workers > MAX_WORKERS)
> +		nr_workers = MAX_WORKERS;
> +
> +	skel->bss->test_use_move_to_local = use_move_to_local;
> +	skel->bss->enqueue_cnt = 0;
> +	skel->bss->dequeue_cnt = 0;
> +	skel->bss->dispatch_dequeue_cnt = 0;
> +	skel->bss->change_dequeue_cnt = 0;
> +	skel->bss->remote_dispatch_cnt = 0;
> +	skel->bss->remote_running_cnt = 0;
> +	skel->bss->missed_dequeue_cnt = 0;
> +	skel->bss->core_sched_exec_dequeue_cnt = 0;

Nit: uei isn't reset between scenarios, so in theory this check in the
second scenario could pass using the exit record of the first one. So maybe we
should add:

    memset(&skel->data->uei, 0, sizeof(skel->data->uei));

Overall, the test looks good.

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

Thanks,
-Andrea