1 // SPDX-License-Identifier: GPL-2.0 2 /* 3 * Verify that ops.dequeue() is called when a task leaves the BPF scheduler's 4 * custody by being moved to the local DSQ of a CPU other than the one whose 5 * rq it's on (move_remote_task_to_local_dsq()). 6 * 7 * ops.enqueue() puts every task into custody and ops.dispatch() moves it to 8 * the dispatching CPU's local DSQ, so most moves cross CPUs. With 9 * @test_use_move_to_local, tasks are queued on a user DSQ and consumed with 10 * scx_bpf_dsq_move_to_local(). Otherwise, they are queued in a BPF queue and 11 * dispatched with SCX_DSQ_LOCAL_ON. 12 * 13 * Copyright (c) 2026 Google LLC. 14 */ 15 16 #include <scx/common.bpf.h> 17 18 #define SHARED_DSQ 0 19 #define MAX_DISPATCH_POPS 8 20 21 char _license[] SEC("license") = "GPL"; 22 23 UEI_DEFINE(uei); 24 25 struct { 26 __uint(type, BPF_MAP_TYPE_QUEUE); 27 __uint(max_entries, 32768); 28 __type(value, s32); 29 } global_queue SEC(".maps"); 30 31 enum task_state { 32 TASK_NONE = 0, 33 TASK_ENQUEUED, /* in BPF custody, waiting for ops.dequeue() */ 34 TASK_DISPATCHED, /* left custody */ 35 }; 36 37 struct task_ctx { 38 enum task_state state; 39 s32 enq_cpu; /* scx_bpf_task_cpu() at ops.enqueue() */ 40 u64 enqueue_seq; 41 }; 42 43 struct { 44 __uint(type, BPF_MAP_TYPE_TASK_STORAGE); 45 __uint(map_flags, BPF_F_NO_PREALLOC); 46 __type(key, int); 47 __type(value, struct task_ctx); 48 } task_ctx_stor SEC(".maps"); 49 50 /* core_cookie only exists with CONFIG_SCHED_CORE */ 51 struct task_struct___core_sched { 52 unsigned long core_cookie; 53 } __attribute__((preserve_access_index)); 54 55 bool test_use_move_to_local; 56 57 u64 enqueue_cnt, dequeue_cnt, dispatch_dequeue_cnt, change_dequeue_cnt; 58 u64 remote_dispatch_cnt, remote_running_cnt, missed_dequeue_cnt; 59 u64 core_sched_exec_dequeue_cnt; 60 61 static struct task_ctx *lookup_task_ctx(struct task_struct *p) 62 { 63 return bpf_task_storage_get(&task_ctx_stor, p, 0, 0); 64 } 65 66 static bool task_has_core_cookie(struct task_struct *p) 67 { 68 struct task_struct___core_sched *t = (void *)p; 69 70 if (!bpf_core_field_exists(t->core_cookie)) 71 return false; 72 return BPF_CORE_READ(t, core_cookie); 73 } 74 75 s32 BPF_STRUCT_OPS(dequeue_remote_select_cpu, struct task_struct *p, 76 s32 prev_cpu, u64 wake_flags) 77 { 78 /* no direct dispatch, always go through ops.enqueue() */ 79 return prev_cpu; 80 } 81 82 void BPF_STRUCT_OPS(dequeue_remote_enqueue, struct task_struct *p, u64 enq_flags) 83 { 84 struct task_ctx *tctx; 85 s32 pid = p->pid; 86 87 tctx = lookup_task_ctx(p); 88 if (!tctx) { 89 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); 90 return; 91 } 92 93 /* the previous custody period must have ended with ops.dequeue() */ 94 if (tctx->state == TASK_ENQUEUED) 95 scx_bpf_error("%d (%s): enqueue while in ENQUEUED state seq=%llu", 96 p->pid, p->comm, tctx->enqueue_seq); 97 98 /* 99 * Mark @p as enqueued before making it visible to ops.dispatch() on 100 * other CPUs, which skips queue entries of tasks not in ENQUEUED 101 * state as stale. 102 */ 103 tctx->state = TASK_ENQUEUED; 104 tctx->enq_cpu = scx_bpf_task_cpu(p); 105 tctx->enqueue_seq++; 106 107 if (test_use_move_to_local) { 108 scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags); 109 } else if (bpf_map_push_elem(&global_queue, &pid, 0)) { 110 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags); 111 tctx->state = TASK_DISPATCHED; 112 tctx->enq_cpu = -1; 113 goto out; 114 } 115 116 __sync_fetch_and_add(&enqueue_cnt, 1); 117 out: 118 scx_bpf_kick_cpu(scx_bpf_task_cpu(p), SCX_KICK_IDLE); 119 } 120 121 void BPF_STRUCT_OPS(dequeue_remote_dequeue, struct task_struct *p, u64 deq_flags) 122 { 123 struct task_ctx *tctx; 124 125 __sync_fetch_and_add(&dequeue_cnt, 1); 126 127 tctx = lookup_task_ctx(p); 128 if (!tctx) 129 return; 130 131 /* 132 * Only core scheduling can pick a task straight out of custody, and 133 * only if the task has a core cookie. Otherwise, the custody exit was 134 * missed when @p was inserted into a local DSQ and got deferred until 135 * @p was picked. 136 */ 137 if ((deq_flags & SCX_DEQ_CORE_SCHED_EXEC) && !task_has_core_cookie(p)) { 138 __sync_fetch_and_add(&core_sched_exec_dequeue_cnt, 1); 139 scx_bpf_error("%d (%s): late ops.dequeue() with SCX_DEQ_CORE_SCHED_EXEC (enq_cpu=%d cpu=%d seq=%llu)", 140 p->pid, p->comm, tctx->enq_cpu, 141 scx_bpf_task_cpu(p), tctx->enqueue_seq); 142 } 143 144 /* ops.dequeue() ends the custody period started by ops.enqueue() */ 145 if (tctx->state != TASK_ENQUEUED) 146 scx_bpf_error("%d (%s): dequeue outside custody deq_flags=0x%llx state=%d seq=%llu", 147 p->pid, p->comm, deq_flags, tctx->state, 148 tctx->enqueue_seq); 149 150 if (deq_flags & SCX_DEQ_SCHED_CHANGE) { 151 __sync_fetch_and_add(&change_dequeue_cnt, 1); 152 tctx->state = TASK_NONE; 153 } else { 154 __sync_fetch_and_add(&dispatch_dequeue_cnt, 1); 155 tctx->state = TASK_DISPATCHED; 156 } 157 } 158 159 void BPF_STRUCT_OPS(dequeue_remote_dispatch, s32 cpu, struct task_struct *prev) 160 { 161 struct task_ctx *tctx; 162 struct task_struct *p; 163 s32 pid; 164 int i; 165 166 if (test_use_move_to_local) { 167 scx_bpf_dsq_move_to_local(SHARED_DSQ, 0); 168 return; 169 } 170 171 /* pop past stale entries so that they don't leave this CPU idle */ 172 bpf_for(i, 0, MAX_DISPATCH_POPS) { 173 if (bpf_map_pop_elem(&global_queue, &pid)) 174 return; 175 176 p = bpf_task_from_pid(pid); 177 if (!p) 178 continue; 179 180 /* 181 * Entries are stale if @p left custody through a property 182 * change dequeue or was dispatched from a duplicate entry. 183 */ 184 tctx = lookup_task_ctx(p); 185 if (!tctx || tctx->state != TASK_ENQUEUED) { 186 bpf_task_release(p); 187 continue; 188 } 189 190 if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) { 191 if (scx_bpf_task_cpu(p) != cpu) 192 __sync_fetch_and_add(&remote_dispatch_cnt, 1); 193 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, 194 SCX_SLICE_DFL, 0); 195 } else { 196 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0); 197 } 198 199 bpf_task_release(p); 200 return; 201 } 202 203 /* out of pops with entries left, retry instead of idling this CPU */ 204 if (!bpf_map_peek_elem(&global_queue, &pid)) 205 scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE); 206 } 207 208 void BPF_STRUCT_OPS(dequeue_remote_running, struct task_struct *p) 209 { 210 struct task_ctx *tctx; 211 212 tctx = lookup_task_ctx(p); 213 if (!tctx) 214 return; 215 216 /* tasks can only run from a local DSQ, i.e. after leaving custody */ 217 if (tctx->state == TASK_ENQUEUED) { 218 __sync_fetch_and_add(&missed_dequeue_cnt, 1); 219 scx_bpf_error("%d (%s): running without ops.dequeue() (enq_cpu=%d cpu=%d seq=%llu)", 220 p->pid, p->comm, tctx->enq_cpu, 221 scx_bpf_task_cpu(p), tctx->enqueue_seq); 222 return; 223 } 224 225 if (tctx->enq_cpu >= 0 && tctx->enq_cpu != scx_bpf_task_cpu(p)) 226 __sync_fetch_and_add(&remote_running_cnt, 1); 227 tctx->enq_cpu = -1; 228 } 229 230 s32 BPF_STRUCT_OPS(dequeue_remote_init_task, struct task_struct *p, 231 struct scx_init_task_args *args) 232 { 233 struct task_ctx *tctx; 234 235 tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 236 BPF_LOCAL_STORAGE_GET_F_CREATE); 237 if (!tctx) 238 return -ENOMEM; 239 240 /* task storage persists across attachments, start from scratch */ 241 tctx->state = TASK_NONE; 242 tctx->enq_cpu = -1; 243 tctx->enqueue_seq = 0; 244 245 return 0; 246 } 247 248 s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_remote_init) 249 { 250 return scx_bpf_create_dsq(SHARED_DSQ, -1); 251 } 252 253 void BPF_STRUCT_OPS(dequeue_remote_exit, struct scx_exit_info *ei) 254 { 255 UEI_RECORD(uei, ei); 256 } 257 258 SEC(".struct_ops.link") 259 struct sched_ext_ops dequeue_remote_ops = { 260 .select_cpu = (void *)dequeue_remote_select_cpu, 261 .enqueue = (void *)dequeue_remote_enqueue, 262 .dequeue = (void *)dequeue_remote_dequeue, 263 .dispatch = (void *)dequeue_remote_dispatch, 264 .running = (void *)dequeue_remote_running, 265 .init_task = (void *)dequeue_remote_init_task, 266 .init = (void *)dequeue_remote_init, 267 .exit = (void *)dequeue_remote_exit, 268 .flags = SCX_OPS_ENQ_LAST, 269 .name = "dequeue_remote", 270 }; 271