1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * A cid-form scheduler checking the cmask cid-form ops.enable() receives: the
4 * header, every cid bit against p->cpus_ptr, and that set_cmask() follows with
5 * the same mask before set_weight() and before the task first becomes runnable,
6 * and never runs before enable().
7 *
8 * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
9 */
10 #include <scx/common.bpf.h>
11
12 char _license[] SEC("license") = "GPL";
13
14 struct {
15 __uint(type, BPF_MAP_TYPE_ARENA);
16 __uint(map_flags, BPF_F_MMAPABLE);
17 __uint(max_entries, 1 << 16);
18 } arena SEC(".maps");
19
20 struct task_ctx {
21 u64 enable_fp; /* fingerprint of the mask enable() received */
22 bool enabled;
23 bool pending; /* enable() ran, the initial set_cmask() hasn't */
24 };
25
26 struct {
27 __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
28 __uint(map_flags, BPF_F_NO_PREALLOC);
29 __type(key, int);
30 __type(value, struct task_ctx);
31 } task_ctx_stor SEC(".maps");
32
33 /* details of a cid bit mismatch, filled by check_mask() */
34 struct mask_mismatch {
35 s32 cid;
36 bool want;
37 bool got;
38 };
39
40 u64 nr_enable, nr_initial_set_cmask, nr_set_cmask, nr_set_weight;
41
42 UEI_DEFINE(uei);
43
lookup_task_ctx(struct task_struct * p)44 static struct task_ctx *lookup_task_ctx(struct task_struct *p)
45 {
46 struct task_ctx *tctx;
47
48 tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
49 if (!tctx)
50 scx_bpf_error("task_ctx lookup failed for %s[%d]", p->comm, p->pid);
51 return tctx;
52 }
53
54 /*
55 * Verify @m's header and every cid bit against @p's cpumask and fingerprint the
56 * bits into @fp. Return 0 on success, -EINVAL on a bad header, -ENOENT on a cid
57 * without a cpu and -EIO on a bit mismatch with the details in @mm.
58 */
check_mask(struct task_struct * p,const struct scx_cmask __arena * m,u64 * fp,struct mask_mismatch * mm)59 static int check_mask(struct task_struct *p, const struct scx_cmask __arena *m, u64 *fp,
60 struct mask_mismatch *mm)
61 {
62 u32 nr_cids = scx_bpf_nr_cids();
63 u64 h = 0;
64 s32 cid;
65
66 if (m->base || m->nr_cids != nr_cids || m->alloc_words != CMASK_NR_WORDS(nr_cids))
67 return -EINVAL;
68
69 bpf_for(cid, 0, nr_cids) {
70 bool want, got;
71 s32 cpu;
72
73 cpu = scx_bpf_cid_to_cpu(cid);
74 if (cpu < 0)
75 return -ENOENT;
76 want = bpf_cpumask_test_cpu(cpu, p->cpus_ptr);
77 got = cmask_test(cid, m);
78 if (want != got) {
79 mm->cid = cid;
80 mm->want = want;
81 mm->got = got;
82 return -EIO;
83 }
84 h = h * 31 + got;
85 }
86
87 *fp = h;
88 return 0;
89 }
90
BPF_STRUCT_OPS_SLEEPABLE(enable_cmask_init_task,struct task_struct * p,struct scx_init_task_args * args)91 s32 BPF_STRUCT_OPS_SLEEPABLE(enable_cmask_init_task, struct task_struct *p,
92 struct scx_init_task_args *args)
93 {
94 if (!bpf_task_storage_get(&task_ctx_stor, p, 0, BPF_LOCAL_STORAGE_GET_F_CREATE))
95 return -ENOMEM;
96 return 0;
97 }
98
BPF_STRUCT_OPS(enable_cmask_enable,struct task_struct * p,struct scx_enable_args * args)99 void BPF_STRUCT_OPS(enable_cmask_enable, struct task_struct *p, struct scx_enable_args *args)
100 {
101 struct scx_cmask __arena *m = (struct scx_cmask __arena *)args->cmask_arena_addr;
102 struct mask_mismatch mm = {};
103 struct task_ctx *tctx;
104 int ret;
105
106 asm volatile("" :: "r"(&arena));
107 tctx = lookup_task_ctx(p);
108 if (!tctx)
109 return;
110
111 __sync_fetch_and_add(&nr_enable, 1);
112 if (tctx->enabled || tctx->pending) {
113 scx_bpf_error("enable: %s[%d] enabled twice", p->comm, p->pid);
114 return;
115 }
116
117 ret = check_mask(p, m, &tctx->enable_fp, &mm);
118 if (ret) {
119 scx_bpf_error("enable: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
120 p->comm, p->pid, ret, mm.cid, mm.want, mm.got);
121 return;
122 }
123 tctx->enabled = true;
124 tctx->pending = true;
125 }
126
BPF_STRUCT_OPS(enable_cmask_set_cmask,struct task_struct * p,struct scx_cmask __arena * m)127 void BPF_STRUCT_OPS(enable_cmask_set_cmask, struct task_struct *p,
128 struct scx_cmask __arena *m)
129 {
130 struct mask_mismatch mm = {};
131 struct task_ctx *tctx;
132 u64 fp;
133 int ret;
134
135 asm volatile("" :: "r"(&arena));
136 tctx = lookup_task_ctx(p);
137 if (!tctx)
138 return;
139
140 __sync_fetch_and_add(&nr_set_cmask, 1);
141 if (!tctx->enabled) {
142 scx_bpf_error("set_cmask: %s[%d] not enabled", p->comm, p->pid);
143 return;
144 }
145
146 ret = check_mask(p, m, &fp, &mm);
147 if (ret) {
148 scx_bpf_error("set_cmask: %s[%d] cmask check failed %d cid=%d want=%d got=%d",
149 p->comm, p->pid, ret, mm.cid, mm.want, mm.got);
150 return;
151 }
152
153 if (tctx->pending) {
154 if (fp != tctx->enable_fp) {
155 scx_bpf_error("set_cmask: %s[%d] initial mask differs from enable()",
156 p->comm, p->pid);
157 return;
158 }
159 tctx->pending = false;
160 __sync_fetch_and_add(&nr_initial_set_cmask, 1);
161 }
162 }
163
BPF_STRUCT_OPS(enable_cmask_set_weight,struct task_struct * p,u32 weight)164 void BPF_STRUCT_OPS(enable_cmask_set_weight, struct task_struct *p, u32 weight)
165 {
166 struct task_ctx *tctx;
167
168 tctx = lookup_task_ctx(p);
169 if (!tctx)
170 return;
171
172 __sync_fetch_and_add(&nr_set_weight, 1);
173 if (tctx->pending)
174 scx_bpf_error("set_weight: %s[%d] before the initial set_cmask()", p->comm,
175 p->pid);
176 }
177
BPF_STRUCT_OPS(enable_cmask_runnable,struct task_struct * p,u64 enq_flags)178 void BPF_STRUCT_OPS(enable_cmask_runnable, struct task_struct *p, u64 enq_flags)
179 {
180 struct task_ctx *tctx;
181
182 tctx = lookup_task_ctx(p);
183 if (!tctx)
184 return;
185
186 if (tctx->pending)
187 scx_bpf_error("runnable: %s[%d] before the initial set_cmask()", p->comm,
188 p->pid);
189 }
190
BPF_STRUCT_OPS(enable_cmask_disable,struct task_struct * p)191 void BPF_STRUCT_OPS(enable_cmask_disable, struct task_struct *p)
192 {
193 struct task_ctx *tctx;
194
195 tctx = lookup_task_ctx(p);
196 if (!tctx)
197 return;
198
199 tctx->enabled = false;
200 tctx->pending = false;
201 }
202
BPF_STRUCT_OPS(enable_cmask_exit,struct scx_exit_info * ei)203 void BPF_STRUCT_OPS(enable_cmask_exit, struct scx_exit_info *ei)
204 {
205 UEI_RECORD(uei, ei);
206 }
207
208 SCX_OPS_CID_DEFINE(enable_cmask_ops,
209 .init_task = (void *)enable_cmask_init_task,
210 .enable = (void *)enable_cmask_enable,
211 .set_cmask = (void *)enable_cmask_set_cmask,
212 .set_weight = (void *)enable_cmask_set_weight,
213 .runnable = (void *)enable_cmask_runnable,
214 .disable = (void *)enable_cmask_disable,
215 .exit = (void *)enable_cmask_exit,
216 .flags = SCX_OPS_SWITCH_PARTIAL,
217 .name = "enable_cmask");
218