diff options
Diffstat (limited to 'tools/testing/selftests/sched_ext/dequeue.bpf.c')
| -rw-r--r-- | tools/testing/selftests/sched_ext/dequeue.bpf.c | 389 |
1 files changed, 389 insertions, 0 deletions
diff --git a/tools/testing/selftests/sched_ext/dequeue.bpf.c b/tools/testing/selftests/sched_ext/dequeue.bpf.c new file mode 100644 index 000000000000..624e2ccb0688 --- /dev/null +++ b/tools/testing/selftests/sched_ext/dequeue.bpf.c @@ -0,0 +1,389 @@ +// SPDX-License-Identifier: GPL-2.0 +/* + * A scheduler that validates ops.dequeue() is called correctly: + * - Tasks dispatched to terminal DSQs (local, global) bypass the BPF + * scheduler entirely: no ops.dequeue() should be called + * - Tasks dispatched to user DSQs from ops.enqueue() enter BPF custody: + * ops.dequeue() must be called when they leave custody + * - Every ops.enqueue() dispatch to non-terminal DSQs is followed by + * exactly one ops.dequeue() (validate 1:1 pairing and state machine) + * + * Copyright (c) 2026 NVIDIA Corporation. + */ + +#include <scx/common.bpf.h> + +#define SHARED_DSQ 0 + +/* + * BPF internal queue. + * + * Tasks are stored here and consumed from ops.dispatch(), validating that + * tasks on BPF internal structures still get ops.dequeue() when they + * leave. + */ +struct { + __uint(type, BPF_MAP_TYPE_QUEUE); + __uint(max_entries, 32768); + __type(value, s32); +} global_queue SEC(".maps"); + +char _license[] SEC("license") = "GPL"; + +UEI_DEFINE(uei); + +/* + * Counters to track the lifecycle of tasks: + * - enqueue_cnt: Number of times ops.enqueue() was called + * - dequeue_cnt: Number of times ops.dequeue() was called (any type) + * - dispatch_dequeue_cnt: Number of regular dispatch dequeues (no flag) + * - change_dequeue_cnt: Number of property change dequeues + * - bpf_queue_full: Number of times the BPF internal queue was full + */ +u64 enqueue_cnt, dequeue_cnt, dispatch_dequeue_cnt, change_dequeue_cnt, bpf_queue_full; + +/* + * Test scenarios: + * 0) Dispatch to local DSQ from ops.select_cpu() (terminal DSQ, bypasses BPF + * scheduler, no dequeue callbacks) + * 1) Dispatch to global DSQ from ops.select_cpu() (terminal DSQ, bypasses BPF + * scheduler, no dequeue callbacks) + * 2) Dispatch to shared user DSQ from ops.select_cpu() (enters BPF scheduler, + * dequeue callbacks expected) + * 3) Dispatch to local DSQ from ops.enqueue() (terminal DSQ, bypasses BPF + * scheduler, no dequeue callbacks) + * 4) Dispatch to global DSQ from ops.enqueue() (terminal DSQ, bypasses BPF + * scheduler, no dequeue callbacks) + * 5) Dispatch to shared user DSQ from ops.enqueue() (enters BPF scheduler, + * dequeue callbacks expected) + * 6) BPF internal queue from ops.enqueue(): store task PIDs in ops.enqueue(), + * consume in ops.dispatch() and dispatch to local DSQ (validates dequeue + * for tasks stored in internal BPF data structures) + */ +u32 test_scenario; + +/* + * Per-task state to track lifecycle and validate workflow semantics. + * State transitions: + * NONE -> ENQUEUED (on enqueue) + * NONE -> DISPATCHED (on direct dispatch to terminal DSQ) + * ENQUEUED -> DISPATCHED (on dispatch dequeue) + * DISPATCHED -> NONE (on property change dequeue or re-enqueue) + * ENQUEUED -> NONE (on property change dequeue before dispatch) + */ +enum task_state { + TASK_NONE = 0, + TASK_ENQUEUED, + TASK_DISPATCHED, +}; + +struct task_ctx { + enum task_state state; /* Current state in the workflow */ + u64 enqueue_seq; /* Sequence number for debugging */ +}; + +struct { + __uint(type, BPF_MAP_TYPE_TASK_STORAGE); + __uint(map_flags, BPF_F_NO_PREALLOC); + __type(key, int); + __type(value, struct task_ctx); +} task_ctx_stor SEC(".maps"); + +static struct task_ctx *try_lookup_task_ctx(struct task_struct *p) +{ + return bpf_task_storage_get(&task_ctx_stor, p, 0, 0); +} + +s32 BPF_STRUCT_OPS(dequeue_select_cpu, struct task_struct *p, + s32 prev_cpu, u64 wake_flags) +{ + struct task_ctx *tctx; + + tctx = try_lookup_task_ctx(p); + if (!tctx) + return prev_cpu; + + switch (test_scenario) { + case 0: + /* + * Direct dispatch to the local DSQ. + * + * Task bypasses BPF scheduler entirely: no enqueue + * tracking, no ops.dequeue() callbacks. + */ + scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0); + tctx->state = TASK_DISPATCHED; + break; + case 1: + /* + * Direct dispatch to the global DSQ. + * + * Task bypasses BPF scheduler entirely: no enqueue + * tracking, no ops.dequeue() callbacks. + */ + scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0); + tctx->state = TASK_DISPATCHED; + break; + case 2: + /* + * Dispatch to a shared user DSQ. + * + * Task enters BPF scheduler management: track + * enqueue/dequeue lifecycle and validate state + * transitions. + */ + if (tctx->state == TASK_ENQUEUED) + scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu", + p->pid, p->comm, tctx->enqueue_seq); + + scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, 0); + + __sync_fetch_and_add(&enqueue_cnt, 1); + + tctx->state = TASK_ENQUEUED; + tctx->enqueue_seq++; + break; + } + + return prev_cpu; +} + +void BPF_STRUCT_OPS(dequeue_enqueue, struct task_struct *p, u64 enq_flags) +{ + struct task_ctx *tctx; + s32 pid = p->pid; + + tctx = try_lookup_task_ctx(p); + if (!tctx) + return; + + switch (test_scenario) { + case 3: + /* + * Direct dispatch to the local DSQ. + * + * Task bypasses BPF scheduler entirely: no enqueue + * tracking, no ops.dequeue() callbacks. + */ + scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, enq_flags); + tctx->state = TASK_DISPATCHED; + break; + case 4: + /* + * Direct dispatch to the global DSQ. + * + * Task bypasses BPF scheduler entirely: no enqueue + * tracking, no ops.dequeue() callbacks. + */ + scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); + tctx->state = TASK_DISPATCHED; + break; + case 5: + /* + * Dispatch to shared user DSQ. + * + * Task enters BPF scheduler management: track + * enqueue/dequeue lifecycle and validate state + * transitions. + */ + if (tctx->state == TASK_ENQUEUED) + scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu", + p->pid, p->comm, tctx->enqueue_seq); + + scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags); + + __sync_fetch_and_add(&enqueue_cnt, 1); + + tctx->state = TASK_ENQUEUED; + tctx->enqueue_seq++; + break; + case 6: + /* + * Store task in BPF internal queue. + * + * Task enters BPF scheduler management: track + * enqueue/dequeue lifecycle and validate state + * transitions. + */ + if (tctx->state == TASK_ENQUEUED) + scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu", + p->pid, p->comm, tctx->enqueue_seq); + + if (bpf_map_push_elem(&global_queue, &pid, 0)) { + scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); + __sync_fetch_and_add(&bpf_queue_full, 1); + + tctx->state = TASK_DISPATCHED; + } else { + __sync_fetch_and_add(&enqueue_cnt, 1); + + tctx->state = TASK_ENQUEUED; + tctx->enqueue_seq++; + } + break; + default: + /* For all other scenarios, dispatch to the global DSQ */ + scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); + tctx->state = TASK_DISPATCHED; + break; + } + + scx_bpf_kick_cpu(scx_bpf_task_cpu(p), SCX_KICK_IDLE); +} + +void BPF_STRUCT_OPS(dequeue_dequeue, struct task_struct *p, u64 deq_flags) +{ + struct task_ctx *tctx; + + __sync_fetch_and_add(&dequeue_cnt, 1); + + tctx = try_lookup_task_ctx(p); + if (!tctx) + return; + + /* + * For scenarios 0, 1, 3, and 4 (terminal DSQs: local and global), + * ops.dequeue() should never be called because tasks bypass the + * BPF scheduler entirely. If we get here, it's a kernel bug. + */ + if (test_scenario == 0 || test_scenario == 3) { + scx_bpf_error("%d (%s): dequeue called for local DSQ scenario", + p->pid, p->comm); + return; + } + + if (test_scenario == 1 || test_scenario == 4) { + scx_bpf_error("%d (%s): dequeue called for global DSQ scenario", + p->pid, p->comm); + return; + } + + if (deq_flags & SCX_DEQ_SCHED_CHANGE) { + /* + * Property change interrupting the workflow. Valid from + * both ENQUEUED and DISPATCHED states. Transitions task + * back to NONE state. + */ + __sync_fetch_and_add(&change_dequeue_cnt, 1); + + /* Validate state transition */ + if (tctx->state != TASK_ENQUEUED && tctx->state != TASK_DISPATCHED) + scx_bpf_error("%d (%s): invalid property change dequeue state=%d seq=%llu", + p->pid, p->comm, tctx->state, tctx->enqueue_seq); + + /* + * Transition back to NONE: task outside scheduler control. + * + * Scenario 6: dispatch() checks tctx->state after popping a + * PID, if the task is in state NONE, it was dequeued by + * property change and must not be dispatched (this + * prevents "target CPU not allowed"). + */ + tctx->state = TASK_NONE; + } else { + /* + * Regular dispatch dequeue: kernel is moving the task from + * BPF custody to a terminal DSQ. Normally we come from + * ENQUEUED state. We can also see TASK_NONE if the task + * was dequeued by property change (SCX_DEQ_SCHED_CHANGE) + * while it was already on a DSQ (dispatched but not yet + * consumed); in that case we just leave state as NONE. + */ + __sync_fetch_and_add(&dispatch_dequeue_cnt, 1); + + /* + * Must be ENQUEUED (normal path) or NONE (already dequeued + * by property change while on a DSQ). + */ + if (tctx->state != TASK_ENQUEUED && tctx->state != TASK_NONE) + scx_bpf_error("%d (%s): dispatch dequeue from state %d seq=%llu", + p->pid, p->comm, tctx->state, tctx->enqueue_seq); + + if (tctx->state == TASK_ENQUEUED) + tctx->state = TASK_DISPATCHED; + + /* NONE: leave as-is, task was already property-change dequeued */ + } +} + +void BPF_STRUCT_OPS(dequeue_dispatch, s32 cpu, struct task_struct *prev) +{ + if (test_scenario == 6) { + struct task_ctx *tctx; + struct task_struct *p; + s32 pid; + + if (bpf_map_pop_elem(&global_queue, &pid)) + return; + + p = bpf_task_from_pid(pid); + if (!p) + return; + + /* + * If the task was dequeued by property change + * (ops.dequeue() set tctx->state = TASK_NONE), skip + * dispatch. + */ + tctx = try_lookup_task_ctx(p); + if (!tctx || tctx->state == TASK_NONE) { + bpf_task_release(p); + return; + } + + /* + * Dispatch to this CPU's local DSQ if allowed, otherwise + * fallback to the global DSQ. + */ + if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) + scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, 0); + else + scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0); + + bpf_task_release(p); + } else { + scx_bpf_dsq_move_to_local(SHARED_DSQ, 0); + } +} + +s32 BPF_STRUCT_OPS(dequeue_init_task, struct task_struct *p, + struct scx_init_task_args *args) +{ + struct task_ctx *tctx; + + tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, + BPF_LOCAL_STORAGE_GET_F_CREATE); + if (!tctx) + return -ENOMEM; + + return 0; +} + +s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_init) +{ + s32 ret; + + ret = scx_bpf_create_dsq(SHARED_DSQ, -1); + if (ret) + return ret; + + return 0; +} + +void BPF_STRUCT_OPS(dequeue_exit, struct scx_exit_info *ei) +{ + UEI_RECORD(uei, ei); +} + +SEC(".struct_ops.link") +struct sched_ext_ops dequeue_ops = { + .select_cpu = (void *)dequeue_select_cpu, + .enqueue = (void *)dequeue_enqueue, + .dequeue = (void *)dequeue_dequeue, + .dispatch = (void *)dequeue_dispatch, + .init_task = (void *)dequeue_init_task, + .init = (void *)dequeue_init, + .exit = (void *)dequeue_exit, + .flags = SCX_OPS_ENQ_LAST, + .name = "dequeue_test", +}; |