scx_lib_selftests/
main.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5mod bpf_skel;
6pub use bpf_skel::*;
7
8use std::mem::MaybeUninit;
9
10use anyhow::bail;
11use anyhow::Context;
12use anyhow::Result;
13
14use std::ffi::c_ulong;
15use std::sync::Arc;
16
17use scx_utils::init_libbpf_logging;
18use scx_utils::Core;
19use scx_utils::Llc;
20use scx_utils::Topology;
21use scx_utils::NR_CPU_IDS;
22
23use simplelog::{ColorChoice, Config as SimplelogConfig, TermLogger, TerminalMode};
24
25use libbpf_rs::skel::OpenSkel;
26use libbpf_rs::skel::SkelBuilder;
27use libbpf_rs::PrintLevel;
28use libbpf_rs::ProgramInput;
29
30fn setup_arenas(skel: &mut BpfSkel<'_>) -> Result<()> {
31    const STATIC_ALLOC_PAGES_GRANULARITY: c_ulong = 512;
32    const TASK_SIZE: c_ulong = 42;
33
34    // Allocate the arena memory from the BPF side so userspace initializes it before starting
35    // the scheduler. Despite the function call's name this is neither a test nor a test run,
36    // it's the recommended way of executing SEC("syscall") probes.
37    let mut args = types::arena_init_args {
38        static_pages: STATIC_ALLOC_PAGES_GRANULARITY,
39        task_ctx_size: TASK_SIZE,
40    };
41
42    let input = ProgramInput {
43        context_in: Some(unsafe {
44            std::slice::from_raw_parts_mut(
45                &mut args as *mut _ as *mut u8,
46                std::mem::size_of_val(&args),
47            )
48        }),
49        ..Default::default()
50    };
51
52    let output = skel.progs.arena_init.test_run(input)?;
53    if output.return_value != 0 {
54        bail!(
55            "Could not initialize arenas, arena_init returned {}",
56            output.return_value as i32
57        );
58    }
59
60    Ok(())
61}
62
63fn setup_topology_node(skel: &mut BpfSkel<'_>, mask: &[u64]) -> Result<()> {
64    let mut args = types::arena_alloc_mask_args {
65        bitmap: 0 as c_ulong,
66    };
67
68    let input = ProgramInput {
69        context_in: Some(unsafe {
70            std::slice::from_raw_parts_mut(
71                &mut args as *mut _ as *mut u8,
72                std::mem::size_of_val(&args),
73            )
74        }),
75        ..Default::default()
76    };
77
78    let output = skel.progs.arena_alloc_mask.test_run(input)?;
79    if output.return_value != 0 {
80        bail!(
81            "Could not initialize arenas, setup_topology_node returned {}",
82            output.return_value as i32
83        );
84    }
85
86    let ptr = unsafe { std::mem::transmute::<u64, &mut [u64; 10]>(args.bitmap) };
87
88    let (valid_mask, _) = ptr.split_at_mut(mask.len());
89    valid_mask.clone_from_slice(mask);
90
91    let mut args = types::arena_topology_node_init_args {
92        bitmap: args.bitmap as c_ulong,
93        data_size: 0 as c_ulong,
94        id: 0 as c_ulong,
95    };
96
97    let input = ProgramInput {
98        context_in: Some(unsafe {
99            std::slice::from_raw_parts_mut(
100                &mut args as *mut _ as *mut u8,
101                std::mem::size_of_val(&args),
102            )
103        }),
104        ..Default::default()
105    };
106
107    let output = skel.progs.arena_topology_node_init.test_run(input)?;
108    if output.return_value != 0 {
109        bail!(
110            "p2dq_topology_node_init returned {}",
111            output.return_value as i32
112        );
113    }
114
115    Ok(())
116}
117
118fn setup_topology(skel: &mut BpfSkel<'_>) -> Result<()> {
119    let topo = Topology::new().expect("Failed to build host topology");
120
121    setup_topology_node(skel, topo.span.as_raw_slice())?;
122
123    for (_, node) in topo.nodes {
124        setup_topology_node(skel, node.span.as_raw_slice())?;
125    }
126
127    for (_, llc) in topo.all_llcs {
128        setup_topology_node(
129            skel,
130            Arc::<Llc>::into_inner(llc)
131                .expect("missing llc")
132                .span
133                .as_raw_slice(),
134        )?;
135    }
136
137    for (_, core) in topo.all_cores {
138        setup_topology_node(
139            skel,
140            Arc::<Core>::into_inner(core)
141                .expect("missing core")
142                .span
143                .as_raw_slice(),
144        )?;
145    }
146    for (_, cpu) in topo.all_cpus {
147        let mut mask = [0; 9];
148        mask[cpu.id.checked_shr(64).unwrap_or(0)] |= 1 << (cpu.id % 64);
149        setup_topology_node(skel, &mask)?;
150    }
151
152    Ok(())
153}
154
155fn main() {
156    TermLogger::init(
157        simplelog::LevelFilter::Info,
158        SimplelogConfig::default(),
159        TerminalMode::Mixed,
160        ColorChoice::Auto,
161    )
162    .unwrap();
163
164    let mut open_object = MaybeUninit::uninit();
165    let mut builder = BpfSkelBuilder::default();
166
167    builder.obj_builder.debug(true);
168    init_libbpf_logging(Some(PrintLevel::Debug));
169
170    let mut skel = builder
171        .open(&mut open_object)
172        .context("Failed to open BPF program")
173        .unwrap();
174
175    skel.maps.rodata_data.as_mut().unwrap().nr_cpu_ids = *NR_CPU_IDS as u32;
176
177    let mut skel = skel.load().context("Failed to load BPF program").unwrap();
178
179    setup_arenas(&mut skel).unwrap();
180    setup_topology(&mut skel).unwrap();
181
182    let input = ProgramInput {
183        ..Default::default()
184    };
185
186    let output = skel.progs.arena_selftest.test_run(input).unwrap();
187    if output.return_value != 0 {
188        println!(
189            "Selftest returned {}, please check bpf tracelog for more details.",
190            output.return_value as i32
191        );
192    }
193}
scx_lib_selftests/main.rs

scx_lib_selftests/
main.rs
