Add --include and --exclude flags to the bench_runtime_local command

Author: Kobzol

collector/benchlib/src/benchmark.rs

@@ -15,9 +15,16 @@ pub fn benchmark_suite<F: FnOnce(&mut BenchmarkSuite)>(define_func: F) {
     suite.run().expect("Benchmark suite has failed");
 }
 
+/// Type-erased function that performs a benchmark.
+struct BenchmarkWrapper {
+    func: Box<dyn Fn() -> anyhow::Result<BenchmarkResult>>,
+}
+
+type BenchmarkMap = HashMap<&'static str, BenchmarkWrapper>;
+
 #[derive(Default)]
 pub struct BenchmarkSuite {
-    benchmarks: HashMap<&'static str, BenchmarkWrapper>,
+    benchmarks: BenchmarkMap,
 }
 
 impl BenchmarkSuite {
@@ -54,33 +61,42 @@ impl BenchmarkSuite {
         let args = parse_cli()?;
         match args {
             Args::Benchmark(args) => {
-                self.run_benchmark(args)?;
+                run_benchmark(args, self.benchmarks)?;
             }
         }
 
         Ok(())
     }
+}
 
-    fn run_benchmark(self, args: BenchmarkArgs) -> anyhow::Result<()> {
-        let mut items: Vec<_> = self.benchmarks.into_iter().collect();
-        items.sort_unstable_by_key(|item| item.0);
+fn run_benchmark(args: BenchmarkArgs, benchmarks: BenchmarkMap) -> anyhow::Result<()> {
+    let mut items: Vec<(&'static str, BenchmarkWrapper)> = benchmarks
+        .into_iter()
+        .filter(|(name, _)| passes_filter(name, args.exclude.as_deref(), args.include.as_deref()))
+        .collect();
+    items.sort_unstable_by_key(|item| item.0);
 
-        let mut results: Vec<BenchmarkResult> = Vec::with_capacity(items.len());
-        for (name, def) in items {
-            for i in 0..args.iterations {
-                let result = (def.func)()?;
-                log::info!("Benchmark (run {i}) `{}` completed: {:?}", name, result);
-                results.push(result);
-            }
+    let mut results: Vec<BenchmarkResult> = Vec::with_capacity(items.len());
+    for (name, def) in items {
+        for i in 0..args.iterations {
+            let result = (def.func)()?;
+            log::info!("Benchmark (run {i}) `{name}` completed: {result:?}");
+            results.push(result);
         }
-
-        println!("{}", serde_json::to_string(&results)?);
-        Ok(())
     }
+
+    println!("{}", serde_json::to_string(&results)?);
+    Ok(())
 }
 
-struct BenchmarkWrapper {
-    func: Box<dyn Fn() -> anyhow::Result<BenchmarkResult>>,
+/// Tests if the name of the benchmark passes through the include and exclude filter flags.
+fn passes_filter(name: &str, exclude: Option<&str>, include: Option<&str>) -> bool {
+    match (exclude, include) {
+        (Some(exclude), Some(include)) => name.starts_with(include) && !name.starts_with(exclude),
+        (None, Some(include)) => name.starts_with(include),
+        (Some(exclude), None) => !name.starts_with(&exclude),
+        (None, None) => true,
+    }
 }
 
 /// Copied from `iai`, so that we don't have to use unstable features.

collector/benchlib/src/cli.rs

@@ -11,6 +11,14 @@ pub struct BenchmarkArgs {
     /// How many times should each benchmark be repeated.
     #[clap(long, default_value = "5")]
     pub iterations: u32,
+
+    /// Exclude all benchmarks matching a prefix in this comma-separated list
+    #[clap(long)]
+    pub exclude: Option<String>,
+
+    /// Include only benchmarks matching a prefix in this comma-separated list
+    #[clap(long)]
+    pub include: Option<String>,
 }
 
 pub fn parse_cli() -> anyhow::Result<Args> {

collector/src/bin/collector.rs

@@ -562,10 +562,19 @@ struct BenchRustcOption {
 enum Commands {
     /// Benchmarks the performance of programs generated by a local rustc
     BenchRuntimeLocal {
+        /// The path to the local rustc to measure
         rustc: String,
         /// Identifier to associate benchmark results with
         #[clap(long)]
         id: Option<String>,
+
+        /// Exclude all benchmarks matching a prefix in this comma-separated list
+        #[clap(long)]
+        exclude: Option<String>,
+
+        /// Include only benchmarks matching a prefix in this comma-separated list
+        #[clap(long)]
+        include: Option<String>,
     },
     /// Benchmarks a local rustc
     BenchLocal {
@@ -692,8 +701,19 @@ fn main_result() -> anyhow::Result<i32> {
     let target_triple = format!("{}-unknown-linux-gnu", std::env::consts::ARCH);
 
     match args.command {
-        Commands::BenchRuntimeLocal { rustc, id } => {
-            bench_runtime(&rustc, id.as_deref(), runtime_benchmark_dir)?;
+        Commands::BenchRuntimeLocal {
+            rustc,
+            id,
+            exclude,
+            include,
+        } => {
+            bench_runtime(
+                &rustc,
+                id.as_deref(),
+                exclude,
+                include,
+                runtime_benchmark_dir,
+            )?;
             Ok(0)
         }
         Commands::BenchLocal {

collector/src/runtime.rs

@@ -10,30 +10,56 @@ struct BenchmarkBinary {
     path: PathBuf,
 }
 
-pub fn bench_runtime(rustc: &str, id: Option<&str>, benchmark_dir: PathBuf) -> anyhow::Result<()> {
+/// Perform a series of runtime benchmarks using the provided `rustc` compiler.
+/// The runtime benchmarks are looked up in `benchmark_dir`, which is expected to be a path
+/// to a Cargo crate. All binaries built by that crate will are expected to be runtime benchmark
+/// suites that leverage `benchlib`.
+pub fn bench_runtime(
+    rustc: &str,
+    id: Option<&str>,
+    exclude: Option<String>,
+    include: Option<String>,
+    benchmark_dir: PathBuf,
+) -> anyhow::Result<()> {
     let toolchain = get_local_toolchain(&[Profile::Opt], rustc, None, None, id, "")?;
     let output = compile_runtime_benchmarks(&toolchain, &benchmark_dir)?;
     let binaries = gather_binaries(&output)?;
 
     for binary in binaries {
         let name = binary.path.file_name().and_then(|s| s.to_str()).unwrap();
 
-        let data: Vec<BenchmarkResult> = execute_runtime_binary(&binary.path, name)?;
+        let data: Vec<BenchmarkResult> =
+            execute_runtime_binary(&binary.path, name, exclude.as_deref(), include.as_deref())?;
         // TODO: do something with the result
         println!("{name}: {:?}", data);
     }
 
     Ok(())
 }
 
-fn execute_runtime_binary(binary: &Path, name: &str) -> anyhow::Result<Vec<BenchmarkResult>> {
+/// Execute a single runtime benchmark suite defined in a binary crate located in
+/// `runtime-benchmarks`. The binary is expected to use benchlib's `BenchmarkSuite` to execute
+/// a set of runtime benchmarks and return a list of `BenchmarkResult`s encoded as JSON.
+fn execute_runtime_binary(
+    binary: &Path,
+    name: &str,
+    exclude: Option<&str>,
+    include: Option<&str>,
+) -> anyhow::Result<Vec<BenchmarkResult>> {
     // Turn off ASLR
     let mut command = Command::new("setarch");
     command.arg(std::env::consts::ARCH);
     command.arg("-R");
     command.arg(binary);
     command.arg("benchmark");
 
+    if let Some(exclude) = exclude {
+        command.args(&["--exclude", exclude]);
+    }
+    if let Some(include) = include {
+        command.args(&["--include", include]);
+    }
+
     let result = command.output()?;
 
     if !result.status.success() {