1 /* SPDX-License-Identifier: GPL-2.0 */ 2 /* 3 * A demo sched_ext flattened cgroup hierarchy scheduler. It implements 4 * hierarchical weight-based cgroup CPU control by flattening the cgroup 5 * hierarchy into a single layer by compounding the active weight share at each 6 * level. Consider the following hierarchy with weights in parentheses: 7 * 8 * R + A (100) + B (100) 9 * | \ C (100) 10 * \ D (200) 11 * 12 * Ignoring the root and threaded cgroups, only B, C and D can contain tasks. 13 * Let's say all three have runnable tasks. The total share that each of these 14 * three cgroups is entitled to can be calculated by compounding its share at 15 * each level. 16 * 17 * For example, B is competing against C and in that competition its share is 18 * 100/(100+100) == 1/2. At its parent level, A is competing against D and A's 19 * share in that competition is 100/(200+100) == 1/3. B's eventual share in the 20 * system can be calculated by multiplying the two shares, 1/2 * 1/3 == 1/6. C's 21 * eventual share is the same at 1/6. D is only competing at the top level and 22 * its share is 200/(100+200) == 2/3. 23 * 24 * So, instead of hierarchically scheduling level-by-level, we can consider it 25 * as B, C and D competing each other with respective share of 1/6, 1/6 and 2/3 26 * and keep updating the eventual shares as the cgroups' runnable states change. 27 * 28 * This flattening of hierarchy can bring a substantial performance gain when 29 * the cgroup hierarchy is nested multiple levels. in a simple benchmark using 30 * wrk[8] on apache serving a CGI script calculating sha1sum of a small file, it 31 * outperforms CFS by ~3% with CPU controller disabled and by ~10% with two 32 * apache instances competing with 2:1 weight ratio nested four level deep. 33 * 34 * However, the gain comes at the cost of not being able to properly handle 35 * thundering herd of cgroups. For example, if many cgroups which are nested 36 * behind a low priority parent cgroup wake up around the same time, they may be 37 * able to consume more CPU cycles than they are entitled to. In many use cases, 38 * this isn't a real concern especially given the performance gain. Also, there 39 * are ways to mitigate the problem further by e.g. introducing an extra 40 * scheduling layer on cgroup delegation boundaries. 41 * 42 * The scheduler first picks the cgroup to run and then schedule the tasks 43 * within by using nested weighted vtime scheduling by default. The 44 * cgroup-internal scheduling can be switched to FIFO with the -f option. 45 */ 46 #include <scx/common.bpf.h> 47 #include "scx_flatcg.h" 48 49 /* 50 * Maximum amount of retries to find a valid cgroup. 51 */ 52 enum { 53 FALLBACK_DSQ = 0, 54 CGROUP_MAX_RETRIES = 1024, 55 }; 56 57 char _license[] SEC("license") = "GPL"; 58 59 const volatile u32 nr_cpus = 32; /* !0 for veristat, set during init */ 60 const volatile u64 cgrp_slice_ns; 61 const volatile bool fifo_sched; 62 63 u64 cvtime_now; 64 UEI_DEFINE(uei); 65 66 struct { 67 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY); 68 __type(key, u32); 69 __type(value, u64); 70 __uint(max_entries, FCG_NR_STATS); 71 } stats SEC(".maps"); 72 73 static void stat_inc(enum fcg_stat_idx idx) 74 { 75 u32 idx_v = idx; 76 77 u64 *cnt_p = bpf_map_lookup_elem(&stats, &idx_v); 78 if (cnt_p) 79 (*cnt_p)++; 80 } 81 82 struct fcg_cpu_ctx { 83 u64 cur_cgid; 84 u64 cur_at; 85 }; 86 87 struct { 88 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY); 89 __type(key, u32); 90 __type(value, struct fcg_cpu_ctx); 91 __uint(max_entries, 1); 92 } cpu_ctx SEC(".maps"); 93 94 struct { 95 __uint(type, BPF_MAP_TYPE_CGRP_STORAGE); 96 __uint(map_flags, BPF_F_NO_PREALLOC); 97 __type(key, int); 98 __type(value, struct fcg_cgrp_ctx); 99 } cgrp_ctx SEC(".maps"); 100 101 struct cgv_node { 102 struct bpf_rb_node rb_node; 103 __u64 cvtime; 104 __u64 cgid; 105 }; 106 107 private(CGV_TREE) struct bpf_spin_lock cgv_tree_lock; 108 private(CGV_TREE) struct bpf_rb_root cgv_tree __contains(cgv_node, rb_node); 109 110 struct cgv_node_stash { 111 struct cgv_node __kptr *node; 112 }; 113 114 struct { 115 __uint(type, BPF_MAP_TYPE_HASH); 116 __uint(max_entries, 16384); 117 __type(key, __u64); 118 __type(value, struct cgv_node_stash); 119 } cgv_node_stash SEC(".maps"); 120 121 struct fcg_task_ctx { 122 u64 bypassed_at; 123 }; 124 125 struct { 126 __uint(type, BPF_MAP_TYPE_TASK_STORAGE); 127 __uint(map_flags, BPF_F_NO_PREALLOC); 128 __type(key, int); 129 __type(value, struct fcg_task_ctx); 130 } task_ctx SEC(".maps"); 131 132 /* gets inc'd on weight tree changes to expire the cached hweights */ 133 u64 hweight_gen = 1; 134 135 static u64 div_round_up(u64 dividend, u64 divisor) 136 { 137 return (dividend + divisor - 1) / divisor; 138 } 139 140 static bool cgv_node_less(struct bpf_rb_node *a, const struct bpf_rb_node *b) 141 { 142 struct cgv_node *cgc_a, *cgc_b; 143 144 cgc_a = container_of(a, struct cgv_node, rb_node); 145 cgc_b = container_of(b, struct cgv_node, rb_node); 146 147 return cgc_a->cvtime < cgc_b->cvtime; 148 } 149 150 static struct fcg_cpu_ctx *find_cpu_ctx(void) 151 { 152 struct fcg_cpu_ctx *cpuc; 153 u32 idx = 0; 154 155 cpuc = bpf_map_lookup_elem(&cpu_ctx, &idx); 156 if (!cpuc) { 157 scx_bpf_error("cpu_ctx lookup failed"); 158 return NULL; 159 } 160 return cpuc; 161 } 162 163 static struct fcg_cgrp_ctx *find_cgrp_ctx(struct cgroup *cgrp) 164 { 165 struct fcg_cgrp_ctx *cgc; 166 167 cgc = bpf_cgrp_storage_get(&cgrp_ctx, cgrp, 0, 0); 168 if (!cgc) { 169 scx_bpf_error("cgrp_ctx lookup failed for cgid %llu", cgrp->kn->id); 170 return NULL; 171 } 172 return cgc; 173 } 174 175 static struct fcg_cgrp_ctx *find_ancestor_cgrp_ctx(struct cgroup *cgrp, int level) 176 { 177 struct fcg_cgrp_ctx *cgc; 178 179 cgrp = bpf_cgroup_ancestor(cgrp, level); 180 if (!cgrp) { 181 scx_bpf_error("ancestor cgroup lookup failed"); 182 return NULL; 183 } 184 185 cgc = find_cgrp_ctx(cgrp); 186 if (!cgc) 187 scx_bpf_error("ancestor cgrp_ctx lookup failed"); 188 bpf_cgroup_release(cgrp); 189 return cgc; 190 } 191 192 static void cgrp_refresh_hweight(struct cgroup *cgrp, struct fcg_cgrp_ctx *cgc) 193 { 194 int level; 195 196 if (!cgc->nr_active) { 197 stat_inc(FCG_STAT_HWT_SKIP); 198 return; 199 } 200 201 if (cgc->hweight_gen == hweight_gen) { 202 stat_inc(FCG_STAT_HWT_CACHE); 203 return; 204 } 205 206 stat_inc(FCG_STAT_HWT_UPDATES); 207 bpf_for(level, 0, cgrp->level + 1) { 208 struct fcg_cgrp_ctx *cgc; 209 bool is_active; 210 211 cgc = find_ancestor_cgrp_ctx(cgrp, level); 212 if (!cgc) 213 break; 214 215 if (!level) { 216 cgc->hweight = FCG_HWEIGHT_ONE; 217 cgc->hweight_gen = hweight_gen; 218 } else { 219 struct fcg_cgrp_ctx *pcgc; 220 221 pcgc = find_ancestor_cgrp_ctx(cgrp, level - 1); 222 if (!pcgc) 223 break; 224 225 /* 226 * We can be opportunistic here and not grab the 227 * cgv_tree_lock and deal with the occasional races. 228 * However, hweight updates are already cached and 229 * relatively low-frequency. Let's just do the 230 * straightforward thing. 231 */ 232 bpf_spin_lock(&cgv_tree_lock); 233 is_active = cgc->nr_active; 234 if (is_active) { 235 cgc->hweight_gen = pcgc->hweight_gen; 236 cgc->hweight = 237 div_round_up(pcgc->hweight * cgc->weight, 238 pcgc->child_weight_sum); 239 } 240 bpf_spin_unlock(&cgv_tree_lock); 241 242 if (!is_active) { 243 stat_inc(FCG_STAT_HWT_RACE); 244 break; 245 } 246 } 247 } 248 } 249 250 static void cgrp_cap_budget(struct cgv_node *cgv_node, struct fcg_cgrp_ctx *cgc) 251 { 252 u64 delta, cvtime, max_budget; 253 254 /* 255 * A node which is on the rbtree can't be pointed to from elsewhere yet 256 * and thus can't be updated and repositioned. Instead, we collect the 257 * vtime deltas separately and apply it asynchronously here. 258 */ 259 delta = __sync_fetch_and_and(&cgc->cvtime_delta, 0); 260 cvtime = cgv_node->cvtime + delta; 261 262 /* 263 * Allow a cgroup to carry the maximum budget proportional to its 264 * hweight such that a full-hweight cgroup can immediately take up half 265 * of the CPUs at the most while staying at the front of the rbtree. 266 */ 267 max_budget = (cgrp_slice_ns * nr_cpus * cgc->hweight) / 268 (2 * FCG_HWEIGHT_ONE); 269 if (time_before(cvtime, cvtime_now - max_budget)) 270 cvtime = cvtime_now - max_budget; 271 272 cgv_node->cvtime = cvtime; 273 } 274 275 static void cgrp_enqueued(struct cgroup *cgrp, struct fcg_cgrp_ctx *cgc) 276 { 277 struct cgv_node_stash *stash; 278 struct cgv_node *cgv_node; 279 u64 cgid = cgrp->kn->id; 280 281 /* paired with cmpxchg in try_pick_next_cgroup() */ 282 if (__sync_val_compare_and_swap(&cgc->queued, 0, 1)) { 283 stat_inc(FCG_STAT_ENQ_SKIP); 284 return; 285 } 286 287 stash = bpf_map_lookup_elem(&cgv_node_stash, &cgid); 288 if (!stash) { 289 scx_bpf_error("cgv_node lookup failed for cgid %llu", cgid); 290 return; 291 } 292 293 /* NULL if the node is already on the rbtree */ 294 cgv_node = bpf_kptr_xchg(&stash->node, NULL); 295 if (!cgv_node) { 296 stat_inc(FCG_STAT_ENQ_RACE); 297 return; 298 } 299 300 bpf_spin_lock(&cgv_tree_lock); 301 cgrp_cap_budget(cgv_node, cgc); 302 bpf_rbtree_add(&cgv_tree, &cgv_node->rb_node, cgv_node_less); 303 bpf_spin_unlock(&cgv_tree_lock); 304 } 305 306 static void set_bypassed_at(struct task_struct *p, struct fcg_task_ctx *taskc) 307 { 308 /* 309 * Tell fcg_stopping() that this bypassed the regular scheduling path 310 * and should be force charged to the cgroup. 0 is used to indicate that 311 * the task isn't bypassing, so if the current runtime is 0, go back by 312 * one nanosecond. 313 */ 314 taskc->bypassed_at = p->se.sum_exec_runtime ?: (u64)-1; 315 } 316 317 s32 BPF_STRUCT_OPS(fcg_select_cpu, struct task_struct *p, s32 prev_cpu, u64 wake_flags) 318 { 319 struct fcg_task_ctx *taskc; 320 bool is_idle = false; 321 s32 cpu; 322 323 cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, wake_flags, &is_idle); 324 325 taskc = bpf_task_storage_get(&task_ctx, p, 0, 0); 326 if (!taskc) { 327 scx_bpf_error("task_ctx lookup failed"); 328 return cpu; 329 } 330 331 /* 332 * If select_cpu_dfl() is recommending local enqueue, the target CPU is 333 * idle. Follow it and charge the cgroup later in fcg_stopping() after 334 * the fact. 335 */ 336 if (is_idle) { 337 set_bypassed_at(p, taskc); 338 stat_inc(FCG_STAT_LOCAL); 339 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0); 340 } 341 342 return cpu; 343 } 344 345 void BPF_STRUCT_OPS(fcg_enqueue, struct task_struct *p, u64 enq_flags) 346 { 347 struct fcg_task_ctx *taskc; 348 struct cgroup *cgrp; 349 struct fcg_cgrp_ctx *cgc; 350 351 taskc = bpf_task_storage_get(&task_ctx, p, 0, 0); 352 if (!taskc) { 353 scx_bpf_error("task_ctx lookup failed"); 354 return; 355 } 356 357 /* 358 * Use the direct dispatching and force charging to deal with tasks with 359 * custom affinities so that we don't have to worry about per-cgroup 360 * dq's containing tasks that can't be executed from some CPUs. 361 */ 362 if (p->nr_cpus_allowed != nr_cpus) { 363 set_bypassed_at(p, taskc); 364 365 /* 366 * The global dq is deprioritized as we don't want to let tasks 367 * to boost themselves by constraining its cpumask. The 368 * deprioritization is rather severe, so let's not apply that to 369 * per-cpu kernel threads. This is ham-fisted. We probably wanna 370 * implement per-cgroup fallback dq's instead so that we have 371 * more control over when tasks with custom cpumask get issued. 372 */ 373 if (p->nr_cpus_allowed == 1 && (p->flags & PF_KTHREAD)) { 374 stat_inc(FCG_STAT_LOCAL); 375 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 376 enq_flags); 377 } else { 378 stat_inc(FCG_STAT_GLOBAL); 379 scx_bpf_dsq_insert(p, FALLBACK_DSQ, SCX_SLICE_DFL, 380 enq_flags); 381 } 382 return; 383 } 384 385 cgrp = scx_bpf_task_cgroup(p); 386 cgc = find_cgrp_ctx(cgrp); 387 if (!cgc) 388 goto out_release; 389 390 if (fifo_sched) { 391 scx_bpf_dsq_insert(p, cgrp->kn->id, SCX_SLICE_DFL, enq_flags); 392 } else { 393 u64 tvtime = p->scx.dsq_vtime; 394 395 /* 396 * Limit the amount of budget that an idling task can accumulate 397 * to one slice. 398 */ 399 if (time_before(tvtime, cgc->tvtime_now - SCX_SLICE_DFL)) 400 tvtime = cgc->tvtime_now - SCX_SLICE_DFL; 401 402 scx_bpf_dsq_insert_vtime(p, cgrp->kn->id, SCX_SLICE_DFL, 403 tvtime, enq_flags); 404 } 405 406 cgrp_enqueued(cgrp, cgc); 407 out_release: 408 bpf_cgroup_release(cgrp); 409 } 410 411 /* 412 * Walk the cgroup tree to update the active weight sums as tasks wake up and 413 * sleep. The weight sums are used as the base when calculating the proportion a 414 * given cgroup or task is entitled to at each level. 415 */ 416 static void update_active_weight_sums(struct cgroup *cgrp, bool runnable) 417 { 418 struct fcg_cgrp_ctx *cgc; 419 bool updated = false; 420 int idx; 421 422 cgc = find_cgrp_ctx(cgrp); 423 if (!cgc) 424 return; 425 426 /* 427 * In most cases, a hot cgroup would have multiple threads going to 428 * sleep and waking up while the whole cgroup stays active. In leaf 429 * cgroups, ->nr_runnable which is updated with __sync operations gates 430 * ->nr_active updates, so that we don't have to grab the cgv_tree_lock 431 * repeatedly for a busy cgroup which is staying active. 432 */ 433 if (runnable) { 434 if (__sync_fetch_and_add(&cgc->nr_runnable, 1)) 435 return; 436 stat_inc(FCG_STAT_ACT); 437 } else { 438 if (__sync_sub_and_fetch(&cgc->nr_runnable, 1)) 439 return; 440 stat_inc(FCG_STAT_DEACT); 441 } 442 443 /* 444 * If @cgrp is becoming runnable, its hweight should be refreshed after 445 * it's added to the weight tree so that enqueue has the up-to-date 446 * value. If @cgrp is becoming quiescent, the hweight should be 447 * refreshed before it's removed from the weight tree so that the usage 448 * charging which happens afterwards has access to the latest value. 449 */ 450 if (!runnable) 451 cgrp_refresh_hweight(cgrp, cgc); 452 453 /* propagate upwards */ 454 bpf_for(idx, 0, cgrp->level) { 455 int level = cgrp->level - idx; 456 struct fcg_cgrp_ctx *cgc, *pcgc = NULL; 457 bool propagate = false; 458 459 cgc = find_ancestor_cgrp_ctx(cgrp, level); 460 if (!cgc) 461 break; 462 if (level) { 463 pcgc = find_ancestor_cgrp_ctx(cgrp, level - 1); 464 if (!pcgc) 465 break; 466 } 467 468 /* 469 * We need the propagation protected by a lock to synchronize 470 * against weight changes. There's no reason to drop the lock at 471 * each level but bpf_spin_lock() doesn't want any function 472 * calls while locked. 473 */ 474 bpf_spin_lock(&cgv_tree_lock); 475 476 if (runnable) { 477 if (!cgc->nr_active++) { 478 updated = true; 479 if (pcgc) { 480 propagate = true; 481 pcgc->child_weight_sum += cgc->weight; 482 } 483 } 484 } else { 485 if (!--cgc->nr_active) { 486 updated = true; 487 if (pcgc) { 488 propagate = true; 489 pcgc->child_weight_sum -= cgc->weight; 490 } 491 } 492 } 493 494 bpf_spin_unlock(&cgv_tree_lock); 495 496 if (!propagate) 497 break; 498 } 499 500 if (updated) 501 __sync_fetch_and_add(&hweight_gen, 1); 502 503 if (runnable) 504 cgrp_refresh_hweight(cgrp, cgc); 505 } 506 507 void BPF_STRUCT_OPS(fcg_runnable, struct task_struct *p, u64 enq_flags) 508 { 509 struct cgroup *cgrp; 510 511 cgrp = scx_bpf_task_cgroup(p); 512 update_active_weight_sums(cgrp, true); 513 bpf_cgroup_release(cgrp); 514 } 515 516 void BPF_STRUCT_OPS(fcg_running, struct task_struct *p) 517 { 518 struct cgroup *cgrp; 519 struct fcg_cgrp_ctx *cgc; 520 521 if (fifo_sched) 522 return; 523 524 cgrp = scx_bpf_task_cgroup(p); 525 cgc = find_cgrp_ctx(cgrp); 526 if (cgc) { 527 /* 528 * @cgc->tvtime_now always progresses forward as tasks start 529 * executing. The test and update can be performed concurrently 530 * from multiple CPUs and thus racy. Any error should be 531 * contained and temporary. Let's just live with it. 532 */ 533 if (time_before(cgc->tvtime_now, p->scx.dsq_vtime)) 534 cgc->tvtime_now = p->scx.dsq_vtime; 535 } 536 bpf_cgroup_release(cgrp); 537 } 538 539 void BPF_STRUCT_OPS(fcg_stopping, struct task_struct *p, bool runnable) 540 { 541 struct fcg_task_ctx *taskc; 542 struct cgroup *cgrp; 543 struct fcg_cgrp_ctx *cgc; 544 545 /* 546 * Scale the execution time by the inverse of the weight and charge. 547 * 548 * Note that the default yield implementation yields by setting 549 * @p->scx.slice to zero and the following would treat the yielding task 550 * as if it has consumed all its slice. If this penalizes yielding tasks 551 * too much, determine the execution time by taking explicit timestamps 552 * instead of depending on @p->scx.slice. 553 */ 554 if (!fifo_sched) { 555 u64 delta = scale_by_task_weight_inverse(p, SCX_SLICE_DFL - p->scx.slice); 556 557 scx_bpf_task_set_dsq_vtime(p, p->scx.dsq_vtime + delta); 558 } 559 560 taskc = bpf_task_storage_get(&task_ctx, p, 0, 0); 561 if (!taskc) { 562 scx_bpf_error("task_ctx lookup failed"); 563 return; 564 } 565 566 if (!taskc->bypassed_at) 567 return; 568 569 cgrp = scx_bpf_task_cgroup(p); 570 cgc = find_cgrp_ctx(cgrp); 571 if (cgc) { 572 __sync_fetch_and_add(&cgc->cvtime_delta, 573 (p->se.sum_exec_runtime - taskc->bypassed_at) * 574 FCG_HWEIGHT_ONE / (cgc->hweight ?: 1)); 575 taskc->bypassed_at = 0; 576 } 577 bpf_cgroup_release(cgrp); 578 } 579 580 void BPF_STRUCT_OPS(fcg_quiescent, struct task_struct *p, u64 deq_flags) 581 { 582 struct cgroup *cgrp; 583 584 cgrp = scx_bpf_task_cgroup(p); 585 update_active_weight_sums(cgrp, false); 586 bpf_cgroup_release(cgrp); 587 } 588 589 void BPF_STRUCT_OPS(fcg_cgroup_set_weight, struct cgroup *cgrp, u32 weight) 590 { 591 struct fcg_cgrp_ctx *cgc, *pcgc = NULL; 592 593 cgc = find_cgrp_ctx(cgrp); 594 if (!cgc) 595 return; 596 597 if (cgrp->level) { 598 pcgc = find_ancestor_cgrp_ctx(cgrp, cgrp->level - 1); 599 if (!pcgc) 600 return; 601 } 602 603 bpf_spin_lock(&cgv_tree_lock); 604 if (pcgc && cgc->nr_active) 605 pcgc->child_weight_sum += (s64)weight - cgc->weight; 606 cgc->weight = weight; 607 bpf_spin_unlock(&cgv_tree_lock); 608 609 /* expire cached hweights so the new weight propagates */ 610 __sync_fetch_and_add(&hweight_gen, 1); 611 } 612 613 static bool try_pick_next_cgroup(u64 *cgidp) 614 { 615 struct bpf_rb_node *rb_node; 616 struct cgv_node_stash *stash; 617 struct cgv_node *cgv_node; 618 struct fcg_cgrp_ctx *cgc; 619 struct cgroup *cgrp; 620 u64 cgid; 621 622 /* pop the front cgroup and wind cvtime_now accordingly */ 623 bpf_spin_lock(&cgv_tree_lock); 624 625 rb_node = bpf_rbtree_first(&cgv_tree); 626 if (!rb_node) { 627 bpf_spin_unlock(&cgv_tree_lock); 628 stat_inc(FCG_STAT_PNC_NO_CGRP); 629 *cgidp = 0; 630 return true; 631 } 632 633 rb_node = bpf_rbtree_remove(&cgv_tree, rb_node); 634 bpf_spin_unlock(&cgv_tree_lock); 635 636 if (!rb_node) { 637 /* 638 * This should never happen. bpf_rbtree_first() was called 639 * above while the tree lock was held, so the node should 640 * always be present. 641 */ 642 scx_bpf_error("node could not be removed"); 643 return true; 644 } 645 646 cgv_node = container_of(rb_node, struct cgv_node, rb_node); 647 cgid = cgv_node->cgid; 648 649 if (time_before(cvtime_now, cgv_node->cvtime)) 650 cvtime_now = cgv_node->cvtime; 651 652 /* 653 * If lookup fails, the cgroup's gone. Free and move on. See 654 * fcg_cgroup_exit(). 655 */ 656 cgrp = bpf_cgroup_from_id(cgid); 657 if (!cgrp) { 658 stat_inc(FCG_STAT_PNC_GONE); 659 goto out_free; 660 } 661 662 cgc = bpf_cgrp_storage_get(&cgrp_ctx, cgrp, 0, 0); 663 if (!cgc) { 664 bpf_cgroup_release(cgrp); 665 stat_inc(FCG_STAT_PNC_GONE); 666 goto out_free; 667 } 668 669 if (!scx_bpf_dsq_move_to_local(cgid, 0)) { 670 bpf_cgroup_release(cgrp); 671 stat_inc(FCG_STAT_PNC_EMPTY); 672 goto out_stash; 673 } 674 675 /* 676 * Successfully consumed from the cgroup. This will be our current 677 * cgroup for the new slice. Refresh its hweight. 678 */ 679 cgrp_refresh_hweight(cgrp, cgc); 680 681 bpf_cgroup_release(cgrp); 682 683 /* 684 * As the cgroup may have more tasks, add it back to the rbtree. Note 685 * that here we charge the full slice upfront and then exact later 686 * according to the actual consumption. This prevents lowpri thundering 687 * herd from saturating the machine. 688 */ 689 bpf_spin_lock(&cgv_tree_lock); 690 cgv_node->cvtime += cgrp_slice_ns * FCG_HWEIGHT_ONE / (cgc->hweight ?: 1); 691 cgrp_cap_budget(cgv_node, cgc); 692 bpf_rbtree_add(&cgv_tree, &cgv_node->rb_node, cgv_node_less); 693 bpf_spin_unlock(&cgv_tree_lock); 694 695 *cgidp = cgid; 696 stat_inc(FCG_STAT_PNC_NEXT); 697 return true; 698 699 out_stash: 700 stash = bpf_map_lookup_elem(&cgv_node_stash, &cgid); 701 if (!stash) { 702 stat_inc(FCG_STAT_PNC_GONE); 703 goto out_free; 704 } 705 706 /* 707 * Paired with cmpxchg in cgrp_enqueued(). If they see the following 708 * transition, they'll enqueue the cgroup. If they are earlier, we'll 709 * see their task in the dq below and requeue the cgroup. 710 */ 711 __sync_val_compare_and_swap(&cgc->queued, 1, 0); 712 713 if (scx_bpf_dsq_nr_queued(cgid)) { 714 bpf_spin_lock(&cgv_tree_lock); 715 bpf_rbtree_add(&cgv_tree, &cgv_node->rb_node, cgv_node_less); 716 bpf_spin_unlock(&cgv_tree_lock); 717 stat_inc(FCG_STAT_PNC_RACE); 718 } else { 719 cgv_node = bpf_kptr_xchg(&stash->node, cgv_node); 720 if (cgv_node) { 721 scx_bpf_error("unexpected !NULL cgv_node stash"); 722 goto out_free; 723 } 724 } 725 726 return false; 727 728 out_free: 729 bpf_obj_drop(cgv_node); 730 return false; 731 } 732 733 void BPF_STRUCT_OPS(fcg_dispatch, s32 cpu, struct task_struct *prev) 734 { 735 struct fcg_cpu_ctx *cpuc; 736 struct fcg_cgrp_ctx *cgc; 737 struct cgroup *cgrp; 738 u64 now = scx_bpf_now(); 739 bool picked_next = false; 740 741 cpuc = find_cpu_ctx(); 742 if (!cpuc) 743 return; 744 745 if (!cpuc->cur_cgid) 746 goto pick_next_cgroup; 747 748 if (time_before(now, cpuc->cur_at + cgrp_slice_ns)) { 749 if (scx_bpf_dsq_move_to_local(cpuc->cur_cgid, 0)) { 750 stat_inc(FCG_STAT_CNS_KEEP); 751 return; 752 } 753 stat_inc(FCG_STAT_CNS_EMPTY); 754 } else { 755 stat_inc(FCG_STAT_CNS_EXPIRE); 756 } 757 758 /* 759 * The current cgroup is expiring. It was already charged a full slice. 760 * Calculate the actual usage and accumulate the delta. 761 */ 762 cgrp = bpf_cgroup_from_id(cpuc->cur_cgid); 763 if (!cgrp) { 764 stat_inc(FCG_STAT_CNS_GONE); 765 goto pick_next_cgroup; 766 } 767 768 cgc = bpf_cgrp_storage_get(&cgrp_ctx, cgrp, 0, 0); 769 if (cgc) { 770 /* 771 * We want to update the vtime delta and then look for the next 772 * cgroup to execute but the latter needs to be done in a loop 773 * and we can't keep the lock held. Oh well... 774 */ 775 s64 delta = now - cpuc->cur_at - cgrp_slice_ns; 776 777 bpf_spin_lock(&cgv_tree_lock); 778 /* keep the dividends positive, BPF division is unsigned */ 779 if (delta >= 0) 780 __sync_fetch_and_add(&cgc->cvtime_delta, 781 (u64)delta * FCG_HWEIGHT_ONE / 782 (cgc->hweight ?: 1)); 783 else 784 __sync_fetch_and_sub(&cgc->cvtime_delta, 785 (u64)-delta * FCG_HWEIGHT_ONE / 786 (cgc->hweight ?: 1)); 787 bpf_spin_unlock(&cgv_tree_lock); 788 } else { 789 stat_inc(FCG_STAT_CNS_GONE); 790 } 791 792 bpf_cgroup_release(cgrp); 793 794 pick_next_cgroup: 795 cpuc->cur_at = now; 796 797 if (scx_bpf_dsq_move_to_local(FALLBACK_DSQ, 0)) { 798 cpuc->cur_cgid = 0; 799 return; 800 } 801 802 bpf_repeat(CGROUP_MAX_RETRIES) { 803 if (try_pick_next_cgroup(&cpuc->cur_cgid)) { 804 picked_next = true; 805 break; 806 } 807 } 808 809 /* 810 * This only happens if try_pick_next_cgroup() races against enqueue 811 * path for more than CGROUP_MAX_RETRIES times, which is extremely 812 * unlikely and likely indicates an underlying bug. There shouldn't be 813 * any stall risk as the race is against enqueue. 814 */ 815 if (!picked_next) 816 stat_inc(FCG_STAT_PNC_FAIL); 817 } 818 819 s32 BPF_STRUCT_OPS(fcg_init_task, struct task_struct *p, 820 struct scx_init_task_args *args) 821 { 822 struct fcg_task_ctx *taskc; 823 struct fcg_cgrp_ctx *cgc; 824 825 /* 826 * @p is new. Let's ensure that its task_ctx is available. We can sleep 827 * in this function and the following will automatically use GFP_KERNEL. 828 */ 829 taskc = bpf_task_storage_get(&task_ctx, p, 0, 830 BPF_LOCAL_STORAGE_GET_F_CREATE); 831 if (!taskc) 832 return -ENOMEM; 833 834 taskc->bypassed_at = 0; 835 836 if (!(cgc = find_cgrp_ctx(args->cgroup))) 837 return -ENOENT; 838 839 scx_bpf_task_set_dsq_vtime(p, cgc->tvtime_now); 840 841 return 0; 842 } 843 844 int BPF_STRUCT_OPS_SLEEPABLE(fcg_cgroup_init, struct cgroup *cgrp, 845 struct scx_cgroup_init_args *args) 846 { 847 struct fcg_cgrp_ctx *cgc; 848 struct cgv_node *cgv_node; 849 struct cgv_node_stash empty_stash = {}, *stash; 850 u64 cgid = cgrp->kn->id; 851 int ret; 852 853 /* 854 * Technically incorrect as cgroup ID is full 64bit while dsq ID is 855 * 63bit. Should not be a problem in practice and easy to spot in the 856 * unlikely case that it breaks. 857 */ 858 ret = scx_bpf_create_dsq(cgid, -1); 859 if (ret) { 860 scx_bpf_error("scx_bpf_create_dsq failed (%d)", ret); 861 return ret; 862 } 863 864 cgc = bpf_cgrp_storage_get(&cgrp_ctx, cgrp, 0, 865 BPF_LOCAL_STORAGE_GET_F_CREATE); 866 if (!cgc) { 867 ret = -ENOMEM; 868 goto err_destroy_dsq; 869 } 870 871 cgc->weight = args->weight; 872 cgc->hweight = FCG_HWEIGHT_ONE; 873 874 ret = bpf_map_update_elem(&cgv_node_stash, &cgid, &empty_stash, 875 BPF_NOEXIST); 876 if (ret) { 877 if (ret != -ENOMEM) 878 scx_bpf_error("unexpected stash creation error (%d)", 879 ret); 880 goto err_destroy_dsq; 881 } 882 883 stash = bpf_map_lookup_elem(&cgv_node_stash, &cgid); 884 if (!stash) { 885 scx_bpf_error("unexpected cgv_node stash lookup failure"); 886 ret = -ENOENT; 887 goto err_destroy_dsq; 888 } 889 890 cgv_node = bpf_obj_new(struct cgv_node); 891 if (!cgv_node) { 892 ret = -ENOMEM; 893 goto err_del_cgv_node; 894 } 895 896 cgv_node->cgid = cgid; 897 cgv_node->cvtime = cvtime_now; 898 899 cgv_node = bpf_kptr_xchg(&stash->node, cgv_node); 900 if (cgv_node) { 901 scx_bpf_error("unexpected !NULL cgv_node stash"); 902 ret = -EBUSY; 903 goto err_drop; 904 } 905 906 return 0; 907 908 err_drop: 909 bpf_obj_drop(cgv_node); 910 err_del_cgv_node: 911 bpf_map_delete_elem(&cgv_node_stash, &cgid); 912 err_destroy_dsq: 913 scx_bpf_destroy_dsq(cgid); 914 return ret; 915 } 916 917 void BPF_STRUCT_OPS(fcg_cgroup_exit, struct cgroup *cgrp) 918 { 919 u64 cgid = cgrp->kn->id; 920 921 /* 922 * For now, there's no way find and remove the cgv_node if it's on the 923 * cgv_tree. Let's drain them in the dispatch path as they get popped 924 * off the front of the tree. 925 */ 926 bpf_map_delete_elem(&cgv_node_stash, &cgid); 927 scx_bpf_destroy_dsq(cgid); 928 } 929 930 void BPF_STRUCT_OPS(fcg_cgroup_move, struct task_struct *p, 931 struct cgroup *from, struct cgroup *to) 932 { 933 struct fcg_cgrp_ctx *from_cgc, *to_cgc; 934 s64 delta; 935 936 /* find_cgrp_ctx() triggers scx_bpf_error() on lookup failures */ 937 if (!(from_cgc = find_cgrp_ctx(from)) || !(to_cgc = find_cgrp_ctx(to))) 938 return; 939 940 delta = (s64)(p->scx.dsq_vtime - from_cgc->tvtime_now); 941 scx_bpf_task_set_dsq_vtime(p, to_cgc->tvtime_now + delta); 942 } 943 944 s32 BPF_STRUCT_OPS_SLEEPABLE(fcg_init) 945 { 946 int ret; 947 948 ret = scx_bpf_create_dsq(FALLBACK_DSQ, -1); 949 if (ret) { 950 scx_bpf_error("failed to create DSQ %d (%d)", FALLBACK_DSQ, ret); 951 return ret; 952 } 953 954 return 0; 955 } 956 957 void BPF_STRUCT_OPS(fcg_exit, struct scx_exit_info *ei) 958 { 959 UEI_RECORD(uei, ei); 960 } 961 962 SCX_OPS_DEFINE(flatcg_ops, 963 .select_cpu = (void *)fcg_select_cpu, 964 .enqueue = (void *)fcg_enqueue, 965 .dispatch = (void *)fcg_dispatch, 966 .runnable = (void *)fcg_runnable, 967 .running = (void *)fcg_running, 968 .stopping = (void *)fcg_stopping, 969 .quiescent = (void *)fcg_quiescent, 970 .init_task = (void *)fcg_init_task, 971 .cgroup_set_weight = (void *)fcg_cgroup_set_weight, 972 .cgroup_init = (void *)fcg_cgroup_init, 973 .cgroup_exit = (void *)fcg_cgroup_exit, 974 .cgroup_move = (void *)fcg_cgroup_move, 975 .init = (void *)fcg_init, 976 .exit = (void *)fcg_exit, 977 .flags = SCX_OPS_ENQ_EXITING, 978 .name = "flatcg"); 979