1 // SPDX-License-Identifier: GPL-2.0
2 /*
3 * A scheduler that validates the behavior of scx_bpf_select_cpu_and() by
4 * selecting idle CPUs strictly within a subset of allowed CPUs.
5 *
6 * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
7 */
8
9 #include <scx/common.bpf.h>
10
11 char _license[] SEC("license") = "GPL";
12
13 UEI_DEFINE(uei);
14
15 private(PREF_CPUS) struct bpf_cpumask __kptr * allowed_cpumask;
16
17 static void
validate_local_idle_state(void)18 validate_local_idle_state(void)
19 {
20 const struct cpumask *idle;
21 struct task_struct *curr;
22 s32 cpu = bpf_get_smp_processor_id();
23 bool cpu_is_idle, curr_is_idle;
24
25 bpf_rcu_read_lock();
26 curr = scx_bpf_cpu_curr(cpu);
27 curr_is_idle = curr && (curr->flags & PF_IDLE);
28 bpf_rcu_read_unlock();
29
30 idle = scx_bpf_get_idle_cpumask();
31 cpu_is_idle = bpf_cpumask_test_cpu(cpu, idle);
32 scx_bpf_put_idle_cpumask(idle);
33
34 /*
35 * Unlike a remote selected CPU, the local CPU cannot go through an
36 * idle re-pick while this callback is running. If it is running a
37 * non-idle scheduling context, it must not be advertised as idle.
38 */
39 if (!curr_is_idle && cpu_is_idle)
40 scx_bpf_error("running CPU %d should be marked as busy", cpu);
41 }
42
43 static void
validate_selected_cpu(const struct task_struct * p,s32 cpu)44 validate_selected_cpu(const struct task_struct *p, s32 cpu)
45 {
46 const struct cpumask *allowed = cast_mask(allowed_cpumask);
47
48 if (!allowed) {
49 scx_bpf_error("allowed domain not initialized");
50 return;
51 }
52
53 if (!bpf_cpumask_test_cpu(cpu, allowed))
54 scx_bpf_error("CPU %d not in the allowed domain for %d (%s)",
55 cpu, p->pid, p->comm);
56
57 if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr))
58 scx_bpf_error("CPU %d not in the affinity mask for %d (%s)",
59 cpu, p->pid, p->comm);
60 }
61
BPF_STRUCT_OPS(allowed_cpus_select_cpu,struct task_struct * p,s32 prev_cpu,u64 wake_flags)62 s32 BPF_STRUCT_OPS(allowed_cpus_select_cpu,
63 struct task_struct *p, s32 prev_cpu, u64 wake_flags)
64 {
65 const struct cpumask *allowed;
66 s32 cpu;
67
68 validate_local_idle_state();
69 allowed = cast_mask(allowed_cpumask);
70 if (!allowed) {
71 scx_bpf_error("allowed domain not initialized");
72 return -EINVAL;
73 }
74
75 /*
76 * Select an idle CPU strictly within the allowed domain.
77 */
78 cpu = scx_bpf_select_cpu_and(p, prev_cpu, wake_flags, allowed, 0);
79 if (cpu >= 0) {
80 validate_selected_cpu(p, cpu);
81 scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
82
83 return cpu;
84 }
85
86 return prev_cpu;
87 }
88
BPF_STRUCT_OPS(allowed_cpus_enqueue,struct task_struct * p,u64 enq_flags)89 void BPF_STRUCT_OPS(allowed_cpus_enqueue, struct task_struct *p, u64 enq_flags)
90 {
91 const struct cpumask *allowed;
92 s32 prev_cpu = scx_bpf_task_cpu(p), cpu;
93
94 scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
95
96 validate_local_idle_state();
97 allowed = cast_mask(allowed_cpumask);
98 if (!allowed) {
99 scx_bpf_error("allowed domain not initialized");
100 return;
101 }
102
103 /*
104 * Use scx_bpf_select_cpu_and() to proactively kick an idle CPU
105 * within @allowed_cpumask, usable by @p.
106 */
107 cpu = scx_bpf_select_cpu_and(p, prev_cpu, 0, allowed, 0);
108 if (cpu >= 0) {
109 validate_selected_cpu(p, cpu);
110 scx_bpf_kick_cpu(cpu, SCX_KICK_IDLE);
111 }
112 }
113
BPF_STRUCT_OPS_SLEEPABLE(allowed_cpus_init)114 s32 BPF_STRUCT_OPS_SLEEPABLE(allowed_cpus_init)
115 {
116 struct bpf_cpumask *mask;
117
118 mask = bpf_cpumask_create();
119 if (!mask)
120 return -ENOMEM;
121
122 mask = bpf_kptr_xchg(&allowed_cpumask, mask);
123 if (mask)
124 bpf_cpumask_release(mask);
125
126 bpf_rcu_read_lock();
127
128 /*
129 * Assign the first online CPU to the allowed domain.
130 */
131 mask = allowed_cpumask;
132 if (mask) {
133 const struct cpumask *online = scx_bpf_get_online_cpumask();
134
135 bpf_cpumask_set_cpu(bpf_cpumask_first(online), mask);
136 scx_bpf_put_cpumask(online);
137 }
138
139 bpf_rcu_read_unlock();
140
141 return 0;
142 }
143
BPF_STRUCT_OPS(allowed_cpus_exit,struct scx_exit_info * ei)144 void BPF_STRUCT_OPS(allowed_cpus_exit, struct scx_exit_info *ei)
145 {
146 UEI_RECORD(uei, ei);
147 }
148
149 struct task_cpu_arg {
150 pid_t pid;
151 };
152
153 SEC("syscall")
select_cpu_from_user(struct task_cpu_arg * input)154 int select_cpu_from_user(struct task_cpu_arg *input)
155 {
156 struct task_struct *p;
157 int cpu;
158
159 p = bpf_task_from_pid(input->pid);
160 if (!p)
161 return -EINVAL;
162
163 bpf_rcu_read_lock();
164 cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0, p->cpus_ptr, 0);
165 bpf_rcu_read_unlock();
166
167 bpf_task_release(p);
168
169 return cpu;
170 }
171
172 SEC(".struct_ops.link")
173 struct sched_ext_ops allowed_cpus_ops = {
174 .select_cpu = (void *)allowed_cpus_select_cpu,
175 .enqueue = (void *)allowed_cpus_enqueue,
176 .init = (void *)allowed_cpus_init,
177 .exit = (void *)allowed_cpus_exit,
178 .name = "allowed_cpus",
179 };
180