Merge pull request #1463 from rust-lang/runtime-cli-progress

Improve runtime benchmark CLI output

Author: Kobzol

collector/benchlib/src/benchmark.rs

@@ -1,46 +1,43 @@
 use crate::cli::{parse_cli, Args, BenchmarkArgs};
+use crate::comm::messages::{BenchmarkMessage, BenchmarkResult, BenchmarkStats};
+use crate::comm::output_message;
 use crate::measure::benchmark_function;
-use crate::messages::BenchmarkResult;
 use crate::process::raise_process_priority;
-use log::LevelFilter;
 use std::collections::HashMap;
 
-/// Create a new benchmark suite. Use the closure argument to define benchmarks.
-pub fn benchmark_suite<F: FnOnce(&mut BenchmarkSuite)>(define_func: F) {
-    env_logger::Builder::from_default_env()
-        .filter_level(LevelFilter::Info)
-        .init();
-    let mut suite = BenchmarkSuite::new();
-    define_func(&mut suite);
-    suite.run().expect("Benchmark suite has failed");
+/// Create a new benchmark group. Use the closure argument to define individual benchmarks.
+pub fn run_benchmark_group<F: FnOnce(&mut BenchmarkGroup)>(define_func: F) {
+    env_logger::init();
+
+    let mut group = BenchmarkGroup::new();
+    define_func(&mut group);
+    group.run().expect("Benchmark group execution has failed");
 }
 
-/// Type-erased function that performs a benchmark.
+/// Type-erased function that executes a single benchmark.
 struct BenchmarkWrapper {
-    func: Box<dyn Fn() -> anyhow::Result<BenchmarkResult>>,
+    func: Box<dyn Fn() -> anyhow::Result<BenchmarkStats>>,
 }
 
-type BenchmarkMap = HashMap<&'static str, BenchmarkWrapper>;
-
 #[derive(Default)]
-pub struct BenchmarkSuite {
-    benchmarks: BenchmarkMap,
+pub struct BenchmarkGroup {
+    benchmarks: HashMap<&'static str, BenchmarkWrapper>,
 }
 
-impl BenchmarkSuite {
+impl BenchmarkGroup {
     pub fn new() -> Self {
         Self::default()
     }
 
     /// Registers a single benchmark.
-    /// `func` should return a closure that will be benchmarked.
+    /// `constructor` should return a closure that will be benchmarked.
     pub fn register<F: Fn() -> Bench + Clone + 'static, R, Bench: FnOnce() -> R + 'static>(
         &mut self,
         name: &'static str,
         constructor: F,
     ) {
         // We want to type-erase the target `func` by wrapping it in a Box.
-        let benchmark_func = Box::new(move || benchmark_function(name, constructor.clone()));
+        let benchmark_func = Box::new(move || benchmark_function(constructor.clone()));
         let benchmark_def = BenchmarkWrapper {
             func: benchmark_func,
         };
@@ -49,62 +46,81 @@ impl BenchmarkSuite {
         }
     }
 
-    /// Execute the benchmark suite. It will parse CLI arguments and decide what to do based on
+    /// Execute the benchmark group. It will parse CLI arguments and decide what to do based on
     /// them.
     pub fn run(self) -> anyhow::Result<()> {
         raise_process_priority();
 
         let args = parse_cli()?;
         match args {
-            Args::Benchmark(args) => {
-                run_benchmark(args, self.benchmarks)?;
+            Args::Run(args) => {
+                self.run_benchmarks(args)?;
             }
+            Args::List => self.list_benchmarks()?,
         }
 
         Ok(())
     }
-}
 
-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}) `{name}` completed: {result:?}");
-            results.push(result);
+    fn run_benchmarks(self, args: BenchmarkArgs) -> anyhow::Result<()> {
+        let mut items: Vec<(&'static str, BenchmarkWrapper)> = self
+            .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 stdout = std::io::stdout().lock();
+
+        for (name, def) in items {
+            let mut stats: Vec<BenchmarkStats> = Vec::with_capacity(args.iterations as usize);
+            for i in 0..args.iterations {
+                let benchmark_stats = (def.func)()?;
+                log::info!("Benchmark (run {i}) `{name}` completed: {benchmark_stats:?}");
+                stats.push(benchmark_stats);
+            }
+            output_message(
+                &mut stdout,
+                BenchmarkMessage::Result(BenchmarkResult {
+                    name: name.to_string(),
+                    stats,
+                }),
+            )?;
         }
+
+        Ok(())
     }
 
-    println!("{}", serde_json::to_string(&results)?);
-    Ok(())
+    fn list_benchmarks(self) -> anyhow::Result<()> {
+        let benchmark_list: Vec<&str> = self.benchmarks.into_keys().collect();
+        serde_json::to_writer(std::io::stdout(), &benchmark_list)?;
+
+        Ok(())
+    }
 }
 
-/// Adds a single benchmark to the benchmark suite.
+/// Adds a single benchmark to the benchmark group.
 /// ```ignore
 /// use benchlib::define_benchmark;
 ///
-/// define_benchmark!(suite, my_bench, {
+/// define_benchmark!(group, my_bench, {
 ///     || do_something()
 /// });
 /// ```
 #[macro_export]
 macro_rules! define_benchmark {
-    ($suite:expr, $name:ident, $fun:expr) => {
+    ($group:expr, $name:ident, $fun:expr) => {
         let func = move || $fun;
-        $suite.register(stringify!($name), func);
+        $group.register(stringify!($name), func);
     };
 }
 
 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

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

collector/benchlib/src/comm/messages.rs

@@ -0,0 +1,27 @@
+//! This module defines messages that are exchanged between a binary that uses `benchlib` and
+//! the `collector` crate.
+
+use std::time::Duration;
+
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+pub enum BenchmarkMessage {
+    Result(BenchmarkResult),
+}
+
+/// Stats gathered by several iterations of a single benchmark.
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+pub struct BenchmarkResult {
+    pub name: String,
+    pub stats: Vec<BenchmarkStats>,
+}
+
+/// The stats gathered by a single benchmark execution.
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+pub struct BenchmarkStats {
+    pub cycles: u64,
+    pub instructions: u64,
+    pub branch_misses: u64,
+    pub cache_misses: u64,
+    pub cache_references: u64,
+    pub wall_time: Duration,
+}

collector/benchlib/src/comm/mod.rs

@@ -0,0 +1,40 @@
+use crate::comm::messages::BenchmarkMessage;
+use std::io::{BufRead, BufReader, Read, Write};
+
+pub mod messages;
+
+/// Messages are communicated as line-delimited JSON.
+pub fn output_message<W: Write>(mut sink: W, message: BenchmarkMessage) -> anyhow::Result<()> {
+    serde_json::to_writer(&mut sink, &message)?;
+    sink.write_all(b"\n")?;
+    Ok(())
+}
+
+pub struct MessageReader<R> {
+    inner: BufReader<R>,
+    line: String,
+}
+
+impl<R: Read> MessageReader<R> {
+    pub fn new(inner: R) -> Self {
+        Self {
+            inner: BufReader::new(inner),
+            line: Default::default(),
+        }
+    }
+}
+
+impl<R: Read> Iterator for MessageReader<R> {
+    type Item = anyhow::Result<BenchmarkMessage>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.inner.read_line(&mut self.line) {
+            Ok(0) => None,
+            Ok(_) => match serde_json::from_str(&self.line) {
+                Ok(value) => Some(Ok(value)),
+                Err(error) => Some(Err(error.into())),
+            },
+            Err(error) => Some(Err(error.into())),
+        }
+    }
+}

collector/benchlib/src/lib.rs

@@ -15,6 +15,6 @@
 
 pub mod benchmark;
 mod cli;
+pub mod comm;
 pub mod measure;
-pub mod messages;
 pub mod process;

collector/benchlib/src/measure/perf_counter/unix.rs

@@ -1,5 +1,5 @@
 use crate::benchmark::black_box;
-use crate::messages::BenchmarkResult;
+use crate::comm::messages::BenchmarkStats;
 use perf_event::events::Hardware;
 use perf_event::{Builder, Counter, Group};
 use std::time::Instant;
@@ -16,9 +16,8 @@ struct Counters {
 /// The function is executed twice, once to gather wall-time measurement and the second time to
 /// gather perf. counters.
 pub fn benchmark_function<F: Fn() -> Bench + 'static, R, Bench: FnOnce() -> R + 'static>(
-    name: &'static str,
     benchmark_constructor: F,
-) -> anyhow::Result<BenchmarkResult> {
+) -> anyhow::Result<BenchmarkStats> {
     let mut group = create_group()?;
     let counters = prepare_counters(&mut group)?;
 
@@ -48,8 +47,7 @@ pub fn benchmark_function<F: Fn() -> Bench + 'static, R, Bench: FnOnce() -> R +
     // Try to avoid optimizing the result out.
     black_box(output);
 
-    let result = BenchmarkResult {
-        name: String::from(name),
+    let result = BenchmarkStats {
         cycles: measurement[&counters.cycles],
         instructions: measurement[&counters.instructions],
         branch_misses: measurement[&counters.branch_misses],

collector/benchlib/src/measure/perf_counter/windows.rs

@@ -1,8 +1,5 @@
-use crate::messages::BenchmarkResult;
+use crate::comm::messages::BenchmarkStats;
 
-pub fn benchmark_function<F: FnOnce() -> R, R>(
-    _name: &'static str,
-    _func: F,
-) -> anyhow::Result<BenchmarkResult> {
+pub fn benchmark_function<F: FnOnce() -> R, R>(_func: F) -> anyhow::Result<BenchmarkStats> {
     panic!("Runtime benchmarking is not supported on Windows");
 }

collector/benchlib/src/messages.rs

@@ -2,7 +2,7 @@ use std::io::{BufRead, BufReader, Write};
 
 use snap::{read::FrameDecoder, write::FrameEncoder};
 
-use benchlib::benchmark::{benchmark_suite, black_box};
+use benchlib::benchmark::{black_box, run_benchmark_group};
 use benchlib::define_benchmark;
 
 const BYTES: usize = 64 * 1024 * 1024;
@@ -11,8 +11,8 @@ fn main() {
     // Inspired by https://github.com/rust-lang/rust/issues/102727
     // The pattern we want is a BufReader which wraps a Read impl where one Read::read call will
     // never fill the whole BufReader buffer.
-    benchmark_suite(|suite| {
-        define_benchmark!(suite, bufreader_snappy, {
+    run_benchmark_group(|group| {
+        define_benchmark!(group, bufreader_snappy, {
             let data = vec![0u8; BYTES];
             move || {
                 let mut compressed = Vec::new();

collector/runtime-benchmarks/bufreader/src/main.rs

@@ -1,11 +1,11 @@
 use benchlib;
-use benchlib::benchmark::benchmark_suite;
+use benchlib::benchmark::run_benchmark_group;
 use benchlib::define_benchmark;
 
 fn main() {
-    benchmark_suite(|suite| {
+    run_benchmark_group(|group| {
         // Measures how long does it take to insert 10 thousand numbers into a `hashbrown` hashmap.
-        define_benchmark!(suite, hashmap_insert_10k, {
+        define_benchmark!(group, hashmap_insert_10k, {
             let mut map = hashbrown::HashMap::with_capacity_and_hasher(
                 10000,
                 fxhash::FxBuildHasher::default(),