tools/testing/selftests/sched_ext/peek_dsq.bpf.c - linux - Git at Google

 // SPDX-License-Identifier: GPL-2.0
 /*
  * A BPF program for testing DSQ operations and peek in particular.
  *
  * Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
  * Copyright (c) 2025 Ryan Newton <ryan.newton@alum.mit.edu>
  */

 #include <scx/common.bpf.h>
 #include <scx/compat.bpf.h>

 char _license[] SEC("license") = "GPL";

 UEI_DEFINE(uei); /* Error handling */

 #define MAX_SAMPLES 100
 #define MAX_CPUS 512
 #define DSQ_POOL_SIZE 8
 int max_samples = MAX_SAMPLES;
 int max_cpus = MAX_CPUS;
 int dsq_pool_size = DSQ_POOL_SIZE;

 /* Global variables to store test results */
 int dsq_peek_result1 = -1;
 long dsq_inserted_pid = -1;
 int insert_test_cpu = -1; /* Set to the cpu that performs the test */
 long dsq_peek_result2 = -1;
 long dsq_peek_result2_pid = -1;
 long dsq_peek_result2_expected = -1;
 int test_dsq_id = 1234; /* Use a simple ID like create_dsq example */
 int real_dsq_id = 1235; /* DSQ for normal operation */
 int enqueue_count = -1;
 int dispatch_count = -1;
 bool debug_ksym_exists;

 /* DSQ pool for stress testing */
 int dsq_pool_base_id = 2000;
 int phase1_complete = -1;
 long total_peek_attempts = -1;
 long successful_peeks = -1;

 /* BPF map for sharing peek results with userspace */
 struct {
 	__uint(type, BPF_MAP_TYPE_ARRAY);
 	__uint(max_entries, MAX_SAMPLES);
 	__type(key, u32);
 	__type(value, long);
 } peek_results SEC(".maps");

 static int get_random_dsq_id(void)
 {
 	u64 time = bpf_ktime_get_ns();

 	return dsq_pool_base_id + (time % DSQ_POOL_SIZE);
 }

 static void record_peek_result(long pid)
 {
 	u32 slot_key;
 	long *slot_pid_ptr;
 	u32 ix;

 	if (pid <= 0)
 		return;

 	/* Find an empty slot or one with the same PID */
 	bpf_for(ix, 0, 10) {
 		slot_key = ((u64)pid + ix) % MAX_SAMPLES;
 		slot_pid_ptr = bpf_map_lookup_elem(&peek_results, &slot_key);
 		if (!slot_pid_ptr)
 			continue;

 		if (*slot_pid_ptr == -1 || *slot_pid_ptr == pid) {
 			*slot_pid_ptr = pid;
 			break;
 		}
 	}
 }

 /* Scan all DSQs in the pool and try to move a task to local */
 static int scan_dsq_pool(void)
 {
 	struct task_struct *task;
 	int moved = 0;
 	int i;

 	bpf_for(i, 0, DSQ_POOL_SIZE) {
 		int dsq_id = dsq_pool_base_id + i;

 		total_peek_attempts++;

 		task = __COMPAT_scx_bpf_dsq_peek(dsq_id);
 		if (task) {
 			successful_peeks++;
 			record_peek_result(task->pid);

 			/* Try to move this task to local */
 			if (!moved && scx_bpf_dsq_move_to_local(dsq_id, 0)) {
 				moved = 1;
 				break;
 			}
 		}
 	}
 	return moved;
 }

 /* Struct_ops scheduler for testing DSQ peek operations */
 void BPF_STRUCT_OPS(peek_dsq_enqueue, struct task_struct *p, u64 enq_flags)
 {
 	struct task_struct *peek_result;
 	int last_insert_test_cpu, cpu;

 	enqueue_count++;
 	cpu = bpf_get_smp_processor_id();
 	last_insert_test_cpu = __sync_val_compare_and_swap(&insert_test_cpu, -1, cpu);

 	/* Phase 1: Simple insert-then-peek test (only on first task) */
 	if (last_insert_test_cpu == -1) {
 		bpf_printk("peek_dsq_enqueue beginning phase 1 peek test on cpu %d", cpu);

 		/* Test 1: Peek empty DSQ - should return NULL */
 		peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
 		dsq_peek_result1 = (long)peek_result; /* Should be 0 (NULL) */

 		/* Test 2: Insert task into test DSQ for testing in dispatch callback */
 		dsq_inserted_pid = p->pid;
 		scx_bpf_dsq_insert(p, test_dsq_id, 0, enq_flags);
 		dsq_peek_result2_expected = (long)p; /* Expected the task we just inserted */
 	} else if (!phase1_complete) {
 		/* Still in phase 1, use real DSQ */
 		scx_bpf_dsq_insert(p, real_dsq_id, 0, enq_flags);
 	} else {
 		/* Phase 2: Random DSQ insertion for stress testing */
 		int random_dsq_id = get_random_dsq_id();

 		scx_bpf_dsq_insert(p, random_dsq_id, 0, enq_flags);
 	}
 }

 void BPF_STRUCT_OPS(peek_dsq_dispatch, s32 cpu, struct task_struct *prev)
 {
 	dispatch_count++;

 	/* Phase 1: Complete the simple peek test if we inserted a task but
 	 * haven't tested peek yet
 	 */
 	if (insert_test_cpu == cpu && dsq_peek_result2 == -1) {
 		struct task_struct *peek_result;

 		bpf_printk("peek_dsq_dispatch completing phase 1 peek test on cpu %d", cpu);

 		/* Test 3: Peek DSQ after insert - should return the task we inserted */
 		peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
 		/* Store the PID of the peeked task for comparison */
 		dsq_peek_result2 = (long)peek_result;
 		dsq_peek_result2_pid = peek_result ? peek_result->pid : -1;

 		/* Now consume the task since we've peeked at it */
 		scx_bpf_dsq_move_to_local(test_dsq_id, 0);

 		/* Mark phase 1 as complete */
 		phase1_complete = 1;
 		bpf_printk("Phase 1 complete, starting phase 2 stress testing");
 	} else if (!phase1_complete) {
 		/* Still in phase 1, use real DSQ */
 		scx_bpf_dsq_move_to_local(real_dsq_id, 0);
 	} else {
 		/* Phase 2: Scan all DSQs in the pool and try to move a task */
 		if (!scan_dsq_pool()) {
 			/* No tasks found in DSQ pool, fall back to real DSQ */
 			scx_bpf_dsq_move_to_local(real_dsq_id, 0);
 		}
 	}
 }

 s32 BPF_STRUCT_OPS_SLEEPABLE(peek_dsq_init)
 {
 	s32 err;
 	int i;

 	/* Always set debug values so we can see which version we're using */
 	debug_ksym_exists = bpf_ksym_exists(scx_bpf_dsq_peek) ? 1 : 0;

 	/* Initialize state first */
 	insert_test_cpu = -1;
 	enqueue_count = 0;
 	dispatch_count = 0;
 	phase1_complete = 0;
 	total_peek_attempts = 0;
 	successful_peeks = 0;

 	/* Create the test and real DSQs */
 	err = scx_bpf_create_dsq(test_dsq_id, -1);
 	if (err) {
 		scx_bpf_error("Failed to create DSQ %d: %d", test_dsq_id, err);
 		return err;
 	}
 	err = scx_bpf_create_dsq(real_dsq_id, -1);
 	if (err) {
 		scx_bpf_error("Failed to create DSQ %d: %d", real_dsq_id, err);
 		return err;
 	}

 	/* Create the DSQ pool for stress testing */
 	bpf_for(i, 0, DSQ_POOL_SIZE) {
 		int dsq_id = dsq_pool_base_id + i;

 		err = scx_bpf_create_dsq(dsq_id, -1);
 		if (err) {
 			scx_bpf_error("Failed to create DSQ pool entry %d: %d", dsq_id, err);
 			return err;
 		}
 	}

 	/* Initialize the peek results map */
 	bpf_for(i, 0, MAX_SAMPLES) {
 		u32 key = i;
 		long pid = -1;

 		bpf_map_update_elem(&peek_results, &key, &pid, BPF_ANY);
 	}

 	return 0;
 }

 void BPF_STRUCT_OPS(peek_dsq_exit, struct scx_exit_info *ei)
 {
 	int i;

 	/* Destroy the primary DSQs */
 	scx_bpf_destroy_dsq(test_dsq_id);
 	scx_bpf_destroy_dsq(real_dsq_id);

 	/* Destroy the DSQ pool */
 	bpf_for(i, 0, DSQ_POOL_SIZE) {
 		int dsq_id = dsq_pool_base_id + i;

 		scx_bpf_destroy_dsq(dsq_id);
 	}

 	UEI_RECORD(uei, ei);
 }

 SEC(".struct_ops.link")
 struct sched_ext_ops peek_dsq_ops = {
 	.enqueue = (void *)peek_dsq_enqueue,
 	.dispatch = (void *)peek_dsq_dispatch,
 	.init = (void *)peek_dsq_init,
 	.exit = (void *)peek_dsq_exit,
 	.name = "peek_dsq",
 };
	// SPDX-License-Identifier: GPL-2.0
	/*
	* A BPF program for testing DSQ operations and peek in particular.
	*
	* Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
	* Copyright (c) 2025 Ryan Newton <ryan.newton@alum.mit.edu>
	*/

	#include <scx/common.bpf.h>
	#include <scx/compat.bpf.h>

	char _license[] SEC("license") = "GPL";

	UEI_DEFINE(uei); /* Error handling */

	#define MAX_SAMPLES 100
	#define MAX_CPUS 512
	#define DSQ_POOL_SIZE 8
	int max_samples = MAX_SAMPLES;
	int max_cpus = MAX_CPUS;
	int dsq_pool_size = DSQ_POOL_SIZE;

	/* Global variables to store test results */
	int dsq_peek_result1 = -1;
	long dsq_inserted_pid = -1;
	int insert_test_cpu = -1; /* Set to the cpu that performs the test */
	long dsq_peek_result2 = -1;
	long dsq_peek_result2_pid = -1;
	long dsq_peek_result2_expected = -1;
	int test_dsq_id = 1234; /* Use a simple ID like create_dsq example */
	int real_dsq_id = 1235; /* DSQ for normal operation */
	int enqueue_count = -1;
	int dispatch_count = -1;
	bool debug_ksym_exists;

	/* DSQ pool for stress testing */
	int dsq_pool_base_id = 2000;
	int phase1_complete = -1;
	long total_peek_attempts = -1;
	long successful_peeks = -1;

	/* BPF map for sharing peek results with userspace */
	struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, MAX_SAMPLES);
	__type(key, u32);
	__type(value, long);
	} peek_results SEC(".maps");

	static int get_random_dsq_id(void)
	{
	u64 time = bpf_ktime_get_ns();

	return dsq_pool_base_id + (time % DSQ_POOL_SIZE);
	}

	static void record_peek_result(long pid)
	{
	u32 slot_key;
	long *slot_pid_ptr;
	u32 ix;

	if (pid <= 0)
	return;

	/* Find an empty slot or one with the same PID */
	bpf_for(ix, 0, 10) {
	slot_key = ((u64)pid + ix) % MAX_SAMPLES;
	slot_pid_ptr = bpf_map_lookup_elem(&peek_results, &slot_key);
	if (!slot_pid_ptr)
	continue;

	if (slot_pid_ptr == -1 \|\| slot_pid_ptr == pid) {
	*slot_pid_ptr = pid;
	break;
	}
	}
	}

	/* Scan all DSQs in the pool and try to move a task to local */
	static int scan_dsq_pool(void)
	{
	struct task_struct *task;
	int moved = 0;
	int i;

	bpf_for(i, 0, DSQ_POOL_SIZE) {
	int dsq_id = dsq_pool_base_id + i;

	total_peek_attempts++;

	task = __COMPAT_scx_bpf_dsq_peek(dsq_id);
	if (task) {
	successful_peeks++;
	record_peek_result(task->pid);

	/* Try to move this task to local */
	if (!moved && scx_bpf_dsq_move_to_local(dsq_id, 0)) {
	moved = 1;
	break;
	}
	}
	}
	return moved;
	}

	/* Struct_ops scheduler for testing DSQ peek operations */
	void BPF_STRUCT_OPS(peek_dsq_enqueue, struct task_struct *p, u64 enq_flags)
	{
	struct task_struct *peek_result;
	int last_insert_test_cpu, cpu;

	enqueue_count++;
	cpu = bpf_get_smp_processor_id();
	last_insert_test_cpu = __sync_val_compare_and_swap(&insert_test_cpu, -1, cpu);

	/* Phase 1: Simple insert-then-peek test (only on first task) */
	if (last_insert_test_cpu == -1) {
	bpf_printk("peek_dsq_enqueue beginning phase 1 peek test on cpu %d", cpu);

	/* Test 1: Peek empty DSQ - should return NULL */
	peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
	dsq_peek_result1 = (long)peek_result; /* Should be 0 (NULL) */

	/* Test 2: Insert task into test DSQ for testing in dispatch callback */
	dsq_inserted_pid = p->pid;
	scx_bpf_dsq_insert(p, test_dsq_id, 0, enq_flags);
	dsq_peek_result2_expected = (long)p; /* Expected the task we just inserted */
	} else if (!phase1_complete) {
	/* Still in phase 1, use real DSQ */
	scx_bpf_dsq_insert(p, real_dsq_id, 0, enq_flags);
	} else {
	/* Phase 2: Random DSQ insertion for stress testing */
	int random_dsq_id = get_random_dsq_id();

	scx_bpf_dsq_insert(p, random_dsq_id, 0, enq_flags);
	}
	}

	void BPF_STRUCT_OPS(peek_dsq_dispatch, s32 cpu, struct task_struct *prev)
	{
	dispatch_count++;

	/* Phase 1: Complete the simple peek test if we inserted a task but
	* haven't tested peek yet
	*/
	if (insert_test_cpu == cpu && dsq_peek_result2 == -1) {
	struct task_struct *peek_result;

	bpf_printk("peek_dsq_dispatch completing phase 1 peek test on cpu %d", cpu);

	/* Test 3: Peek DSQ after insert - should return the task we inserted */
	peek_result = __COMPAT_scx_bpf_dsq_peek(test_dsq_id);
	/* Store the PID of the peeked task for comparison */
	dsq_peek_result2 = (long)peek_result;
	dsq_peek_result2_pid = peek_result ? peek_result->pid : -1;

	/* Now consume the task since we've peeked at it */
	scx_bpf_dsq_move_to_local(test_dsq_id, 0);

	/* Mark phase 1 as complete */
	phase1_complete = 1;
	bpf_printk("Phase 1 complete, starting phase 2 stress testing");
	} else if (!phase1_complete) {
	/* Still in phase 1, use real DSQ */
	scx_bpf_dsq_move_to_local(real_dsq_id, 0);
	} else {
	/* Phase 2: Scan all DSQs in the pool and try to move a task */
	if (!scan_dsq_pool()) {
	/* No tasks found in DSQ pool, fall back to real DSQ */
	scx_bpf_dsq_move_to_local(real_dsq_id, 0);
	}
	}
	}

	s32 BPF_STRUCT_OPS_SLEEPABLE(peek_dsq_init)
	{
	s32 err;
	int i;

	/* Always set debug values so we can see which version we're using */
	debug_ksym_exists = bpf_ksym_exists(scx_bpf_dsq_peek) ? 1 : 0;

	/* Initialize state first */
	insert_test_cpu = -1;
	enqueue_count = 0;
	dispatch_count = 0;
	phase1_complete = 0;
	total_peek_attempts = 0;
	successful_peeks = 0;

	/* Create the test and real DSQs */
	err = scx_bpf_create_dsq(test_dsq_id, -1);
	if (err) {
	scx_bpf_error("Failed to create DSQ %d: %d", test_dsq_id, err);
	return err;
	}
	err = scx_bpf_create_dsq(real_dsq_id, -1);
	if (err) {
	scx_bpf_error("Failed to create DSQ %d: %d", real_dsq_id, err);
	return err;
	}

	/* Create the DSQ pool for stress testing */
	bpf_for(i, 0, DSQ_POOL_SIZE) {
	int dsq_id = dsq_pool_base_id + i;

	err = scx_bpf_create_dsq(dsq_id, -1);
	if (err) {
	scx_bpf_error("Failed to create DSQ pool entry %d: %d", dsq_id, err);
	return err;
	}
	}

	/* Initialize the peek results map */
	bpf_for(i, 0, MAX_SAMPLES) {
	u32 key = i;
	long pid = -1;

	bpf_map_update_elem(&peek_results, &key, &pid, BPF_ANY);
	}

	return 0;
	}

	void BPF_STRUCT_OPS(peek_dsq_exit, struct scx_exit_info *ei)
	{
	int i;

	/* Destroy the primary DSQs */
	scx_bpf_destroy_dsq(test_dsq_id);
	scx_bpf_destroy_dsq(real_dsq_id);

	/* Destroy the DSQ pool */
	bpf_for(i, 0, DSQ_POOL_SIZE) {
	int dsq_id = dsq_pool_base_id + i;

	scx_bpf_destroy_dsq(dsq_id);
	}

	UEI_RECORD(uei, ei);
	}

	SEC(".struct_ops.link")
	struct sched_ext_ops peek_dsq_ops = {
	.enqueue = (void *)peek_dsq_enqueue,
	.dispatch = (void *)peek_dsq_dispatch,
	.init = (void *)peek_dsq_init,
	.exit = (void *)peek_dsq_exit,
	.name = "peek_dsq",
	};