scx_arena/
arenalib.rs

1// SPDX-License-Identifier: GPL-2.0
2//
3// Copyright (c) 2025 Meta Platforms
4// Author: Emil Tsalapatis <etsal@meta.com>
5
6// This software may be used and distributed according to the terms of the
7// GNU General Public License version 2.
8
9pub use crate::bpf_skel::types;
10
11use scx_utils::Topology;
12use scx_utils::{Core, Llc};
13
14use std::ffi::CString;
15use std::os::raw::c_ulong;
16use std::sync::Arc;
17
18use anyhow::bail;
19use anyhow::Result;
20
21use libbpf_rs::libbpf_sys;
22use libbpf_rs::AsRawLibbpf;
23use libbpf_rs::Object;
24use libbpf_rs::ProgramInput;
25use libbpf_rs::ProgramMut;
26
27// MAX_CPU_ARRSZ has to be big enough to accomodate all present CPUs.
28// Even if it's larger than the size of cpumask_t, we truncate any
29// invalid data when passing it to the kernel's topology init functions.
30/// Maximum length of CPU mask supported by the library in bits.
31const MAX_CPU_SUPPORTED: usize = 640;
32
33/// Holds state related to BPF arenas in the program.
34#[derive(Debug)]
35pub struct ArenaLib<'a> {
36    task_size: usize,
37    obj: &'a mut Object,
38}
39
40impl<'a> ArenaLib<'a> {
41    /// Maximum CPU mask size, derived from MAX_CPU_SUPPORTED.
42    const MAX_CPU_ARRSZ: usize = (MAX_CPU_SUPPORTED + 63) / 64;
43
44    /// Amount of pages allocated at once form the BPF map. by the static stack allocator.
45    const STATIC_ALLOC_PAGES_GRANULARITY: c_ulong = 8;
46
47    fn run_prog_by_name(&self, name: &str, input: ProgramInput) -> Result<i32> {
48        let c_name = CString::new(name)?;
49        let ptr = unsafe {
50            libbpf_sys::bpf_object__find_program_by_name(
51                self.obj.as_libbpf_object().as_ptr(),
52                c_name.as_ptr(),
53            )
54        };
55        if ptr as u64 == 0 as u64 {
56            bail!("No program with name {} found in object", name);
57        }
58
59        let bpfprog = unsafe { &mut *ptr };
60        let prog = ProgramMut::new_mut(bpfprog);
61
62        let output = prog.test_run(input)?;
63
64        // Reach into the object and get the fd of the program
65        // Get the fd of the test program to run
66
67        return Ok(output.return_value as i32);
68    }
69
70    /// Set up basic library state.
71    fn setup_arena(&self) -> Result<()> {
72        // Allocate the arena memory from the BPF side so userspace initializes it before starting
73        // the scheduler. Despite the function call's name this is neither a test nor a test run,
74        // it's the recommended way of executing SEC("syscall") probes.
75        let mut args = types::arena_init_args {
76            static_pages: Self::STATIC_ALLOC_PAGES_GRANULARITY as c_ulong,
77            task_ctx_size: self.task_size as c_ulong,
78        };
79
80        let input = ProgramInput {
81            context_in: Some(unsafe {
82                std::slice::from_raw_parts_mut(
83                    &mut args as *mut _ as *mut u8,
84                    std::mem::size_of_val(&args),
85                )
86            }),
87            ..Default::default()
88        };
89
90        let ret = self.run_prog_by_name("arena_init", input)?;
91        if ret != 0 {
92            bail!("Could not initialize arenas, setup_arenas returned {}", ret);
93        }
94
95        Ok(())
96    }
97
98    fn setup_topology_node(&self, mask: &[u64], id: usize) -> Result<()> {
99        let mut args = types::arena_alloc_mask_args {
100            bitmap: 0 as c_ulong,
101        };
102
103        let input = ProgramInput {
104            context_in: Some(unsafe {
105                std::slice::from_raw_parts_mut(
106                    &mut args as *mut _ as *mut u8,
107                    std::mem::size_of_val(&args),
108                )
109            }),
110            ..Default::default()
111        };
112
113        let ret = self.run_prog_by_name("arena_alloc_mask", input)?;
114
115        if ret != 0 {
116            bail!(
117                "Could not initialize arenas, setup_topology_node returned {}",
118                ret
119            );
120        }
121
122        let ptr = unsafe {
123            &mut *std::ptr::with_exposed_provenance_mut::<[u64; 640]>(
124                args.bitmap.try_into().unwrap(),
125            )
126        };
127
128        let (valid_mask, _) = ptr.split_at_mut(mask.len());
129        valid_mask.clone_from_slice(mask);
130
131        let mut args = types::arena_topology_node_init_args {
132            bitmap: args.bitmap as c_ulong,
133            data_size: 0 as c_ulong,
134            id: id as c_ulong,
135        };
136
137        let input = ProgramInput {
138            context_in: Some(unsafe {
139                std::slice::from_raw_parts_mut(
140                    &mut args as *mut _ as *mut u8,
141                    std::mem::size_of_val(&args),
142                )
143            }),
144            ..Default::default()
145        };
146
147        let ret = self.run_prog_by_name("arena_topology_node_init", input)?;
148        if ret != 0 {
149            bail!("arena_topology_node_init returned {}", ret);
150        }
151
152        Ok(())
153    }
154
155    fn setup_topology(&self) -> Result<()> {
156        let topo = Topology::new().expect("Failed to build host topology");
157
158        // Top level - ID 0 is fine as there's only one top-level node
159        self.setup_topology_node(topo.span.as_raw_slice(), 0)?;
160
161        for (node_id, node) in topo.nodes {
162            self.setup_topology_node(node.span.as_raw_slice(), node_id)?;
163        }
164
165        // LLCs need to use their actual LLC ID for proper indexing in topo_nodes
166        for (llc_id, llc) in topo.all_llcs {
167            self.setup_topology_node(
168                Arc::<Llc>::into_inner(llc)
169                    .expect("missing llc")
170                    .span
171                    .as_raw_slice(),
172                llc_id,
173            )?;
174        }
175
176        for (core_id, core) in topo.all_cores {
177            self.setup_topology_node(
178                Arc::<Core>::into_inner(core)
179                    .expect("missing core")
180                    .span
181                    .as_raw_slice(),
182                core_id,
183            )?;
184        }
185        for (_, cpu) in topo.all_cpus {
186            let mut mask = [0; Self::MAX_CPU_ARRSZ - 1];
187            mask[cpu.id / 64] |= 1 << (cpu.id % 64);
188            self.setup_topology_node(&mask, cpu.id)?;
189        }
190
191        Ok(())
192    }
193
194    /// Create an Arenalib object This call only initializes the Rust side of Arenalib.
195    pub fn init(obj: &'a mut Object, task_size: usize, nr_cpus: usize) -> Result<Self> {
196        if nr_cpus >= MAX_CPU_SUPPORTED {
197            bail!("Scheduler specifies too many CPUs");
198        }
199
200        Ok(Self { task_size, obj })
201    }
202
203    /// Set up the BPF arena library state.
204    pub fn setup(&self) -> Result<()> {
205        self.setup_arena()?;
206        self.setup_topology()?;
207
208        Ok(())
209    }
210}
scx_arena/arenalib.rs

scx_arena/
arenalib.rs
