1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * Verify that ops.dequeue() is called when a task leaves the BPF scheduler's
4 * custody by being moved to the local DSQ of a CPU other than the one whose
5 * rq it's on (move_remote_task_to_local_dsq()).
6 *
7 * ops.enqueue() puts every task into custody and ops.dispatch() moves it to
8 * the dispatching CPU's local DSQ, so most moves cross CPUs. With
9 * @test_use_move_to_local, tasks are queued on a user DSQ and consumed with
10 * scx_bpf_dsq_move_to_local(). Otherwise, they are queued in a BPF queue and
11 * dispatched with SCX_DSQ_LOCAL_ON.
12 *
13 * Copyright (c) 2026 Google LLC.
14 */
15
16 #include <scx/common.bpf.h>
17
18 #define SHARED_DSQ 0
19 #define MAX_DISPATCH_POPS 8
20
21 char _license[] SEC("license") = "GPL";
22
23 UEI_DEFINE(uei);
24
25 struct {
26 __uint(type, BPF_MAP_TYPE_QUEUE);
27 __uint(max_entries, 32768);
28 __type(value, s32);
29 } global_queue SEC(".maps");
30
31 enum task_state {
32 TASK_NONE = 0,
33 TASK_ENQUEUED, /* in BPF custody, waiting for ops.dequeue() */
34 TASK_DISPATCHED, /* left custody */
35 };
36
37 struct task_ctx {
38 enum task_state state;
39 s32 enq_cpu; /* scx_bpf_task_cpu() at ops.enqueue() */
40 u64 enqueue_seq;
41 };
42
43 struct {
44 __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
45 __uint(map_flags, BPF_F_NO_PREALLOC);
46 __type(key, int);
47 __type(value, struct task_ctx);
48 } task_ctx_stor SEC(".maps");
49
50 /* core_cookie only exists with CONFIG_SCHED_CORE */
51 struct task_struct___core_sched {
52 unsigned long core_cookie;
53 } __attribute__((preserve_access_index));
54
55 bool test_use_move_to_local;
56
57 u64 enqueue_cnt, dequeue_cnt, dispatch_dequeue_cnt, change_dequeue_cnt;
58 u64 remote_dispatch_cnt, remote_running_cnt, missed_dequeue_cnt;
59 u64 core_sched_exec_dequeue_cnt;
60
lookup_task_ctx(struct task_struct * p)61 static struct task_ctx *lookup_task_ctx(struct task_struct *p)
62 {
63 return bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
64 }
65
task_has_core_cookie(struct task_struct * p)66 static bool task_has_core_cookie(struct task_struct *p)
67 {
68 struct task_struct___core_sched *t = (void *)p;
69
70 if (!bpf_core_field_exists(t->core_cookie))
71 return false;
72 return BPF_CORE_READ(t, core_cookie);
73 }
74
BPF_STRUCT_OPS(dequeue_remote_select_cpu,struct task_struct * p,s32 prev_cpu,u64 wake_flags)75 s32 BPF_STRUCT_OPS(dequeue_remote_select_cpu, struct task_struct *p,
76 s32 prev_cpu, u64 wake_flags)
77 {
78 /* no direct dispatch, always go through ops.enqueue() */
79 return prev_cpu;
80 }
81
BPF_STRUCT_OPS(dequeue_remote_enqueue,struct task_struct * p,u64 enq_flags)82 void BPF_STRUCT_OPS(dequeue_remote_enqueue, struct task_struct *p, u64 enq_flags)
83 {
84 struct task_ctx *tctx;
85 s32 pid = p->pid;
86
87 tctx = lookup_task_ctx(p);
88 if (!tctx) {
89 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
90 return;
91 }
92
93 /* the previous custody period must have ended with ops.dequeue() */
94 if (tctx->state == TASK_ENQUEUED)
95 scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu",
96 p->pid, p->comm, tctx->enqueue_seq);
97
98 /*
99 * Mark @p as enqueued before making it visible to ops.dispatch() on
100 * other CPUs, which skips queue entries of tasks not in ENQUEUED
101 * state as stale.
102 */
103 tctx->state = TASK_ENQUEUED;
104 tctx->enq_cpu = scx_bpf_task_cpu(p);
105 tctx->enqueue_seq++;
106
107 if (test_use_move_to_local) {
108 scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags);
109 } else if (bpf_map_push_elem(&global_queue, &pid, 0)) {
110 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
111 tctx->state = TASK_DISPATCHED;
112 tctx->enq_cpu = -1;
113 goto out;
114 }
115
116 __sync_fetch_and_add(&enqueue_cnt, 1);
117 out:
118 scx_bpf_kick_cpu(scx_bpf_task_cpu(p), SCX_KICK_IDLE);
119 }
120
BPF_STRUCT_OPS(dequeue_remote_dequeue,struct task_struct * p,u64 deq_flags)121 void BPF_STRUCT_OPS(dequeue_remote_dequeue, struct task_struct *p, u64 deq_flags)
122 {
123 struct task_ctx *tctx;
124
125 __sync_fetch_and_add(&dequeue_cnt, 1);
126
127 tctx = lookup_task_ctx(p);
128 if (!tctx)
129 return;
130
131 /*
132 * Only core scheduling can pick a task straight out of custody, and
133 * only if the task has a core cookie. Otherwise, the custody exit was
134 * missed when @p was inserted into a local DSQ and got deferred until
135 * @p was picked.
136 */
137 if ((deq_flags & SCX_DEQ_CORE_SCHED_EXEC) && !task_has_core_cookie(p)) {
138 __sync_fetch_and_add(&core_sched_exec_dequeue_cnt, 1);
139 scx_bpf_error("%d (%s): late ops.dequeue() with SCX_DEQ_CORE_SCHED_EXEC (enq_cpu=%d cpu=%d seq=%llu)",
140 p->pid, p->comm, tctx->enq_cpu,
141 scx_bpf_task_cpu(p), tctx->enqueue_seq);
142 }
143
144 /* ops.dequeue() ends the custody period started by ops.enqueue() */
145 if (tctx->state != TASK_ENQUEUED)
146 scx_bpf_error("%d (%s): dequeue outside custody deq_flags=0x%llx state=%d seq=%llu",
147 p->pid, p->comm, deq_flags, tctx->state,
148 tctx->enqueue_seq);
149
150 if (deq_flags & SCX_DEQ_SCHED_CHANGE) {
151 __sync_fetch_and_add(&change_dequeue_cnt, 1);
152 tctx->state = TASK_NONE;
153 } else {
154 __sync_fetch_and_add(&dispatch_dequeue_cnt, 1);
155 tctx->state = TASK_DISPATCHED;
156 }
157 }
158
BPF_STRUCT_OPS(dequeue_remote_dispatch,s32 cpu,struct task_struct * prev)159 void BPF_STRUCT_OPS(dequeue_remote_dispatch, s32 cpu, struct task_struct *prev)
160 {
161 struct task_ctx *tctx;
162 struct task_struct *p;
163 s32 pid;
164 int i;
165
166 if (test_use_move_to_local) {
167 scx_bpf_dsq_move_to_local(SHARED_DSQ, 0);
168 return;
169 }
170
171 /* pop past stale entries so that they don't leave this CPU idle */
172 bpf_for(i, 0, MAX_DISPATCH_POPS) {
173 if (bpf_map_pop_elem(&global_queue, &pid))
174 return;
175
176 p = bpf_task_from_pid(pid);
177 if (!p)
178 continue;
179
180 /*
181 * Entries are stale if @p left custody through a property
182 * change dequeue or was dispatched from a duplicate entry.
183 */
184 tctx = lookup_task_ctx(p);
185 if (!tctx || tctx->state != TASK_ENQUEUED) {
186 bpf_task_release(p);
187 continue;
188 }
189
190 if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) {
191 if (scx_bpf_task_cpu(p) != cpu)
192 __sync_fetch_and_add(&remote_dispatch_cnt, 1);
193 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu,
194 SCX_SLICE_DFL, 0);
195 } else {
196 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
197 }
198
199 bpf_task_release(p);
200 return;
201 }
202
203 /* out of pops with entries left, retry instead of idling this CPU */
204 if (!bpf_map_peek_elem(&global_queue, &pid))
205 scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE);
206 }
207
BPF_STRUCT_OPS(dequeue_remote_running,struct task_struct * p)208 void BPF_STRUCT_OPS(dequeue_remote_running, struct task_struct *p)
209 {
210 struct task_ctx *tctx;
211
212 tctx = lookup_task_ctx(p);
213 if (!tctx)
214 return;
215
216 /* tasks can only run from a local DSQ, i.e. after leaving custody */
217 if (tctx->state == TASK_ENQUEUED) {
218 __sync_fetch_and_add(&missed_dequeue_cnt, 1);
219 scx_bpf_error("%d (%s): running without ops.dequeue() (enq_cpu=%d cpu=%d seq=%llu)",
220 p->pid, p->comm, tctx->enq_cpu,
221 scx_bpf_task_cpu(p), tctx->enqueue_seq);
222 return;
223 }
224
225 if (tctx->enq_cpu >= 0 && tctx->enq_cpu != scx_bpf_task_cpu(p))
226 __sync_fetch_and_add(&remote_running_cnt, 1);
227 tctx->enq_cpu = -1;
228 }
229
BPF_STRUCT_OPS(dequeue_remote_init_task,struct task_struct * p,struct scx_init_task_args * args)230 s32 BPF_STRUCT_OPS(dequeue_remote_init_task, struct task_struct *p,
231 struct scx_init_task_args *args)
232 {
233 struct task_ctx *tctx;
234
235 tctx = bpf_task_storage_get(&task_ctx_stor, p, 0,
236 BPF_LOCAL_STORAGE_GET_F_CREATE);
237 if (!tctx)
238 return -ENOMEM;
239
240 /* task storage persists across attachments, start from scratch */
241 tctx->state = TASK_NONE;
242 tctx->enq_cpu = -1;
243 tctx->enqueue_seq = 0;
244
245 return 0;
246 }
247
BPF_STRUCT_OPS_SLEEPABLE(dequeue_remote_init)248 s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_remote_init)
249 {
250 return scx_bpf_create_dsq(SHARED_DSQ, -1);
251 }
252
BPF_STRUCT_OPS(dequeue_remote_exit,struct scx_exit_info * ei)253 void BPF_STRUCT_OPS(dequeue_remote_exit, struct scx_exit_info *ei)
254 {
255 UEI_RECORD(uei, ei);
256 }
257
258 SEC(".struct_ops.link")
259 struct sched_ext_ops dequeue_remote_ops = {
260 .select_cpu = (void *)dequeue_remote_select_cpu,
261 .enqueue = (void *)dequeue_remote_enqueue,
262 .dequeue = (void *)dequeue_remote_dequeue,
263 .dispatch = (void *)dequeue_remote_dispatch,
264 .running = (void *)dequeue_remote_running,
265 .init_task = (void *)dequeue_remote_init_task,
266 .init = (void *)dequeue_remote_init,
267 .exit = (void *)dequeue_remote_exit,
268 .flags = SCX_OPS_ENQ_LAST,
269 .name = "dequeue_remote",
270 };
271