scx_utils/
bpf_builder.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.
5
6use crate::clang_info::ClangInfo;
7use anyhow::anyhow;
8use anyhow::Context;
9use anyhow::Result;
10use glob::glob;
11use libbpf_cargo::SkeletonBuilder;
12use libbpf_rs::Linker;
13use std::collections::BTreeSet;
14use std::env;
15use std::path::Path;
16use std::path::PathBuf;
17use tracing::Level;
18use tracing_subscriber::{filter, layer::SubscriberExt, Layer};
19
20#[derive(Debug)]
21/// # Build helpers for sched_ext schedulers with Rust userspace component
22///
23/// This is to be used from `build.rs` of a cargo project which implements a
24/// [sched_ext](https://github.com/sched-ext/scx) scheduler with C BPF
25/// component and Rust userspace component. `BpfBuilder` provides everything
26/// necessary to build the BPF component and generate Rust bindings.
27/// BpfBuilder provides the following.
28///
29/// 1. *`vmlinux.h` and other common BPF header files*
30///
31/// All sched_ext BPF implementations require `vmlinux.h` and many make use
32/// of common constructs such as
33/// [`user_exit_info`](https://github.com/sched-ext/scx/blob/main/scheds/include/common/user_exit_info.h).
34/// `BpfBuilder` makes these headers available when compiling BPF source
35/// code and generating bindings for it. The included headers can be browsed
36/// at <https://github.com/sched-ext/scx/tree/main/scheds/include>.
37///
38/// These headers can be superseded using environment variables which will
39/// be discussed later.
40///
41/// 2. *Header bindings using `bindgen`*
42///
43/// If enabled with `.enable_intf()`, the input `.h` file is processed by
44/// `bindgen` to generate Rust bindings. This is useful in establishing
45/// shared constants and data types between the BPF and user components.
46///
47/// Note that the types generated with `bindgen` are different from the
48/// types used by the BPF skeleton even when they are the same types in BPF.
49/// This is a source of ugliness and we are hoping to address it by
50/// improving `libbpf-cargo` in the future.
51///
52/// 3. *BPF compilation and generation of the skeleton and its bindings*
53///
54/// If enabled with `.enable_skel()`, the input `.bpf.c` file is compiled
55/// and its skeleton and bindings are generated using `libbpf-cargo`.
56///
57/// ## An Example
58///
59/// This section shows how `BpfBuilder` can be used in an example project.
60/// For a concrete example, take a look at
61/// [`scx_rusty`](https://github.com/sched-ext/scx/tree/main/scheds/rust/scx_rusty).
62///
63/// A minimal source tree using all features would look like the following:
64///
65/// ```text
66/// scx_hello_world
67/// |-- Cargo.toml
68/// |-- build.rs
69/// \-- src
70///     |-- main.rs
71///     |-- bpf_intf.rs
72///     |-- bpf_skel.rs
73///     \-- bpf
74///         |-- intf.h
75///         \-- main.c
76/// ```
77///
78/// The following three files would contain the actual implementation:
79///
80/// - `src/main.rs`: Rust userspace component which loads the BPF blob and
81/// interacts it using the generated bindings.
82///
83/// - `src/bpf/intf.h`: C header file definining constants and structs
84/// that will be used by both the BPF and userspace components.
85///
86/// - `src/bpf/main.c`: C source code implementing the BPF component -
87/// including `struct sched_ext_ops`.
88///
89/// And then there are boilerplates to generate the bindings and make them
90/// available as modules to `main.rs`.
91///
92/// - `Cargo.toml`: Includes `scx_utils` in the `[build-dependencies]`
93/// section.
94///
95/// - `build.rs`: Uses `scx_utils::BpfBuilder` to build and generate
96/// bindings for the BPF component. For this project, it can look like the
97/// following.
98///
99/// ```should_panic
100/// fn main() {
101///     scx_utils::BpfBuilder::new()
102///         .unwrap()
103///         .enable_intf("src/bpf/intf.h", "bpf_intf.rs")
104///         .enable_skel("src/bpf/main.bpf.c", "bpf")
105///         .build()
106///         .unwrap();
107/// }
108/// ```
109///
110/// - `bpf_intf.rs`: Import the bindings generated by `bindgen` into a
111/// module. Above, we told `.enable_intf()` to generate the bindings into
112/// `bpf_intf.rs`, so the file would look like the following. The `allow`
113/// directives are useful if the header is including `vmlinux.h`.
114///
115/// ```ignore
116/// #![allow(non_upper_case_globals)]
117/// #![allow(non_camel_case_types)]
118/// #![allow(non_snake_case)]
119/// #![allow(dead_code)]
120///
121/// include!(concat!(env!("OUT_DIR"), "/bpf_intf.rs"));
122/// ```
123///
124/// - `bpf_skel.rs`: Import the BPF skeleton bindings generated by
125/// `libbpf-cargo` into a module. Above, we told `.enable_skel()` to use the
126/// skeleton name `bpf`, so the file would look like the following.
127///
128/// ```ignore
129/// include!(concat!(env!("OUT_DIR"), "/bpf_skel.rs"));
130/// ```
131///
132/// ## Compiler Flags and Environment Variables
133///
134/// BPF being its own CPU architecture and independent runtime environment,
135/// build environment and steps are already rather complex. The need to
136/// interface between two different languages - C and Rust - adds further
137/// complexities. `BpfBuilder` automates most of the process. The determined
138/// build environment is recorded in the `build.rs` output and can be
139/// obtained with a command like the following:
140///
141/// ```text
142/// $ grep '^scx_utils:clang=' target/release/build/scx_rusty-*/output
143/// ```
144///
145/// While the automatic settings should work most of the time, there can be
146/// times when overriding them is necessary. The following environment
147/// variables can be used to customize the build environment.
148///
149/// - `BPF_CLANG`: The clang command to use. (Default: `clang`)
150///
151/// - `BPF_CFLAGS`: Compiler flags to use when building BPF source code. If
152///   specified, the flags from this variable are the only flags passed to
153///   the compiler. `BpfBuilder` won't generate any flags including `-I`
154///   flags for the common header files and other `CFLAGS` related variables
155///   are ignored.
156///
157/// - `BPF_BASE_CFLAGS`: Override the non-include part of cflags.
158///
159/// - `BPF_EXTRA_CFLAGS_PRE_INCL`: Add cflags before the automic include
160///   search path options. Header files in the search paths added by this
161///   variable will supercede the automatic ones.
162///
163/// - `BPF_EXTRA_CFLAGS_POST_INCL`: Add cflags after the automic include
164///   search path options. Header paths added by this variable will be
165///   searched only if the target header file can't be found in the
166///   automatic header paths.
167///
168/// - `RUSTFLAGS`: This is a generic `cargo` flag and can be useful for
169///   specifying extra linker flags.
170///
171/// A common case for using the above flags is using the latest `libbpf`
172/// from the kernel tree. Let's say the kernel tree is at `$KERNEL` and
173/// `libbpf`. The following builds `libbpf` shipped with the kernel:
174///
175/// ```test
176/// $ cd $KERNEL
177/// $ make -C tools/bpf/bpftool
178/// ```
179///
180/// To link the scheduler against the resulting `libbpf`:
181///
182/// ```test
183/// $ env BPF_EXTRA_CFLAGS_POST_INCL=$KERNEL/tools/bpf/bpftool/libbpf/include \
184///   RUSTFLAGS="-C link-args=-lelf -C link-args=-lz -C link-args=-lzstd \
185///   -L$KERNEL/tools/bpf/bpftool/libbpf" cargo build --release
186/// ```
187pub struct BpfBuilder {
188    clang: ClangInfo,
189    cflags: Vec<String>,
190    out_dir: PathBuf,
191    sources: BTreeSet<String>,
192
193    intf_input_output: Option<(String, String)>,
194    skel_input_name: Option<(String, String)>,
195}
196
197impl BpfBuilder {
198    const BPF_H_TAR: &'static [u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bpf_h.tar"));
199
200    fn install_bpf_h<P: AsRef<Path>>(dest: P) -> Result<()> {
201        let mut ar = tar::Archive::new(Self::BPF_H_TAR);
202        ar.unpack(dest)?;
203        Ok(())
204    }
205
206    fn determine_cflags<P>(clang: &ClangInfo, out_dir: P) -> Result<Vec<String>>
207    where
208        P: AsRef<Path> + std::fmt::Debug,
209    {
210        let bpf_h = out_dir
211            .as_ref()
212            .join("scx_utils-bpf_h")
213            .to_str()
214            .ok_or(anyhow!(
215                "{:?}/scx_utils-bph_h can't be converted to str",
216                &out_dir
217            ))?
218            .to_string();
219        Self::install_bpf_h(&bpf_h)?;
220
221        let mut cflags = Vec::<String>::new();
222
223        cflags.append(&mut match env::var("BPF_BASE_CFLAGS") {
224            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
225            _ => clang.determine_base_cflags()?,
226        });
227
228        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_PRE_INCL") {
229            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
230            _ => vec![],
231        });
232
233        cflags.push(format!(
234            "-I{}/arch/{}",
235            &bpf_h,
236            &clang.kernel_target().unwrap()
237        ));
238        cflags.push(format!("-I{}", &bpf_h));
239        cflags.push(format!("-I{}/bpf-compat", &bpf_h));
240
241        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_POST_INCL") {
242            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
243            _ => vec![],
244        });
245
246        Ok(cflags)
247    }
248
249    /// Create a new `BpfBuilder` struct. Call `enable` and `set` methods to
250    /// configure and `build` method to compile and generate bindings. See
251    /// the struct documentation for details.
252    pub fn new() -> Result<Self> {
253        let out_dir = PathBuf::from(env::var("OUT_DIR")?);
254
255        let clang = ClangInfo::new()?;
256        let cflags = match env::var("BPF_CFLAGS") {
257            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
258            _ => Self::determine_cflags(&clang, &out_dir)?,
259        };
260
261        println!("scx_utils:clang={:?} {:?}", &clang, &cflags);
262
263        Ok(Self {
264            clang,
265            cflags,
266            out_dir,
267
268            sources: BTreeSet::new(),
269            intf_input_output: None,
270            skel_input_name: None,
271        })
272    }
273
274    /// Enable generation of header bindings using `bindgen`. `@input` is
275    /// the `.h` file defining the constants and types to be shared between
276    /// BPF and Rust components. `@output` is the `.rs` file to be
277    /// generated.
278    pub fn enable_intf(&mut self, input: &str, output: &str) -> &mut Self {
279        self.intf_input_output = Some((input.into(), output.into()));
280        self
281    }
282
283    /// Enable compilation of BPF code and generation of the skeleton and
284    /// its Rust bindings. `@input` is the `.bpf.c` file containing the BPF
285    /// source code and `@output` is the `.rs` file to be generated.
286    pub fn enable_skel(&mut self, input: &str, name: &str) -> &mut Self {
287        self.skel_input_name = Some((input.into(), name.into()));
288        self.sources.insert(input.into());
289
290        self
291    }
292
293    fn input_insert_deps(&self, deps: &mut BTreeSet<String>) -> () {
294        let (input, _) = match &self.intf_input_output {
295            Some(pair) => pair,
296            None => return,
297        };
298
299        // Tell cargo to invalidate the built crate whenever the wrapper changes
300        deps.insert(input.to_string());
301    }
302
303    fn bindgen_bpf_intf(&self) -> Result<()> {
304        let (input, output) = match &self.intf_input_output {
305            Some(pair) => pair,
306            None => return Ok(()),
307        };
308
309        // The bindgen::Builder is the main entry point to bindgen, and lets
310        // you build up options for the resulting bindings.
311        let bindings = bindgen::Builder::default()
312            // Should run clang with the same -I options as BPF compilation.
313            .clang_args(
314                self.cflags
315                    .iter()
316                    .chain(["-target".into(), "bpf".into()].iter()),
317            )
318            // The input header we would like to generate bindings for.
319            .header(input)
320            // Tell cargo to invalidate the built crate whenever any of the
321            // included header files changed.
322            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
323            .generate()
324            .context("Unable to generate bindings")?;
325
326        bindings
327            .write_to_file(self.out_dir.join(output))
328            .context("Couldn't write bindings")
329    }
330
331    pub fn add_source(&mut self, input: &str) -> &mut Self {
332        self.sources.insert(input.into());
333        self
334    }
335
336    pub fn compile_link_gen(&mut self) -> Result<()> {
337        let input = match &self.skel_input_name {
338            Some((name, _)) => name,
339            None => return Ok(()),
340        };
341
342        // Better to hardcode the name because it gets embedded in
343        // all the skeleton struct/member definitions.
344        let linkobj = self.out_dir.join("bpf.bpf.o");
345        let mut linker = Linker::new(&linkobj)?;
346
347        for filename in self.sources.iter() {
348            let name = Path::new(filename).file_name().unwrap().to_str().unwrap();
349            let obj = self.out_dir.join(name.replace(".bpf.c", ".bpf.o"));
350
351            with_clang_warnings(|| {
352                SkeletonBuilder::new()
353                    .source(filename)
354                    .rustfmt("disable_rustfmt")
355                    .obj(&obj)
356                    .clang(&self.clang.clang)
357                    .clang_args(&self.cflags)
358                    .build()
359            })?;
360
361            linker.add_file(&obj)?;
362        }
363
364        linker.link()?;
365
366        self.bindgen_bpf_intf()?;
367
368        let skel_path = self.out_dir.join("bpf_skel.rs");
369
370        SkeletonBuilder::new()
371            .obj(&linkobj)
372            .clang(&self.clang.clang)
373            .clang_args(&self.cflags)
374            .rustfmt("disable_rustfmt")
375            .generate(&skel_path)?;
376
377        let mut deps = BTreeSet::new();
378        self.add_src_deps(&mut deps, input)?;
379        for filename in self.sources.iter() {
380            deps.insert(filename.to_string());
381        }
382
383        self.gen_cargo_reruns(Some(&deps))?;
384
385        Ok(())
386    }
387
388    fn gen_bpf_skel(&self, deps: &mut BTreeSet<String>) -> Result<()> {
389        let (input, name) = match &self.skel_input_name {
390            Some(pair) => pair,
391            None => return Ok(()),
392        };
393
394        let obj = self.out_dir.join(format!("{name}.bpf.o"));
395        let skel_path = self.out_dir.join(format!("{name}_skel.rs"));
396
397        with_clang_warnings(|| {
398            SkeletonBuilder::new()
399                .source(input)
400                .obj(&obj)
401                .clang(&self.clang.clang)
402                .clang_args(&self.cflags)
403                .rustfmt("disable_rustfmt")
404                .build_and_generate(&skel_path)
405        })?;
406
407        self.add_src_deps(deps, input)?;
408
409        Ok(())
410    }
411
412    fn add_src_deps(&self, deps: &mut BTreeSet<String>, input: &str) -> Result<()> {
413        let c_path = PathBuf::from(input);
414        let dir = c_path
415            .parent()
416            .ok_or(anyhow!("Source {:?} doesn't have parent dir", c_path))?
417            .to_str()
418            .ok_or(anyhow!("Parent dir of {:?} isn't a UTF-8 string", c_path))?;
419
420        for path in glob(&format!("{dir}/*.[hc]"))?.filter_map(Result::ok) {
421            deps.insert(
422                path.to_str()
423                    .ok_or(anyhow!("Path {:?} is not a valid string", path))?
424                    .to_string(),
425            );
426        }
427
428        Ok(())
429    }
430
431    fn gen_cargo_reruns(&self, dependencies: Option<&BTreeSet<String>>) -> Result<()> {
432        println!("cargo:rerun-if-env-changed=BPF_CLANG");
433        println!("cargo:rerun-if-env-changed=BPF_CFLAGS");
434        println!("cargo:rerun-if-env-changed=BPF_BASE_CFLAGS");
435        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_PRE_INCL");
436        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_POST_INCL");
437        if let Some(deps) = dependencies {
438            for dep in deps.iter() {
439                println!("cargo:rerun-if-changed={dep}");
440            }
441        };
442
443        for source in self.sources.iter() {
444            println!("cargo:rerun-if-changed={source}");
445        }
446
447        Ok(())
448    }
449
450    /// Build and generate the enabled bindings.
451    pub fn build(&self) -> Result<()> {
452        let mut deps = BTreeSet::new();
453
454        self.input_insert_deps(&mut deps);
455
456        self.bindgen_bpf_intf()?;
457        self.gen_bpf_skel(&mut deps)?;
458        self.gen_cargo_reruns(Some(&deps))?;
459        Ok(())
460    }
461}
462
463// Helper function to set up tracing and output compiler warnings
464fn with_clang_warnings<F, R>(f: F) -> Result<R>
465where
466    F: FnOnce() -> Result<R>,
467{
468    let fmt =
469        tracing_subscriber::fmt::layer().event_format(libbpf_cargo::util::CargoWarningFormatter);
470    let filter = filter::Targets::new()
471        .with_target("compiler-stdout", Level::INFO)
472        .with_target("compiler-stderr", Level::INFO);
473
474    let subscriber = tracing_subscriber::registry().with(fmt.with_filter(filter));
475
476    // Execute the closure with a tracing guard
477    Ok({
478        let _guard = tracing::subscriber::set_default(subscriber);
479        f()
480    }?)
481}
482
483#[cfg(test)]
484mod tests {
485    use regex::Regex;
486    use sscanf::sscanf;
487
488    use crate::builder::ClangInfo;
489
490    #[test]
491    fn test_bpf_builder_new() {
492        let td = tempfile::tempdir().unwrap();
493        unsafe { std::env::set_var("OUT_DIR", td.path()) };
494
495        let res = super::BpfBuilder::new();
496        assert!(res.is_ok(), "Failed to create BpfBuilder ({res:?})");
497    }
498
499    #[test]
500    fn test_vmlinux_h_ver_sha1() {
501        let clang_info = ClangInfo::new().unwrap();
502
503        let mut ar = tar::Archive::new(super::BpfBuilder::BPF_H_TAR);
504        let mut found = false;
505
506        let pattern = Regex::new(r"arch\/.*\/vmlinux-.*.h").unwrap();
507
508        for entry in ar.entries().unwrap() {
509            let entry = entry.unwrap();
510            let file_name = entry.header().path().unwrap();
511            let file_name_str = file_name.to_string_lossy().to_owned();
512            if file_name_str.contains(&clang_info.kernel_target().unwrap()) {
513                found = true;
514            }
515            if !pattern.find(&file_name_str).is_some() {
516                continue;
517            }
518
519            println!("checking {file_name_str}");
520
521            let (arch, ver, sha1) =
522                sscanf!(file_name_str, "arch/{String}/vmlinux-v{String}-g{String}.h").unwrap();
523            println!(
524                "vmlinux.h: arch={:?} ver={:?} sha1={:?}",
525                &arch, &ver, &sha1,
526            );
527
528            assert!(regex::Regex::new(r"^([1-9][0-9]*\.[1-9][0-9][a-z0-9-]*)$")
529                .unwrap()
530                .is_match(&ver));
531            assert!(regex::Regex::new(r"^[0-9a-z]{12}$")
532                .unwrap()
533                .is_match(&sha1));
534        }
535
536        assert!(found);
537    }
538}