1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * A BPF program for testing DSQ operations and peek in particular.
4 *
5 * Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
6 * Copyright (c) 2025 Ryan Newton <ryan.newton@alum.mit.edu>
7 */
8
9 #include <scx/common.bpf.h>
10 #include <scx/compat.bpf.h>
11
12 char _license[] SEC("license") = "GPL";
13
14 UEI_DEFINE(uei); /* Error handling */
15
16 #define MAX_SAMPLES 100
17 #define MAX_CPUS 512
18 #define DSQ_POOL_SIZE 8
19 int max_samples = MAX_SAMPLES;
20 int max_cpus = MAX_CPUS;
21 int dsq_pool_size = DSQ_POOL_SIZE;
22
23 /* Global variables to store test results */
24 int dsq_peek_result1 = -1;
25 long dsq_inserted_pid = -1;
26 int insert_test_cpu = -1; /* Set to the cpu that performs the test */
27 long dsq_peek_result2 = -1;
28 long dsq_peek_result2_pid = -1;
29 long dsq_peek_result2_expected = -1;
30 int test_dsq_id = 1234; /* Use a simple ID like create_dsq example */
31 int real_dsq_id = 1235; /* DSQ for normal operation */
32 int enqueue_count = -1;
33 int dispatch_count = -1;
34 bool debug_ksym_exists;
35
36 /* DSQ pool for stress testing */
37 int dsq_pool_base_id = 2000;
38 int phase1_complete = -1;
39 long total_peek_attempts = -1;
40 long successful_peeks = -1;
41
42 /* BPF map for sharing peek results with userspace */
43 struct {
44 __uint(type, BPF_MAP_TYPE_ARRAY);
45 __uint(max_entries, MAX_SAMPLES);
46 __type(key, u32);
47 __type(value, long);
48 } peek_results SEC(".maps");
49
get_random_dsq_id(void)50 static int get_random_dsq_id(void)
51 {
52 u64 time = bpf_ktime_get_ns();
53
54 return dsq_pool_base_id + (time % DSQ_POOL_SIZE);
55 }
56
record_peek_result(long pid)57 static void record_peek_result(long pid)
58 {
59 u32 slot_key;
60 long *slot_pid_ptr;
61 int ix;
62
63 if (pid <= 0)
64 return;
65
66 /* Find an empty slot or one with the same PID */
67 bpf_for(ix, 0, 10) {
68 slot_key = (pid + ix) % MAX_SAMPLES;
69 slot_pid_ptr = bpf_map_lookup_elem(&peek_results, &slot_key);
70 if (!slot_pid_ptr)
71 continue;
72
73 if (*slot_pid_ptr == -1 || *slot_pid_ptr == pid) {
74 *slot_pid_ptr = pid;
75 break;
76 }
77 }
78 }
79
80 /* Scan all DSQs in the pool and try to move a task to local */
scan_dsq_pool(void)81 static int scan_dsq_pool(void)
82 {
83 struct task_struct *task;
84 int moved = 0;
85 int i;
86
87 bpf_for(i, 0, DSQ_POOL_SIZE) {
88 int dsq_id = dsq_pool_base_id + i;
89
90 total_peek_attempts++;
91
92 task = __COMPAT_scx_bpf_dsq_peek(dsq_id);
93 if (task) {
94 successful_peeks++;
95 record_peek_result(task->pid);
96
97 /* Try to move this task to local */
98 if (!moved && scx_bpf_dsq_move_to_local(dsq_id) == 0) {
99 moved = 1;
100 break;
101 }
102 }
103 }
104 return moved;
105 }
106
107 /* Struct_ops scheduler for testing DSQ peek operations */
BPF_STRUCT_OPS(peek_dsq_enqueue,struct task_struct * p,u64 enq_flags)108 void BPF_STRUCT_OPS(peek_dsq_enqueue, struct task_struct *p, u64 enq_flags)
109 {
110 struct task_struct *peek_result;
111 int last_insert_test_cpu, cpu;
112
113 enqueue_count++;
114 cpu = bpf_get_smp_processor_id();
115 last_insert_test_cpu = __sync_val_compare_and_swap(&insert_test_cpu, -1, cpu);
116
117 /* Phase 1: Simple insert-then-peek test (only on first task) */
118 if (last_insert_test_cpu == -1) {
119 bpf_printk("peek_dsq_enqueue beginning phase 1 peek test on cpu %d", cpu);
120
121 /* Test 1: Peek empty DSQ - should return NULL */
122 peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
123 dsq_peek_result1 = (long)peek_result; /* Should be 0 (NULL) */
124
125 /* Test 2: Insert task into test DSQ for testing in dispatch callback */
126 dsq_inserted_pid = p->pid;
127 scx_bpf_dsq_insert(p, test_dsq_id, 0, enq_flags);
128 dsq_peek_result2_expected = (long)p; /* Expected the task we just inserted */
129 } else if (!phase1_complete) {
130 /* Still in phase 1, use real DSQ */
131 scx_bpf_dsq_insert(p, real_dsq_id, 0, enq_flags);
132 } else {
133 /* Phase 2: Random DSQ insertion for stress testing */
134 int random_dsq_id = get_random_dsq_id();
135
136 scx_bpf_dsq_insert(p, random_dsq_id, 0, enq_flags);
137 }
138 }
139
BPF_STRUCT_OPS(peek_dsq_dispatch,s32 cpu,struct task_struct * prev)140 void BPF_STRUCT_OPS(peek_dsq_dispatch, s32 cpu, struct task_struct *prev)
141 {
142 dispatch_count++;
143
144 /* Phase 1: Complete the simple peek test if we inserted a task but
145 * haven't tested peek yet
146 */
147 if (insert_test_cpu == cpu && dsq_peek_result2 == -1) {
148 struct task_struct *peek_result;
149
150 bpf_printk("peek_dsq_dispatch completing phase 1 peek test on cpu %d", cpu);
151
152 /* Test 3: Peek DSQ after insert - should return the task we inserted */
153 peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
154 /* Store the PID of the peeked task for comparison */
155 dsq_peek_result2 = (long)peek_result;
156 dsq_peek_result2_pid = peek_result ? peek_result->pid : -1;
157
158 /* Now consume the task since we've peeked at it */
159 scx_bpf_dsq_move_to_local(test_dsq_id);
160
161 /* Mark phase 1 as complete */
162 phase1_complete = 1;
163 bpf_printk("Phase 1 complete, starting phase 2 stress testing");
164 } else if (!phase1_complete) {
165 /* Still in phase 1, use real DSQ */
166 scx_bpf_dsq_move_to_local(real_dsq_id);
167 } else {
168 /* Phase 2: Scan all DSQs in the pool and try to move a task */
169 if (!scan_dsq_pool()) {
170 /* No tasks found in DSQ pool, fall back to real DSQ */
171 scx_bpf_dsq_move_to_local(real_dsq_id);
172 }
173 }
174 }
175
BPF_STRUCT_OPS_SLEEPABLE(peek_dsq_init)176 s32 BPF_STRUCT_OPS_SLEEPABLE(peek_dsq_init)
177 {
178 s32 err;
179 int i;
180
181 /* Always set debug values so we can see which version we're using */
182 debug_ksym_exists = bpf_ksym_exists(scx_bpf_dsq_peek) ? 1 : 0;
183
184 /* Initialize state first */
185 insert_test_cpu = -1;
186 enqueue_count = 0;
187 dispatch_count = 0;
188 phase1_complete = 0;
189 total_peek_attempts = 0;
190 successful_peeks = 0;
191
192 /* Create the test and real DSQs */
193 err = scx_bpf_create_dsq(test_dsq_id, -1);
194 if (err) {
195 scx_bpf_error("Failed to create DSQ %d: %d", test_dsq_id, err);
196 return err;
197 }
198 err = scx_bpf_create_dsq(real_dsq_id, -1);
199 if (err) {
200 scx_bpf_error("Failed to create DSQ %d: %d", test_dsq_id, err);
201 return err;
202 }
203
204 /* Create the DSQ pool for stress testing */
205 bpf_for(i, 0, DSQ_POOL_SIZE) {
206 int dsq_id = dsq_pool_base_id + i;
207
208 err = scx_bpf_create_dsq(dsq_id, -1);
209 if (err) {
210 scx_bpf_error("Failed to create DSQ pool entry %d: %d", dsq_id, err);
211 return err;
212 }
213 }
214
215 /* Initialize the peek results map */
216 bpf_for(i, 0, MAX_SAMPLES) {
217 u32 key = i;
218 long pid = -1;
219
220 bpf_map_update_elem(&peek_results, &key, &pid, BPF_ANY);
221 }
222
223 return 0;
224 }
225
BPF_STRUCT_OPS(peek_dsq_exit,struct scx_exit_info * ei)226 void BPF_STRUCT_OPS(peek_dsq_exit, struct scx_exit_info *ei)
227 {
228 int i;
229
230 /* Destroy the primary DSQs */
231 scx_bpf_destroy_dsq(test_dsq_id);
232 scx_bpf_destroy_dsq(real_dsq_id);
233
234 /* Destroy the DSQ pool */
235 bpf_for(i, 0, DSQ_POOL_SIZE) {
236 int dsq_id = dsq_pool_base_id + i;
237
238 scx_bpf_destroy_dsq(dsq_id);
239 }
240
241 UEI_RECORD(uei, ei);
242 }
243
244 SEC(".struct_ops.link")
245 struct sched_ext_ops peek_dsq_ops = {
246 .enqueue = (void *)peek_dsq_enqueue,
247 .dispatch = (void *)peek_dsq_dispatch,
248 .init = (void *)peek_dsq_init,
249 .exit = (void *)peek_dsq_exit,
250 .name = "peek_dsq",
251 };
252