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 accommodate 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/// Live BPF arena library state. Returned by setup() and must be kept alive
34/// for as long as the scheduler instance uses the arena: dropping it stops
35/// and joins the library's background threads.
36#[must_use]
37#[derive(Debug)]
38pub struct ArenaLib {
39    _watcher: crate::Daemon,
40    _urcu: Option<crate::Daemon>,
41}
42
43impl ArenaLib {
44    /// Maximum CPU mask size, derived from MAX_CPU_SUPPORTED.
45    const MAX_CPU_ARRSZ: usize = (MAX_CPU_SUPPORTED + 63) / 64;
46
47    /// Amount of pages allocated at once form the BPF map. by the static stack allocator.
48    const STATIC_ALLOC_PAGES_GRANULARITY: c_ulong = 8;
49
50    fn run_prog_by_name(obj: &Object, name: &str, input: ProgramInput) -> Result<i32> {
51        let c_name = CString::new(name)?;
52        let ptr = unsafe {
53            libbpf_sys::bpf_object__find_program_by_name(
54                obj.as_libbpf_object().as_ptr(),
55                c_name.as_ptr(),
56            )
57        };
58        if ptr as u64 == 0 as u64 {
59            bail!("No program with name {} found in object", name);
60        }
61
62        let bpfprog = unsafe { &mut *ptr };
63        let prog = ProgramMut::new_mut(bpfprog);
64
65        let output = prog.test_run(input)?;
66
67        // Reach into the object and get the fd of the program
68        // Get the fd of the test program to run
69
70        return Ok(output.return_value as i32);
71    }
72
73    /// Set up basic library state.
74    fn setup_arena(obj: &Object, task_size: usize, task_align: usize) -> Result<()> {
75        // Allocate the arena memory from the BPF side so userspace initializes it before starting
76        // the scheduler. Despite the function call's name this is neither a test nor a test run,
77        // it's the recommended way of executing SEC("syscall") probes.
78        let mut args = types::arena_init_args {
79            static_pages: Self::STATIC_ALLOC_PAGES_GRANULARITY as c_ulong,
80            task_ctx_size: task_size as c_ulong,
81            task_ctx_align: task_align as c_ulong,
82        };
83
84        let input = ProgramInput {
85            context_in: Some(unsafe {
86                std::slice::from_raw_parts_mut(
87                    &mut args as *mut _ as *mut u8,
88                    std::mem::size_of_val(&args),
89                )
90            }),
91            ..Default::default()
92        };
93
94        let ret = Self::run_prog_by_name(obj, "arena_init", input)?;
95        if ret != 0 {
96            bail!("Could not initialize arenas, setup_arenas returned {}", ret);
97        }
98
99        Ok(())
100    }
101
102    fn setup_topology_node(obj: &Object, mask: &[u64], id: usize) -> Result<()> {
103        let mut args = types::arena_alloc_mask_args {
104            bitmap: 0 as c_ulong,
105        };
106
107        // Exclude memory-only NUMA nodes
108        if mask.into_iter().all(|&b| b == 0) {
109            return Ok(());
110        }
111
112        let input = ProgramInput {
113            context_in: Some(unsafe {
114                std::slice::from_raw_parts_mut(
115                    &mut args as *mut _ as *mut u8,
116                    std::mem::size_of_val(&args),
117                )
118            }),
119            ..Default::default()
120        };
121
122        let ret = Self::run_prog_by_name(obj, "arena_alloc_mask", input)?;
123
124        if ret != 0 {
125            bail!(
126                "Could not initialize arenas, setup_topology_node returned {}",
127                ret
128            );
129        }
130
131        let ptr = unsafe {
132            &mut *std::ptr::with_exposed_provenance_mut::<[u64; 640]>(
133                args.bitmap.try_into().unwrap(),
134            )
135        };
136
137        let (valid_mask, _) = ptr.split_at_mut(mask.len());
138        valid_mask.clone_from_slice(mask);
139
140        let mut args = types::arena_topology_node_init_args {
141            bitmap: args.bitmap as c_ulong,
142            data_size: 0 as c_ulong,
143            id: id as c_ulong,
144        };
145
146        let input = ProgramInput {
147            context_in: Some(unsafe {
148                std::slice::from_raw_parts_mut(
149                    &mut args as *mut _ as *mut u8,
150                    std::mem::size_of_val(&args),
151                )
152            }),
153            ..Default::default()
154        };
155
156        let ret = Self::run_prog_by_name(obj, "arena_topology_node_init", input)?;
157        if ret != 0 {
158            bail!("arena_topology_node_init returned {}", ret);
159        }
160
161        Ok(())
162    }
163
164    /// Set the per-level maximum number of children before registering topology
165    /// nodes. Each topology node at level L is allocated with
166    /// topo_max_children[L] child pointer slots, so these values must be set
167    /// before any arena_topology_node_init() calls. The sizes are derived from
168    /// the actual host topology to keep per-node allocation as small as
169    /// possible.
170    ///
171    /// NOTE: rust/scx_arena/selftests/src/main.rs::setup_topology() contains
172    /// equivalent logic and must be kept in sync with this function.
173    fn setup_topology_max_children(obj: &Object, topo: &Topology) -> Result<()> {
174        // Compute the maximum number of children at each topology level.
175        // TOPO_TOP  (0): children are NUMA nodes
176        // TOPO_NODE (1): children are LLCs
177        // TOPO_LLC  (2): children are cores
178        // TOPO_CORE (3): children are CPUs (SMT threads)
179        // TOPO_CPU  (4): leaf nodes, no children
180        let max_children: [u32; 5] = [
181            topo.nodes.len() as u32,
182            topo.nodes.values().map(|n| n.llcs.len()).max().unwrap_or(0) as u32,
183            topo.all_llcs
184                .values()
185                .map(|l| l.cores.len())
186                .max()
187                .unwrap_or(0) as u32,
188            topo.all_cores
189                .values()
190                .map(|c| c.cpus.len())
191                .max()
192                .unwrap_or(0) as u32,
193            0,
194        ];
195
196        let mut args = types::arena_topology_init_args { max_children };
197
198        let input = ProgramInput {
199            context_in: Some(unsafe {
200                std::slice::from_raw_parts_mut(
201                    &mut args as *mut _ as *mut u8,
202                    std::mem::size_of_val(&args),
203                )
204            }),
205            ..Default::default()
206        };
207
208        let ret = Self::run_prog_by_name(obj, "arena_topology_init", input)?;
209        if ret != 0 {
210            bail!("arena_topology_init returned {}", ret);
211        }
212
213        Ok(())
214    }
215
216    fn setup_topology(obj: &Object) -> Result<()> {
217        let topo = Topology::new().expect("Failed to build host topology");
218
219        Self::setup_topology_max_children(obj, &topo)?;
220
221        // Top level - ID 0 is fine as there's only one top-level node
222        Self::setup_topology_node(obj, topo.span.as_raw_slice(), 0)?;
223
224        for (node_id, node) in topo.nodes {
225            Self::setup_topology_node(obj, node.span.as_raw_slice(), node_id)?;
226        }
227
228        // LLCs need to use their actual LLC ID for proper indexing in topo_nodes
229        for (llc_id, llc) in topo.all_llcs {
230            Self::setup_topology_node(
231                obj,
232                Arc::<Llc>::into_inner(llc)
233                    .expect("missing llc")
234                    .span
235                    .as_raw_slice(),
236                llc_id,
237            )?;
238        }
239
240        for (core_id, core) in topo.all_cores {
241            Self::setup_topology_node(
242                obj,
243                Arc::<Core>::into_inner(core)
244                    .expect("missing core")
245                    .span
246                    .as_raw_slice(),
247                core_id,
248            )?;
249        }
250        for (_, cpu) in topo.all_cpus {
251            let mut mask = [0; Self::MAX_CPU_ARRSZ - 1];
252            mask[cpu.id / 64] |= 1 << (cpu.id % 64);
253            Self::setup_topology_node(obj, &mask, cpu.id)?;
254        }
255
256        Ok(())
257    }
258
259    /// Set up the BPF arena library state and, when the object carries the
260    /// scx_urcu doorbell, spawn the reclaim daemon. The returned ArenaLib
261    /// owns the library's background threads.
262    /// @task_align: task ctx element alignment, 0 for word alignment.
263    pub fn setup(
264        obj: &Object,
265        task_size: usize,
266        task_align: usize,
267        nr_cpus: usize,
268    ) -> Result<ArenaLib> {
269        if nr_cpus >= MAX_CPU_SUPPORTED {
270            bail!("Scheduler specifies too many CPUs");
271        }
272
273        Self::setup_arena(obj, task_size, task_align)?;
274        Self::setup_topology(obj)?;
275
276        Self::start(obj)
277    }
278
279    /// Start the userspace services for BPF arena state initialized by the
280    /// caller. The returned ArenaLib must be kept alive for as long as the BPF
281    /// object uses the arena.
282    ///
283    /// Use this instead of setup() when a scheduler has its own BPF-side arena
284    /// initialization and does not use the generic arena topology.
285    pub fn start(obj: &Object) -> Result<ArenaLib> {
286        Ok(ArenaLib {
287            _watcher: crate::stream_watcher_spawn(obj)?,
288            _urcu: crate::urcu_spawn(obj)?,
289        })
290    }
291}
scx_arena/arenalib.rs

scx_arena/
arenalib.rs
