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