peek_dsq.bpf.c

  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
 12char _license[] SEC("license") = "GPL";
 13
 14UEI_DEFINE(uei); /* Error handling */
 15
 16#define MAX_SAMPLES 100
 17#define MAX_CPUS 512
 18#define DSQ_POOL_SIZE 8
 19int max_samples = MAX_SAMPLES;
 20int max_cpus = MAX_CPUS;
 21int dsq_pool_size = DSQ_POOL_SIZE;
 22
 23/* Global variables to store test results */
 24int dsq_peek_result1 = -1;
 25long dsq_inserted_pid = -1;
 26int insert_test_cpu = -1; /* Set to the cpu that performs the test */
 27long dsq_peek_result2 = -1;
 28long dsq_peek_result2_pid = -1;
 29long dsq_peek_result2_expected = -1;
 30int test_dsq_id = 1234; /* Use a simple ID like create_dsq example */
 31int real_dsq_id = 1235; /* DSQ for normal operation */
 32int enqueue_count = -1;
 33int dispatch_count = -1;
 34bool debug_ksym_exists;
 35
 36/* DSQ pool for stress testing */
 37int dsq_pool_base_id = 2000;
 38int phase1_complete = -1;
 39long total_peek_attempts = -1;
 40long successful_peeks = -1;
 41
 42/* BPF map for sharing peek results with userspace */
 43struct {
 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
 50static 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
 57static 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 */
 81static 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 */
108void 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
140void 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
176s32 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
226void 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
244SEC(".struct_ops.link")
245struct 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};