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 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 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 */ 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 */ 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 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 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 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