projects/linux·blob master · refs

270 lines · 7.3 KiB

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
21char _license[] SEC("license") = "GPL";
22
23UEI_DEFINE(uei);
24
25struct {
26 __uint(type, BPF_MAP_TYPE_QUEUE);
27 __uint(max_entries, 32768);
28 __type(value, s32);
29} global_queue SEC(".maps");
30
31enum task_state {
32 TASK_NONE = 0,
33 TASK_ENQUEUED, /* in BPF custody, waiting for ops.dequeue() */
34 TASK_DISPATCHED, /* left custody */
35};
36
37struct task_ctx {
38 enum task_state state;
39 s32 enq_cpu; /* scx_bpf_task_cpu() at ops.enqueue() */
40 u64 enqueue_seq;
41};
42
43struct {
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 */
51struct task_struct___core_sched {
52 unsigned long core_cookie;
53} __attribute__((preserve_access_index));
54
55bool test_use_move_to_local;
56
57u64 enqueue_cnt, dequeue_cnt, dispatch_dequeue_cnt, change_dequeue_cnt;
58u64 remote_dispatch_cnt, remote_running_cnt, missed_dequeue_cnt;
59u64 core_sched_exec_dequeue_cnt;
60
61static 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
66static 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
75s32 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
82void 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);
117out:
118 scx_bpf_kick_cpu(scx_bpf_task_cpu(p), SCX_KICK_IDLE);
119}
120
121void 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
159void 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
208void 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
230s32 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
248s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_remote_init)
249{
250 return scx_bpf_create_dsq(SHARED_DSQ, -1);
251}
252
253void BPF_STRUCT_OPS(dequeue_remote_exit, struct scx_exit_info *ei)
254{
255 UEI_RECORD(uei, ei);
256}
257
258SEC(".struct_ops.link")
259struct 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};