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

scx_utils/
bpf_builder.rs
