tools/testing/selftests/sched_ext/select_cpu_dfl_nodispatch.bpf.c

   1 /* SPDX-License-Identifier: GPL-2.0 */
   2 /*
   3  * A scheduler that validates the behavior of direct dispatching with a default
   4  * select_cpu implementation, and with the SCX_OPS_ENQ_DFL_NO_DISPATCH ops flag
   5  * specified.
   6  *
   7  * Copyright (c) 2023 Meta Platforms, Inc. and affiliates.
   8  * Copyright (c) 2023 David Vernet <dvernet@meta.com>
   9  * Copyright (c) 2023 Tejun Heo <tj@kernel.org>
  10  */
  11
  12 #include <scx/common.bpf.h>
  13
  14 char _license[] SEC("license") = "GPL";
  15
  16 bool saw_local = false;
  17
  18 /* Per-task scheduling context */
  19 struct task_ctx {
  20         bool    force_local;    /* CPU changed by ops.select_cpu() */
  21 };
  22
  23 struct {
  24         __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
  25         __uint(map_flags, BPF_F_NO_PREALLOC);
  26         __type(key, int);
  27         __type(value, struct task_ctx);
  28 } task_ctx_stor SEC(".maps");
  29
  30 /* Manually specify the signature until the kfunc is added to the scx repo. */
  31 s32 scx_bpf_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
  32                            bool *found) __ksym;
  33
  34 s32 BPF_STRUCT_OPS(select_cpu_dfl_nodispatch_select_cpu, struct task_struct *p,
  35                    s32 prev_cpu, u64 wake_flags)
  36 {
  37         struct task_ctx *tctx;
  38         s32 cpu;
  39
  40         tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
  41         if (!tctx) {
  42                 scx_bpf_error("task_ctx lookup failed");
  43                 return -ESRCH;
  44         }
  45
  46         cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, wake_flags,
  47                                      &tctx->force_local);
  48
  49         return cpu;
  50 }
  51
  52 void BPF_STRUCT_OPS(select_cpu_dfl_nodispatch_enqueue, struct task_struct *p,
  53                     u64 enq_flags)
  54 {
  55         u64 dsq_id = SCX_DSQ_GLOBAL;
  56         struct task_ctx *tctx;
  57
  58         tctx = bpf_task_storage_get(&task_ctx_stor, p, 0, 0);
  59         if (!tctx) {
  60                 scx_bpf_error("task_ctx lookup failed");
  61                 return;
  62         }
  63
  64         if (tctx->force_local) {
  65                 dsq_id = SCX_DSQ_LOCAL;
  66                 tctx->force_local = false;
  67                 saw_local = true;
  68         }
  69
  70         scx_bpf_dispatch(p, dsq_id, SCX_SLICE_DFL, enq_flags);
  71 }
  72
  73 s32 BPF_STRUCT_OPS(select_cpu_dfl_nodispatch_init_task,
  74                    struct task_struct *p, struct scx_init_task_args *args)
  75 {
  76         if (bpf_task_storage_get(&task_ctx_stor, p, 0,
  77                                  BPF_LOCAL_STORAGE_GET_F_CREATE))
  78                 return 0;
  79         else
  80                 return -ENOMEM;
  81 }
  82
  83 SEC(".struct_ops.link")
  84 struct sched_ext_ops select_cpu_dfl_nodispatch_ops = {
  85         .select_cpu             = (void *) select_cpu_dfl_nodispatch_select_cpu,
  86         .enqueue                = (void *) select_cpu_dfl_nodispatch_enqueue,
  87         .init_task              = (void *) select_cpu_dfl_nodispatch_init_task,
  88         .name                   = "select_cpu_dfl_nodispatch",
  89 };