1use anyhow::{anyhow, bail, Context, Result};
7use libbpf_rs::libbpf_sys::*;
8use libbpf_rs::{AsRawLibbpf, OpenProgramImpl};
9use log::warn;
10use std::env;
11use std::ffi::c_void;
12use std::ffi::CStr;
13use std::ffi::CString;
14use std::io;
15use std::io::BufRead;
16use std::io::BufReader;
17use std::mem::size_of;
18use std::slice::from_raw_parts;
19
20const PROCFS_MOUNTS: &str = "/proc/mounts";
21const TRACEFS: &str = "tracefs";
22const DEBUGFS: &str = "debugfs";
23
24lazy_static::lazy_static! {
25 pub static ref SCX_OPS_KEEP_BUILTIN_IDLE: u64 =
26 read_enum("scx_ops_flags", "SCX_OPS_KEEP_BUILTIN_IDLE").unwrap_or(0);
27 pub static ref SCX_OPS_ENQ_LAST: u64 =
28 read_enum("scx_ops_flags", "SCX_OPS_ENQ_LAST").unwrap_or(0);
29 pub static ref SCX_OPS_ENQ_EXITING: u64 =
30 read_enum("scx_ops_flags", "SCX_OPS_ENQ_EXITING").unwrap_or(0);
31 pub static ref SCX_OPS_SWITCH_PARTIAL: u64 =
32 read_enum("scx_ops_flags", "SCX_OPS_SWITCH_PARTIAL").unwrap_or(0);
33 pub static ref SCX_OPS_ENQ_MIGRATION_DISABLED: u64 =
34 read_enum("scx_ops_flags", "SCX_OPS_ENQ_MIGRATION_DISABLED").unwrap_or(0);
35 pub static ref SCX_OPS_ALLOW_QUEUED_WAKEUP: u64 =
36 read_enum("scx_ops_flags", "SCX_OPS_ALLOW_QUEUED_WAKEUP").unwrap_or(0);
37 pub static ref SCX_OPS_BUILTIN_IDLE_PER_NODE: u64 =
38 read_enum("scx_ops_flags", "SCX_OPS_BUILTIN_IDLE_PER_NODE").unwrap_or(0);
39
40 pub static ref SCX_PICK_IDLE_CORE: u64 =
41 read_enum("scx_pick_idle_cpu_flags", "SCX_PICK_IDLE_CORE").unwrap_or(0);
42 pub static ref SCX_PICK_IDLE_IN_NODE: u64 =
43 read_enum("scx_pick_idle_cpu_flags", "SCX_PICK_IDLE_IN_NODE").unwrap_or(0);
44
45 pub static ref ROOT_PREFIX: String =
46 env::var("SCX_SYSFS_PREFIX").unwrap_or("".to_string());
47}
48
49fn load_vmlinux_btf() -> &'static mut btf {
50 let btf = unsafe { btf__load_vmlinux_btf() };
51 if btf.is_null() {
52 panic!("btf__load_vmlinux_btf() returned NULL, was CONFIG_DEBUG_INFO_BTF enabled?")
53 }
54 unsafe { &mut *btf }
55}
56
57lazy_static::lazy_static! {
58 static ref VMLINUX_BTF: &'static mut btf = load_vmlinux_btf();
59}
60
61fn btf_kind(t: &btf_type) -> u32 {
62 (t.info >> 24) & 0x1f
63}
64
65fn btf_vlen(t: &btf_type) -> u32 {
66 t.info & 0xffff
67}
68
69fn btf_type_plus_1(t: &btf_type) -> *const c_void {
70 let ptr_val = t as *const btf_type as usize;
71 (ptr_val + size_of::<btf_type>()) as *const c_void
72}
73
74fn btf_enum(t: &btf_type) -> &[btf_enum] {
75 let ptr = btf_type_plus_1(t);
76 unsafe { from_raw_parts(ptr as *const btf_enum, btf_vlen(t) as usize) }
77}
78
79fn btf_enum64(t: &btf_type) -> &[btf_enum64] {
80 let ptr = btf_type_plus_1(t);
81 unsafe { from_raw_parts(ptr as *const btf_enum64, btf_vlen(t) as usize) }
82}
83
84fn btf_members(t: &btf_type) -> &[btf_member] {
85 let ptr = btf_type_plus_1(t);
86 unsafe { from_raw_parts(ptr as *const btf_member, btf_vlen(t) as usize) }
87}
88
89fn btf_name_str_by_offset(btf: &btf, name_off: u32) -> Result<&str> {
90 let n = unsafe { btf__name_by_offset(btf, name_off) };
91 if n.is_null() {
92 bail!("btf__name_by_offset() returned NULL");
93 }
94 Ok(unsafe { CStr::from_ptr(n) }
95 .to_str()
96 .with_context(|| format!("Failed to convert {:?} to string", n))?)
97}
98
99pub fn read_enum(type_name: &str, name: &str) -> Result<u64> {
100 let btf: &btf = *VMLINUX_BTF;
101
102 let type_name = CString::new(type_name).unwrap();
103 let tid = unsafe { btf__find_by_name(btf, type_name.as_ptr()) };
104 if tid < 0 {
105 bail!("type {:?} doesn't exist, ret={}", type_name, tid);
106 }
107
108 let t = unsafe { btf__type_by_id(btf, tid as _) };
109 if t.is_null() {
110 bail!("btf__type_by_id({}) returned NULL", tid);
111 }
112 let t = unsafe { &*t };
113
114 match btf_kind(t) {
115 BTF_KIND_ENUM => {
116 for e in btf_enum(t).iter() {
117 if btf_name_str_by_offset(btf, e.name_off)? == name {
118 return Ok(e.val as u64);
119 }
120 }
121 }
122 BTF_KIND_ENUM64 => {
123 for e in btf_enum64(t).iter() {
124 if btf_name_str_by_offset(btf, e.name_off)? == name {
125 return Ok(((e.val_hi32 as u64) << 32) | (e.val_lo32) as u64);
126 }
127 }
128 }
129 _ => (),
130 }
131
132 Err(anyhow!("{:?} doesn't exist in {:?}", name, type_name))
133}
134
135pub fn struct_has_field(type_name: &str, field: &str) -> Result<bool> {
136 let btf: &btf = *VMLINUX_BTF;
137
138 let type_name = CString::new(type_name).unwrap();
139 let tid = unsafe { btf__find_by_name_kind(btf, type_name.as_ptr(), BTF_KIND_STRUCT) };
140 if tid < 0 {
141 bail!("type {:?} doesn't exist, ret={}", type_name, tid);
142 }
143
144 let t = unsafe { btf__type_by_id(btf, tid as _) };
145 if t.is_null() {
146 bail!("btf__type_by_id({}) returned NULL", tid);
147 }
148 let t = unsafe { &*t };
149
150 for m in btf_members(t).iter() {
151 if btf_name_str_by_offset(btf, m.name_off)? == field {
152 return Ok(true);
153 }
154 }
155
156 Ok(false)
157}
158
159pub fn ksym_exists(ksym: &str) -> Result<bool> {
160 let btf: &btf = *VMLINUX_BTF;
161
162 let ksym_name = CString::new(ksym).unwrap();
163 let tid = unsafe { btf__find_by_name(btf, ksym_name.as_ptr()) };
164 Ok(tid >= 0)
165}
166
167pub fn in_kallsyms(ksym: &str) -> Result<bool> {
168 let file = std::fs::File::open("/proc/kallsyms")?;
169 let reader = std::io::BufReader::new(file);
170
171 for line in reader.lines() {
172 for sym in line.unwrap().split_whitespace() {
173 if ksym == sym {
174 return Ok(true);
175 }
176 }
177 }
178
179 Ok(false)
180}
181
182pub fn get_fs_mount(mount_type: &str) -> Result<Vec<std::path::PathBuf>> {
184 let proc_mounts_path = std::path::Path::new(PROCFS_MOUNTS);
185
186 let file = std::fs::File::open(proc_mounts_path)
187 .with_context(|| format!("Failed to open {}", proc_mounts_path.display()))?;
188
189 let reader = BufReader::new(file);
190
191 let mut mounts = Vec::new();
192 for line in reader.lines() {
193 let line = line.context("Failed to read line from /proc/mounts")?;
194 let mount_info: Vec<&str> = line.split_whitespace().collect();
195
196 if mount_info.len() > 3 && mount_info[2] == mount_type {
197 let mount_path = std::path::PathBuf::from(mount_info[1]);
198 mounts.push(mount_path);
199 }
200 }
201
202 Ok(mounts)
203}
204
205pub fn tracefs_mount() -> Result<std::path::PathBuf> {
207 let mounts = get_fs_mount(TRACEFS)?;
208 mounts.into_iter().next().context("No tracefs mount found")
209}
210
211pub fn debugfs_mount() -> Result<std::path::PathBuf> {
213 let mounts = get_fs_mount(DEBUGFS)?;
214 mounts.into_iter().next().context("No debugfs mount found")
215}
216
217pub fn tracer_available(tracer: &str) -> Result<bool> {
218 let base_path = tracefs_mount().unwrap_or_else(|_| debugfs_mount().unwrap().join("tracing"));
219 let file = match std::fs::File::open(base_path.join("available_tracers")) {
220 Ok(f) => f,
221 Err(_) => return Ok(false),
222 };
223 let reader = std::io::BufReader::new(file);
224
225 for line in reader.lines() {
226 for tc in line.unwrap().split_whitespace() {
227 if tracer == tc {
228 return Ok(true);
229 }
230 }
231 }
232
233 Ok(false)
234}
235
236pub fn tracepoint_exists(tracepoint: &str) -> Result<bool> {
237 let base_path = tracefs_mount().unwrap_or_else(|_| debugfs_mount().unwrap().join("tracing"));
238 let file = match std::fs::File::open(base_path.join("available_events")) {
239 Ok(f) => f,
240 Err(_) => return Ok(false),
241 };
242 let reader = std::io::BufReader::new(file);
243
244 for line in reader.lines() {
245 for tp in line.unwrap().split_whitespace() {
246 if tracepoint == tp {
247 return Ok(true);
248 }
249 }
250 }
251
252 Ok(false)
253}
254
255pub fn cond_kprobe_enable<T>(sym: &str, prog_ptr: &OpenProgramImpl<T>) -> Result<bool> {
256 if in_kallsyms(sym)? {
257 unsafe {
258 bpf_program__set_autoload(prog_ptr.as_libbpf_object().as_ptr(), true);
259 }
260 return Ok(true);
261 } else {
262 warn!("symbol {sym} is missing, kprobe not loaded");
263 }
264
265 Ok(false)
266}
267
268pub fn cond_kprobes_enable<T>(kprobes: Vec<(&str, &OpenProgramImpl<T>)>) -> Result<bool> {
269 for (sym, _) in kprobes.iter() {
271 if in_kallsyms(sym)? == false {
272 warn!("symbol {sym} is missing, kprobe not loaded");
273 return Ok(false);
274 }
275 }
276
277 for (_, ptr) in kprobes.iter() {
279 unsafe {
280 bpf_program__set_autoload(ptr.as_libbpf_object().as_ptr(), true);
281 }
282 }
283
284 Ok(true)
285}
286
287pub fn cond_tracepoint_enable<T>(tracepoint: &str, prog_ptr: &OpenProgramImpl<T>) -> Result<bool> {
288 if tracepoint_exists(tracepoint)? {
289 unsafe {
290 bpf_program__set_autoload(prog_ptr.as_libbpf_object().as_ptr(), true);
291 }
292 return Ok(true);
293 } else {
294 warn!("tradepoint {tracepoint} is missing, tracepoint not loaded");
295 }
296
297 Ok(false)
298}
299
300pub fn cond_tracepoints_enable<T>(tracepoints: Vec<(&str, &OpenProgramImpl<T>)>) -> Result<bool> {
301 for (tp, _) in tracepoints.iter() {
303 if tracepoint_exists(tp)? == false {
304 warn!("tradepoint {tp} is missing, tracepoint not loaded");
305 return Ok(false);
306 }
307 }
308
309 for (_, ptr) in tracepoints.iter() {
311 unsafe {
312 bpf_program__set_autoload(ptr.as_libbpf_object().as_ptr(), true);
313 }
314 }
315
316 Ok(true)
317}
318
319pub fn is_sched_ext_enabled() -> io::Result<bool> {
320 let content = std::fs::read_to_string("/sys/kernel/sched_ext/state")?;
321
322 match content.trim() {
323 "enabled" => Ok(true),
324 "disabled" => Ok(false),
325 _ => {
326 Err(io::Error::new(
328 io::ErrorKind::InvalidData,
329 "Unexpected content in /sys/kernel/sched_ext/state",
330 ))
331 }
332 }
333}
334
335#[macro_export]
336macro_rules! unwrap_or_break {
337 ($expr: expr, $label: lifetime) => {{
338 match $expr {
339 Ok(val) => val,
340 Err(e) => break $label Err(e),
341 }
342 }};
343}
344
345pub fn check_min_requirements() -> Result<()> {
346 if let Ok(false) | Err(_) = struct_has_field("sched_ext_ops", "dump") {
349 bail!("sched_ext_ops.dump() missing, kernel too old?");
350 }
351 Ok(())
352}
353
354#[rustfmt::skip]
359#[macro_export]
360macro_rules! scx_ops_open {
361 ($builder: expr, $obj_ref: expr, $ops: ident, $open_opts: expr) => { 'block: {
362 scx_utils::paste! {
363 scx_utils::unwrap_or_break!(scx_utils::compat::check_min_requirements(), 'block);
364 use ::anyhow::Context;
365 use ::libbpf_rs::skel::SkelBuilder;
366
367 let mut skel = match $open_opts {
368 Some(opts_ref) => { match $builder.open_opts(opts_ref, $obj_ref).context("Failed to open BPF program with options") {
370 Ok(val) => val,
371 Err(e) => break 'block Err(e),
372 }
373 }
374 None => {
375 match $builder.open($obj_ref).context("Failed to open BPF program") {
376 Ok(val) => val,
377 Err(e) => break 'block Err(e),
378 }
379 }
380 };
381
382 let ops = skel.struct_ops.[<$ops _mut>]();
383 let path = std::path::Path::new("/sys/kernel/sched_ext/hotplug_seq");
384
385 let val = match std::fs::read_to_string(&path) {
386 Ok(val) => val,
387 Err(_) => {
388 break 'block Err(anyhow::anyhow!("Failed to open or read file {:?}", path));
389 }
390 };
391
392 ops.hotplug_seq = match val.trim().parse::<u64>() {
393 Ok(parsed) => parsed,
394 Err(_) => {
395 break 'block Err(anyhow::anyhow!("Failed to parse hotplug seq {}", val));
396 }
397 };
398
399 if let Ok(s) = ::std::env::var("SCX_TIMEOUT_MS") {
400 skel.struct_ops.[<$ops _mut>]().timeout_ms = match s.parse::<u32>() {
401 Ok(ms) => {
402 ::scx_utils::info!("Setting timeout_ms to {} based on environment", ms);
403 ms
404 },
405 Err(e) => {
406 break 'block anyhow::Result::Err(e).context("SCX_TIMEOUT_MS has invalid value");
407 },
408 };
409 }
410
411 {
412 let ops = skel.struct_ops.[<$ops _mut>]();
413
414 let name_field = &mut ops.name;
415
416 let version_suffix = ::scx_utils::build_id::ops_version_suffix(env!("CARGO_PKG_VERSION"));
417 let bytes = version_suffix.as_bytes();
418 let mut i = 0;
419 let mut bytes_idx = 0;
420 let mut found_null = false;
421
422 while i < name_field.len() - 1 {
423 found_null |= name_field[i] == 0;
424 if !found_null {
425 i += 1;
426 continue;
427 }
428
429 if bytes_idx < bytes.len() {
430 name_field[i] = bytes[bytes_idx] as i8;
431 bytes_idx += 1;
432 } else {
433 break;
434 }
435 i += 1;
436 }
437 name_field[i] = 0;
438 }
439
440 $crate::import_enums!(skel);
441
442 let result = ::anyhow::Result::Ok(skel);
443
444 result
445 }
446 }};
447}
448
449#[rustfmt::skip]
454#[macro_export]
455macro_rules! scx_ops_load {
456 ($skel: expr, $ops: ident, $uei: ident) => { 'block: {
457 scx_utils::paste! {
458 use ::anyhow::Context;
459 use ::libbpf_rs::skel::OpenSkel;
460
461 scx_utils::uei_set_size!($skel, $ops, $uei);
462 $skel.load().context("Failed to load BPF program")
463 }
464 }};
465}
466
467#[rustfmt::skip]
469#[macro_export]
470macro_rules! scx_ops_attach {
471 ($skel: expr, $ops: ident) => { 'block: {
472 use ::anyhow::Context;
473 use ::libbpf_rs::skel::Skel;
474
475 if scx_utils::compat::is_sched_ext_enabled().unwrap_or(false) {
476 break 'block Err(anyhow::anyhow!(
477 "another sched_ext scheduler is already running"
478 ));
479 }
480 $skel
481 .attach()
482 .context("Failed to attach non-struct_ops BPF programs")
483 .and_then(|_| {
484 $skel
485 .maps
486 .$ops
487 .attach_struct_ops()
488 .context("Failed to attach struct_ops BPF programs")
489 })
490 }};
491}
492
493#[cfg(test)]
494mod tests {
495 #[test]
496 fn test_read_enum() {
497 assert_eq!(super::read_enum("pid_type", "PIDTYPE_TGID").unwrap(), 1);
498 }
499
500 #[test]
501 fn test_struct_has_field() {
502 assert!(super::struct_has_field("task_struct", "flags").unwrap());
503 assert!(!super::struct_has_field("task_struct", "NO_SUCH_FIELD").unwrap());
504 assert!(super::struct_has_field("NO_SUCH_STRUCT", "NO_SUCH_FIELD").is_err());
505 }
506
507 #[test]
508 fn test_ksym_exists() {
509 assert!(super::ksym_exists("bpf_task_acquire").unwrap());
510 assert!(!super::ksym_exists("NO_SUCH_KFUNC").unwrap());
511 }
512}