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

   1 /* SPDX-License-Identifier: GPL-2.0 */
   2 /*
   3  * Copyright (c) 2024 Meta Platforms, Inc. and affiliates.
   4  * Copyright (c) 2024 David Vernet <dvernet@meta.com>
   5  */
   6 #include <scx/common.bpf.h>
   7
   8 char _license[] SEC("license") = "GPL";
   9 const volatile s32 nr_cpus;
  10
  11 UEI_DEFINE(uei);
  12
  13 struct {
  14         __uint(type, BPF_MAP_TYPE_QUEUE);
  15         __uint(max_entries, 8192);
  16         __type(value, s32);
  17 } queue SEC(".maps");
  18
  19 s32 BPF_STRUCT_OPS(dsp_local_on_select_cpu, struct task_struct *p,
  20                    s32 prev_cpu, u64 wake_flags)
  21 {
  22         return prev_cpu;
  23 }
  24
  25 void BPF_STRUCT_OPS(dsp_local_on_enqueue, struct task_struct *p,
  26                     u64 enq_flags)
  27 {
  28         s32 pid = p->pid;
  29
  30         if (bpf_map_push_elem(&queue, &pid, 0))
  31                 scx_bpf_error("Failed to enqueue %s[%d]", p->comm, p->pid);
  32 }
  33
  34 void BPF_STRUCT_OPS(dsp_local_on_dispatch, s32 cpu, struct task_struct *prev)
  35 {
  36         s32 pid, target;
  37         struct task_struct *p;
  38
  39         if (bpf_map_pop_elem(&queue, &pid))
  40                 return;
  41
  42         p = bpf_task_from_pid(pid);
  43         if (!p)
  44                 return;
  45
  46         target = bpf_get_prandom_u32() % nr_cpus;
  47
  48         scx_bpf_dispatch(p, SCX_DSQ_LOCAL_ON | target, SCX_SLICE_DFL, 0);
  49         bpf_task_release(p);
  50 }
  51
  52 void BPF_STRUCT_OPS(dsp_local_on_exit, struct scx_exit_info *ei)
  53 {
  54         UEI_RECORD(uei, ei);
  55 }
  56
  57 SEC(".struct_ops.link")
  58 struct sched_ext_ops dsp_local_on_ops = {
  59         .select_cpu             = (void *) dsp_local_on_select_cpu,
  60         .enqueue                = (void *) dsp_local_on_enqueue,
  61         .dispatch               = (void *) dsp_local_on_dispatch,
  62         .exit                   = (void *) dsp_local_on_exit,
  63         .name                   = "dsp_local_on",
  64         .timeout_ms             = 1000U,
  65 };