1 files changed, 217 insertions, 0 deletions
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..0068f3c7ab3c
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/enable_cmask.bpf.c
@@ -0,0 +1,217 @@
+// 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>
+
+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");
+
+/* details of a cid bit mismatch, filled by check_mask() */
+struct mask_mismatch {
+	s32	cid;
+	bool	want;
+	bool	got;
+};
+
+u64 nr_enable, nr_initial_set_cmask, nr_set_cmask, nr_set_weight;
+
+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 @mm.
+ */
+static int check_mask(struct task_struct *p, const struct scx_cmask __arena *m, u64 *fp,
+		      struct mask_mismatch *mm)
+{
+	u32 nr_cids = scx_bpf_nr_cids();
+	u64 h = 0;
+	s32 cid;
+
+	if (m->base || m->nr_cids != nr_cids || m->alloc_words != CMASK_NR_WORDS(nr_cids))
+		return -EINVAL;
+
+	bpf_for(cid, 0, nr_cids) {
+		bool want, got;
+		s32 cpu;
+
+		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) {
+			mm->cid = cid;
+			mm->want = want;
+			mm->got = got;
+			return -EIO;
+		}
+		h = h * 31 + got;
+	}
+
+	*fp = h;
+	return 0;
+}
+
+s32 BPF_STRUCT_OPS_SLEEPABLE(enable_cmask_init_task, struct task_struct *p,
+			     struct scx_init_task_args *args)
+{
+	if (!bpf_task_storage_get(&task_ctx_stor, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE))
+		return -ENOMEM;
+	return 0;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_enable, struct task_struct *p, struct scx_enable_args *args)
+{
+	struct scx_cmask __arena *m = (struct scx_cmask __arena *)args->cmask_arena_addr;
+	struct mask_mismatch mm = {};
+	struct task_ctx *tctx;
+	int ret;
+
+	asm volatile("" :: "r"(&arena));
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_enable, 1);
+	if (tctx->enabled || tctx->pending) {
+		scx_bpf_error("enable: %s[%d] enabled twice", p->comm, p->pid);
+		return;
+	}
+
+	ret = check_mask(p, m, &tctx->enable_fp, &mm);
+	if (ret) {
+		scx_bpf_error("enable: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
+			      p->comm, p->pid, ret, mm.cid, mm.want, mm.got);
+		return;
+	}
+	tctx->enabled = true;
+	tctx->pending = true;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_set_cmask, struct task_struct *p,
+		    struct scx_cmask __arena *m)
+{
+	struct mask_mismatch mm = {};
+	struct task_ctx *tctx;
+	u64 fp;
+	int ret;
+
+	asm volatile("" :: "r"(&arena));
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_set_cmask, 1);
+	if (!tctx->enabled) {
+		scx_bpf_error("set_cmask: %s[%d] not enabled", p->comm, p->pid);
+		return;
+	}
+
+	ret = check_mask(p, m, &fp, &mm);
+	if (ret) {
+		scx_bpf_error("set_cmask: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
+			      p->comm, p->pid, ret, mm.cid, mm.want, mm.got);
+		return;
+	}
+
+	if (tctx->pending) {
+		if (fp != tctx->enable_fp) {
+			scx_bpf_error("set_cmask: %s[%d] initial mask differs from enable()",
+				      p->comm, p->pid);
+			return;
+		}
+		tctx->pending = false;
+		__sync_fetch_and_add(&nr_initial_set_cmask, 1);
+	}
+}
+
+void BPF_STRUCT_OPS(enable_cmask_set_weight, struct task_struct *p, u32 weight)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	__sync_fetch_and_add(&nr_set_weight, 1);
+	if (tctx->pending)
+		scx_bpf_error("set_weight: %s[%d] before the initial set_cmask()", p->comm,
+			      p->pid);
+}
+
+void BPF_STRUCT_OPS(enable_cmask_runnable, struct task_struct *p, u64 enq_flags)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	if (tctx->pending)
+		scx_bpf_error("runnable: %s[%d] before the initial set_cmask()", p->comm,
+			      p->pid);
+}
+
+void BPF_STRUCT_OPS(enable_cmask_disable, struct task_struct *p)
+{
+	struct task_ctx *tctx;
+
+	tctx = lookup_task_ctx(p);
+	if (!tctx)
+		return;
+
+	tctx->enabled = false;
+	tctx->pending = false;
+}
+
+void BPF_STRUCT_OPS(enable_cmask_exit, struct scx_exit_info *ei)
+{
+	UEI_RECORD(uei, ei);
+}
+
+SCX_OPS_CID_DEFINE(enable_cmask_ops,
+		   .init_task	= (void *)enable_cmask_init_task,
+		   .enable	= (void *)enable_cmask_enable,
+		   .set_cmask	= (void *)enable_cmask_set_cmask,
+		   .set_weight	= (void *)enable_cmask_set_weight,
+		   .runnable	= (void *)enable_cmask_runnable,
+		   .disable	= (void *)enable_cmask_disable,
+		   .exit	= (void *)enable_cmask_exit,
+		   .flags	= SCX_OPS_SWITCH_PARTIAL,
+		   .name	= "enable_cmask");