scx_arena_libarena_selftests/
main.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
6mod bpf_skel;
7use bpf_skel::*;
8
9use std::io::IsTerminal;
10use std::mem::MaybeUninit;
11
12use anyhow::{Context, Result, bail};
13use clap::Parser;
14use libbpf_rs::ProgramInput;
15use libbpf_rs::skel::{OpenSkel, Skel, SkelBuilder};
16use scx_arena::ScxLibArena;
17use scx_utils::{NR_CPU_IDS, init_libbpf_logging};
18use simplelog::{ColorChoice, Config as SimplelogConfig, TermLogger, TerminalMode};
19
20const COLOR_BRIGHT_GREEN: &str = "\x1b[92m";
21const COLOR_BRIGHT_RED: &str = "\x1b[91m";
22const COLOR_RESET: &str = "\x1b[0m";
23
24const SUITES: &[&str] = &["atq", "btree", "lvqueue", "minheap", "rbtree", "topology"];
25
26#[derive(Debug, Parser)]
27#[clap(about = "scx_arena library selftests")]
28struct Opts {
29    #[clap(long)]
30    list: bool,
31
32    #[clap(long = "test", value_name = "NAME", num_args(1..))]
33    tests: Vec<String>,
34}
35
36fn available_tests() -> String {
37    SUITES
38        .iter()
39        .map(|name| format!("  {}", name))
40        .collect::<Vec<_>>()
41        .join("\n")
42}
43
44fn colorize(text: &str, color: &str, is_tty: bool) -> String {
45    if is_tty {
46        format!("{}{}{}", color, text, COLOR_RESET)
47    } else {
48        text.to_string()
49    }
50}
51
52fn suite_prog<'a, 'obj>(
53    skel: &'a BpfSkel<'obj>,
54    name: &str,
55) -> Result<&'a libbpf_rs::ProgramMut<'obj>> {
56    Ok(match name {
57        "atq" => &skel.progs.arena_selftest_atq,
58        "btree" => &skel.progs.arena_selftest_btree,
59        "lvqueue" => &skel.progs.arena_selftest_lvqueue,
60        "minheap" => &skel.progs.arena_selftest_minheap,
61        "rbtree" => &skel.progs.arena_selftest_rbtree,
62        "topology" => &skel.progs.arena_selftest_topology,
63        _ => bail!("unknown test: {}", name),
64    })
65}
66
67fn run_test(skel: &BpfSkel<'_>, name: &str) -> Result<i32> {
68    let output = suite_prog(skel, name)?.test_run(ProgramInput::default())?;
69    Ok(output.return_value as i32)
70}
71
72fn main() {
73    TermLogger::init(
74        simplelog::LevelFilter::Info,
75        SimplelogConfig::default(),
76        TerminalMode::Mixed,
77        ColorChoice::Auto,
78    )
79    .unwrap();
80
81    let opts = Opts::parse();
82    if opts.list {
83        println!("Available test cases:\n{}", available_tests());
84        return;
85    }
86    for name in &opts.tests {
87        if !SUITES.contains(&name.as_str()) {
88            eprintln!(
89                "Unknown test: '{}'.\nAvailable tests:\n{}",
90                name,
91                available_tests()
92            );
93            std::process::exit(1);
94        }
95    }
96
97    let mut open_object = MaybeUninit::uninit();
98    let mut builder = BpfSkelBuilder::default();
99    builder.obj_builder.debug(true);
100    init_libbpf_logging(Some(libbpf_rs::PrintLevel::Debug));
101
102    let mut open = builder
103        .open(&mut open_object)
104        .context("Failed to open BPF program")
105        .unwrap();
106    open.maps.rodata_data.as_mut().unwrap().nr_cpu_ids = *NR_CPU_IDS as u32;
107    let skel = open.load().context("Failed to load BPF program").unwrap();
108
109    let _arena = ScxLibArena::setup(skel.object(), 42, 0, *NR_CPU_IDS)
110        .context("Failed to set up arena")
111        .unwrap();
112
113    let to_run: Vec<&str> = if opts.tests.is_empty() {
114        SUITES.to_vec()
115    } else {
116        opts.tests.iter().map(String::as_str).collect()
117    };
118    let stdout_tty = std::io::stdout().is_terminal();
119    let stderr_tty = std::io::stderr().is_terminal();
120    let pass_label = colorize("[ PASS ]", COLOR_BRIGHT_GREEN, stdout_tty);
121    let fail_label = colorize("[ FAIL ]", COLOR_BRIGHT_RED, stderr_tty);
122    let mut any_failed = false;
123
124    for name in to_run {
125        match run_test(&skel, name) {
126            Ok(0) => println!("{} {}", pass_label, name),
127            Ok(ret) => {
128                eprintln!("{} {} (returned {})", fail_label, name, ret);
129                any_failed = true;
130            }
131            Err(err) => {
132                eprintln!("{} {} ({})", fail_label, name, err);
133                any_failed = true;
134            }
135        }
136    }
137
138    if any_failed {
139        eprintln!(
140            "{}",
141            colorize(
142                "One or more selftests failed.",
143                COLOR_BRIGHT_RED,
144                stderr_tty
145            )
146        );
147        std::process::exit(1);
148    }
149    println!(
150        "{}",
151        colorize("All selftests passed.", COLOR_BRIGHT_GREEN, stdout_tty)
152    );
153}
scx_arena_libarena_selftests/main.rs

scx_arena_libarena_selftests/
main.rs
