1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * Copyright (c) 2023 Meta Platforms, Inc. and affiliates.
4 * Copyright (c) 2023 David Vernet <dvernet@meta.com>
5 * Copyright (c) 2023 Tejun Heo <tj@kernel.org>
6 */
7
8 #include <scx/common.bpf.h>
9
10 char _license[] SEC("license") = "GPL";
11
12 UEI_DEFINE(uei);
13
BPF_STRUCT_OPS(enq_select_cpu_select_cpu,struct task_struct * p,s32 prev_cpu,u64 wake_flags)14 s32 BPF_STRUCT_OPS(enq_select_cpu_select_cpu, struct task_struct *p,
15 s32 prev_cpu, u64 wake_flags)
16 {
17 /* Bounce all tasks to ops.enqueue() */
18 return prev_cpu;
19 }
20
BPF_STRUCT_OPS(enq_select_cpu_enqueue,struct task_struct * p,u64 enq_flags)21 void BPF_STRUCT_OPS(enq_select_cpu_enqueue, struct task_struct *p,
22 u64 enq_flags)
23 {
24 s32 cpu, prev_cpu = scx_bpf_task_cpu(p);
25 bool found = false;
26
27 cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, 0, &found);
28 if (found) {
29 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cpu, SCX_SLICE_DFL, enq_flags);
30 return;
31 }
32
33 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags);
34 }
35
BPF_STRUCT_OPS(enq_select_cpu_exit,struct scx_exit_info * ei)36 void BPF_STRUCT_OPS(enq_select_cpu_exit, struct scx_exit_info *ei)
37 {
38 UEI_RECORD(uei, ei);
39 }
40
41 struct task_cpu_arg {
42 pid_t pid;
43 };
44
45 SEC("syscall")
select_cpu_from_user(struct task_cpu_arg * input)46 int select_cpu_from_user(struct task_cpu_arg *input)
47 {
48 struct task_struct *p;
49 bool found = false;
50 s32 cpu;
51
52 p = bpf_task_from_pid(input->pid);
53 if (!p)
54 return -EINVAL;
55
56 bpf_rcu_read_lock();
57 cpu = scx_bpf_select_cpu_dfl(p, bpf_get_smp_processor_id(), 0, &found);
58 if (!found)
59 cpu = -EBUSY;
60 bpf_rcu_read_unlock();
61
62 bpf_task_release(p);
63
64 return cpu;
65 }
66
67 SEC(".struct_ops.link")
68 struct sched_ext_ops enq_select_cpu_ops = {
69 .select_cpu = (void *)enq_select_cpu_select_cpu,
70 .enqueue = (void *)enq_select_cpu_enqueue,
71 .exit = (void *)enq_select_cpu_exit,
72 .name = "enq_select_cpu",
73 .timeout_ms = 1000U,
74 };
75