Add command for listing runtime benchmarks

Author: Kobzol

collector/benchlib/src/benchmark.rs

@@ -59,12 +59,20 @@ impl BenchmarkSuite {
             Args::Benchmark(args) => {
                 run_benchmark(args, self.benchmarks)?;
             }
+            Args::ListBenchmarks => list_benchmarks(self.benchmarks)?,
         }
 
         Ok(())
     }
 }
 
+fn list_benchmarks(benchmarks: BenchmarkMap) -> anyhow::Result<()> {
+    let benchmark_list: Vec<&str> = benchmarks.into_keys().collect();
+    serde_json::to_writer(std::io::stdout(), &benchmark_list)?;
+
+    Ok(())
+}
+
 fn run_benchmark(args: BenchmarkArgs, benchmarks: BenchmarkMap) -> anyhow::Result<()> {
     let mut items: Vec<(&'static str, BenchmarkWrapper)> = benchmarks
         .into_iter()
@@ -104,7 +112,7 @@ macro_rules! define_benchmark {
 pub use define_benchmark;
 
 /// 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 {
+pub 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),

collector/benchlib/src/cli.rs

@@ -4,6 +4,8 @@ use clap::{FromArgMatches, IntoApp};
 pub enum Args {
     /// Benchmark all benchmarks in this benchmark suite and print the results as JSON.
     Benchmark(BenchmarkArgs),
+    /// List benchmarks that are defined in the current suite as a JSON array.
+    ListBenchmarks,
 }
 
 #[derive(clap::Parser, Debug)]

collector/src/bin/collector.rs

@@ -26,7 +26,7 @@ use tokio::runtime::Runtime;
 
 use collector::execute::bencher::BenchProcessor;
 use collector::execute::profiler::{ProfileProcessor, Profiler};
-use collector::runtime::bench_runtime;
+use collector::runtime::{bench_runtime, BenchmarkFilter};
 use collector::toolchain::{get_local_toolchain, Compiler, Sysroot};
 
 fn n_normal_benchmarks_remaining(n: usize) -> String {
@@ -710,8 +710,7 @@ fn main_result() -> anyhow::Result<i32> {
             bench_runtime(
                 &rustc,
                 id.as_deref(),
-                exclude,
-                include,
+                BenchmarkFilter::new(exclude, include),
                 runtime_benchmark_dir,
             )?;
             Ok(0)

collector/src/runtime.rs

@@ -1,5 +1,6 @@
 use crate::benchmark::profile::Profile;
 use crate::toolchain::{get_local_toolchain, LocalToolchain};
+use benchlib::benchmark::passes_filter;
 use benchlib::messages::BenchmarkResult;
 use cargo_metadata::Message;
 use std::path::{Path, PathBuf};
@@ -8,6 +9,46 @@ use std::process::Command;
 #[derive(Debug)]
 struct BenchmarkBinary {
     path: PathBuf,
+    benchmark_names: Vec<String>,
+}
+
+#[derive(Debug)]
+struct BenchmarkDatabase {
+    binaries: Vec<BenchmarkBinary>,
+}
+
+impl BenchmarkDatabase {
+    fn benchmark_names(&self) -> impl Iterator<Item = &str> {
+        self.binaries
+            .iter()
+            .flat_map(|binary| binary.benchmark_names.iter().map(|n| n.as_ref()))
+    }
+
+    fn total_benchmark_count(&self) -> u64 {
+        self.benchmark_names().count() as u64
+    }
+    fn filtered_benchmark_count(&self, filter: &BenchmarkFilter) -> u64 {
+        self.benchmark_names()
+            .filter(|benchmark| {
+                passes_filter(
+                    &benchmark,
+                    filter.exclude.as_deref(),
+                    filter.include.as_deref(),
+                )
+            })
+            .count() as u64
+    }
+}
+
+pub struct BenchmarkFilter {
+    exclude: Option<String>,
+    include: Option<String>,
+}
+
+impl BenchmarkFilter {
+    pub fn new(exclude: Option<String>, include: Option<String>) -> BenchmarkFilter {
+        Self { exclude, include }
+    }
 }
 
 /// Perform a series of runtime benchmarks using the provided `rustc` compiler.
@@ -17,19 +58,25 @@ struct BenchmarkBinary {
 pub fn bench_runtime(
     rustc: &str,
     id: Option<&str>,
-    exclude: Option<String>,
-    include: Option<String>,
+    filter: BenchmarkFilter,
     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)?;
+    let benchmark_db = discover_benchmarks(&output)?;
+
+    let total_benchmark_count = benchmark_db.total_benchmark_count();
+    let filtered = benchmark_db.filtered_benchmark_count(&filter);
+    println!(
+        "Executing {} benchmarks ({} filtered out)",
+        filtered,
+        total_benchmark_count - filtered
+    );
 
-    for binary in binaries {
+    for binary in benchmark_db.binaries {
         let name = binary.path.file_name().and_then(|s| s.to_str()).unwrap();
 
-        let data: Vec<BenchmarkResult> =
-            execute_runtime_binary(&binary.path, name, exclude.as_deref(), include.as_deref())?;
+        let data: Vec<BenchmarkResult> = execute_runtime_binary(&binary.path, name, &filter)?;
         // TODO: do something with the result
         println!("{name}: {:?}", data);
     }
@@ -43,8 +90,7 @@ pub fn bench_runtime(
 fn execute_runtime_binary(
     binary: &Path,
     name: &str,
-    exclude: Option<&str>,
-    include: Option<&str>,
+    filter: &BenchmarkFilter,
 ) -> anyhow::Result<Vec<BenchmarkResult>> {
     // Turn off ASLR
     let mut command = Command::new("setarch");
@@ -53,10 +99,10 @@ fn execute_runtime_binary(
     command.arg(binary);
     command.arg("benchmark");
 
-    if let Some(exclude) = exclude {
+    if let Some(ref exclude) = filter.exclude {
         command.args(&["--exclude", exclude]);
     }
-    if let Some(include) = include {
+    if let Some(ref include) = filter.include {
         command.args(&["--include", include]);
     }
 
@@ -99,7 +145,8 @@ fn compile_runtime_benchmarks(toolchain: &LocalToolchain, dir: &Path) -> anyhow:
 }
 
 /// Parse Cargo JSON output and find all compiled binaries.
-fn gather_binaries(cargo_stdout: &[u8]) -> anyhow::Result<Vec<BenchmarkBinary>> {
+/// Then execute each benchmark with the `list-benchmarks` command to find out its benchmark names.
+fn discover_benchmarks(cargo_stdout: &[u8]) -> anyhow::Result<BenchmarkDatabase> {
     let mut binaries = vec![];
 
     for message in Message::parse_stream(cargo_stdout) {
@@ -109,7 +156,16 @@ fn gather_binaries(cargo_stdout: &[u8]) -> anyhow::Result<Vec<BenchmarkBinary>>
                 if let Some(ref executable) = artifact.executable {
                     if artifact.target.kind.iter().any(|k| k == "bin") {
                         let path = executable.as_std_path().to_path_buf();
-                        binaries.push(BenchmarkBinary { path });
+                        let benchmarks = gather_benchmarks(&path).map_err(|err| {
+                            anyhow::anyhow!(
+                                "Cannot gather benchmarks from `{}`: {err:?}",
+                                path.display()
+                            )
+                        })?;
+                        binaries.push(BenchmarkBinary {
+                            path,
+                            benchmark_names: benchmarks,
+                        });
                     }
                 }
             }
@@ -119,5 +175,12 @@ fn gather_binaries(cargo_stdout: &[u8]) -> anyhow::Result<Vec<BenchmarkBinary>>
 
     log::debug!("Found binaries: {:?}", binaries);
 
-    Ok(binaries)
+    Ok(BenchmarkDatabase { binaries })
+}
+
+/// Uses the `list-benchmarks` command from `benchlib` to find the benchmark names from the given
+/// benchmark binary.
+fn gather_benchmarks(binary: &Path) -> anyhow::Result<Vec<String>> {
+    let output = Command::new(binary).arg("list-benchmarks").output()?;
+    Ok(serde_json::from_slice(&output.stdout)?)
 }