Re: [PATCH 2/2] selftests/sched_ext: Check the cmask cid-form ops.enable() receives

[Andrea Righi <arighi@nvidia.com>]

Hi Tejun,

nits below.

On Fri, Sep 18, 2026 at 02:28:38PM -1000, Tejun Heo wrote:
> cid-form ops.enable() now hands the task's cmask to the scheduler and
> set_cmask() repeats it right after. Add a cid-form selftest that checks both
> against p->cpus_ptr, that they match each other, that the initial
> set_cmask() lands before set_weight() and before the task first becomes
> runnable, and that set_cmask() never precedes enable(), across class-switch
> enables, fork-path enables and live affinity changes.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>
> ---
>  tools/testing/selftests/sched_ext/Makefile    |   1 +
>  .../selftests/sched_ext/enable_cmask.bpf.c    | 213 ++++++++++++++++++
>  .../selftests/sched_ext/enable_cmask.c        | 138 ++++++++++++
>  3 files changed, 352 insertions(+)
>  create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.bpf.c
>  create mode 100644 tools/testing/selftests/sched_ext/enable_cmask.c
> 
> diff --git a/tools/testing/selftests/sched_ext/Makefile b/tools/testing/selftests/sched_ext/Makefile
> index 49897727f535..4e06d0baaeec 100644
> --- a/tools/testing/selftests/sched_ext/Makefile
> +++ b/tools/testing/selftests/sched_ext/Makefile
> @@ -169,6 +169,7 @@ auto-test-targets :=			\
>  	ddsp_bogus_dsq_fail		\
>  	ddsp_vtimelocal_fail		\
>  	dsp_local_on			\
> +	enable_cmask			\
>  	enq_select_cpu			\
>  	exit				\
>  	hotplug				\
> diff --git a/tools/testing/selftests/sched_ext/enable_cmask.bpf.c b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
> new file mode 100644
> index 000000000000..bbbf989d2c69
> --- /dev/null
> +++ b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
> @@ -0,0 +1,213 @@
> +// SPDX-License-Identifier: GPL-2.0
> +/*
> + * A cid-form scheduler checking the cmask cid-form ops.enable() receives: the
> + * header, every cid bit against p->cpus_ptr, and that set_cmask() follows with
> + * the same mask before set_weight() and before the task first becomes runnable,
> + * and never runs before enable().
> + *
> + * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
> + */
> +#include <scx/common.bpf.h>
> +
> +#define MAX_CPUS 1024
> +
> +char _license[] SEC("license") = "GPL";
> +
> +struct {
> +	__uint(type, BPF_MAP_TYPE_ARENA);
> +	__uint(map_flags, BPF_F_MMAPABLE);
> +	__uint(max_entries, 1 << 16);
> +} arena SEC(".maps");
> +
> +struct task_ctx {
> +	u64	enable_fp;	/* fingerprint of the mask enable() received */
> +	bool	enabled;
> +	bool	pending;	/* enable() ran, the initial set_cmask() hasn't */
> +};
> +
> +struct {
> +	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
> +	__uint(map_flags, BPF_F_NO_PREALLOC);
> +	__type(key, int);
> +	__type(value, struct task_ctx);
> +} task_ctx_stor SEC(".maps");
> +
> +u64 nr_enable, nr_initial_set_cmask, nr_set_cmask, nr_set_weight;
> +s32 bad_cid;
> +bool bad_want, bad_got;

As Sashiko also reported, these globals can be written concurrently by
ops.enable() and ops.set_cmask() and could make a failure report misleading. It
should be possible to move these in a caller-local stack struct.

> +
> +UEI_DEFINE(uei);
> +
> +static struct task_ctx *lookup_task_ctx(struct task_struct *p)
> +{
> +	struct task_ctx *tctx;
> +
> +	tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
> +	if (!tctx)
> +		scx_bpf_error("task_ctx lookup failed for %s[%d]", p->comm, p->pid);
> +	return tctx;
> +}
> +
> +/*
> + * Verify @m's header and every cid bit against @p's cpumask and fingerprint the
> + * bits into @fp. Return 0 on success, -EINVAL on a bad header, -ENOENT on a cid
> + * without a cpu and -EIO on a bit mismatch with the details in @bad_*.
> + */
> +static int check_mask(struct task_struct *p, const struct scx_cmask __arena *m, u64 *fp)
> +{
> +	u32 nr_cids = scx_bpf_nr_cids();
> +	u64 h = 0;
> +	s32 cid;
> +
> +	if (m->base || m->nr_cids != nr_cids)
> +		return -EINVAL;

Since this is checking the mask header, should this validate alloc_words too?

	if (m->base || m->nr_cids != nr_cids ||
	    m->alloc_words != CMASK_NR_WORDS(nr_cids))
		return -EINVAL;

> +
> +	bpf_for(cid, 0, MAX_CPUS) {
> +		bool want, got;
> +		s32 cpu;
> +
> +		if (cid >= nr_cids)
> +			break;

Can we iterate directly to nr_cids?

> +		cpu = scx_bpf_cid_to_cpu(cid);
> +		if (cpu < 0)
> +			return -ENOENT;
> +		want = bpf_cpumask_test_cpu(cpu, p->cpus_ptr);
> +		got = cmask_test(cid, m);
> +		if (want != got) {
> +			bad_cid = cid;
> +			bad_want = want;
> +			bad_got = got;
> +			return -EIO;
> +		}
> +		h = h * 31 + got;
> +	}
> +
> +	*fp = h;
> +	return 0;
> +}

Thanks,
-Andrea