enq_select_cpu.bpf.c

 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
10char _license[] SEC("license") = "GPL";
11
12UEI_DEFINE(uei);
13
14s32 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
21void 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
36void BPF_STRUCT_OPS(enq_select_cpu_exit, struct scx_exit_info *ei)
37{
38	UEI_RECORD(uei, ei);
39}
40
41struct task_cpu_arg {
42	pid_t pid;
43};
44
45SEC("syscall")
46int 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
67SEC(".struct_ops.link")
68struct 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};