custom metrics formatting

Cargo.toml

@@ -50,6 +50,7 @@ table_reporter = ["prettytable-rs"]
 branch_predictor = ["bpu_trasher"]
 default = ["branch_predictor"]
 
+
 [[bench]]
 name = "bench"
 harness = false

src/bench.rs

@@ -1,12 +1,6 @@
 use std::sync::atomic;
 
-use crate::{
-    bench_id::BenchId,
-    black_box,
-    output_value::OutputValue,
-    plugins::{alloc::*, *},
-    stats::*,
-};
+use crate::{bench_id::BenchId, black_box, output_value::OutputValue, plugins::*, stats::*};
 use quanta::Clock;
 
 /// The trait which typically wraps a InputWithBenchmark and allows to hide the generics.
@@ -65,8 +59,8 @@ pub struct BenchResult {
     pub input_size_in_bytes: Option<usize>,
     /// The size of the output returned by the bench. Enables reporting.
     pub output_value: Option<String>,
-    /// Memory tracking is enabled and the peak memory consumption is reported.
-    pub tracked_memory: bool,
+    /// Formatted custom metrics for reporting
+    pub formatted_custom_metrics: Vec<(&'static str, String)>,
 }
 
 /// Bundle of input and benchmark for running benchmarks
@@ -123,11 +117,8 @@ impl<'a, I, O: OutputValue> Bench<'a> for InputWithBenchmark<'a, I, O> {
     fn get_results(&mut self, plugins: &mut PluginManager) -> BenchResult {
         let num_iter = self.get_num_iter_or_fail();
         let total_num_iter = self.bench.num_group_iter as u64 * num_iter as u64;
-        let memory_consumption: Option<&Vec<usize>> = plugins
-            .downcast_plugin::<PeakMemAllocPlugin>(ALLOC_EVENT_LISTENER_NAME)
-            .and_then(|counters| counters.get_by_bench_id(&self.bench.bench_id));
-        let stats = compute_stats(&self.results, memory_consumption);
-        let tracked_memory = memory_consumption.is_some();
+        let custom_metrics = plugins.get_custom_metrics(&self.bench.bench_id);
+        let stats = compute_stats(&self.results, custom_metrics);
 
         let perf_counter = get_perf_counter(plugins, &self.bench.bench_id, total_num_iter);
         let output_value = (self.bench.fun)(self.input);
@@ -136,7 +127,7 @@ impl<'a, I, O: OutputValue> Bench<'a> for InputWithBenchmark<'a, I, O> {
             stats,
             perf_counter,
             input_size_in_bytes: self.input_size_in_bytes,
-            tracked_memory,
+            formatted_custom_metrics: Vec::new(),
             output_value: output_value.format(),
             old_stats: None,
             old_perf_counter: None,

src/plugins/alloc.rs

@@ -1,4 +1,5 @@
 use std::any::Any;
+use yansi::Paint;
 
 use peakmem_alloc::PeakMemAllocTrait;
 
@@ -55,4 +56,50 @@ impl EventListener for PeakMemAllocPlugin {
             _ => {}
         }
     }
+
+    fn custom_metrics(&self, bench_id: &BenchId, metrics: &mut Vec<(&'static str, f64)>) {
+        if let Some(perf) = self.get_by_bench_id(bench_id) {
+            let total_memory: usize = perf.iter().copied().sum();
+            let avg_memory = if perf.is_empty() {
+                0
+            } else {
+                total_memory / perf.len()
+            };
+            metrics.push(("Memory", avg_memory as f64));
+        }
+    }
+
+    fn custom_metric_keys(&self) -> &[&'static str] {
+        &["Memory"]
+    }
+
+    fn format_custom_metrics(
+        &self,
+        stats: &BenchStats,
+        other: Option<&crate::stats::BenchStats>,
+    ) -> Vec<(&'static str, String)> {
+        let avg_memory = stats
+            .custom_metrics
+            .iter()
+            .find(|(k, _)| *k == "Memory")
+            .map(|(_, v)| *v)
+            .unwrap_or(0.0) as u64;
+        let mem_diff = crate::stats::compute_diff(stats, None, other, |stats| {
+            stats
+                .custom_metrics
+                .iter()
+                .find(|(k, _)| *k == "Memory")
+                .map(|(_, v)| *v)
+                .unwrap_or(0.0) as u64
+        });
+
+        let s = format!(
+            "Memory: {} {}",
+            crate::report::format::bytes_to_string(avg_memory)
+                .bright_cyan()
+                .bold(),
+            mem_diff,
+        );
+        vec![("Memory", s)]
+    }
 }

src/plugins/events.rs

@@ -7,7 +7,7 @@
 //! Any type that implements the [EventListener] trait can be added to [PluginManager].
 //!
 
-use crate::{bench::BenchResult, bench_id::BenchId};
+use crate::{bench::BenchResult, bench_id::BenchId, stats::BenchStats};
 use std::any::Any;
 
 /// Events that can be emitted by the benchmark runner.
@@ -88,6 +88,24 @@ pub trait EventListener: Any {
     fn on_event(&mut self, event: PluginEvents);
     /// Downcast the listener to `Any`.
     fn as_any(&mut self) -> &mut dyn Any;
+
+    /// Append custom metrics to be reported in the stats.
+    /// This is called once per benchmark run.
+    fn custom_metrics(&self, _bench_id: &BenchId, _metrics: &mut Vec<(&'static str, f64)>) {}
+
+    /// Returns a list of metric keys this plugin will report.
+    fn custom_metric_keys(&self) -> &[&'static str] {
+        &[]
+    }
+
+    /// Formats the custom metrics. Returns a list of (key, formatted_string).
+    fn format_custom_metrics(
+        &self,
+        _stats: &BenchStats,
+        _other: Option<&BenchStats>,
+    ) -> Vec<(&'static str, String)> {
+        Vec::new()
+    }
 }
 
 /// [PluginManager] is responsible for managing plugins and emitting events.
@@ -160,6 +178,35 @@ impl PluginManager {
             }
         }
     }
+
+    /// Collect all custom metrics from plugins for a specific benchmark.
+    pub fn get_custom_metrics(&self, bench_id: &BenchId) -> Vec<(&'static str, f64)> {
+        let capacity: usize = self
+            .listeners
+            .iter()
+            .map(|(_, l)| l.custom_metric_keys().len())
+            .sum();
+        let mut metrics = Vec::with_capacity(capacity);
+        for (_listener_name, listener) in self.listeners.iter() {
+            listener.custom_metrics(bench_id, &mut metrics);
+        }
+        metrics
+    }
+
+    /// Ask plugins to format their custom metrics for reporting.
+    pub fn format_custom_metrics(
+        &self,
+        stats: &BenchStats,
+        other: Option<&BenchStats>,
+    ) -> Vec<(&'static str, String)> {
+        let mut formatted = Vec::new();
+        for (_, listener) in self.listeners.iter() {
+            for (k, v) in listener.format_custom_metrics(stats, other) {
+                formatted.push((k, v));
+            }
+        }
+        formatted
+    }
 }
 
 impl Default for PluginManager {

src/report/mod.rs

@@ -35,7 +35,7 @@ pub use table_reporter::TableReporter;
 
 use yansi::Paint;
 
-use format::{bytes_to_string, format_duration_or_throughput};
+use format::format_duration_or_throughput;
 
 use crate::{
     bench::Bench,
@@ -65,6 +65,12 @@ pub(crate) fn report_group<'a>(
         fetch_previous_run_and_write_results_to_disk(&mut result);
         results.push(result);
     }
+
+    for result in results.iter_mut() {
+        result.formatted_custom_metrics =
+            events.format_custom_metrics(&result.stats, result.old_stats.as_ref());
+    }
+
     events.emit(PluginEvents::GroupStop {
         runner_name,
         group_name,
@@ -76,7 +82,7 @@ pub(crate) fn report_group<'a>(
 pub(crate) fn avg_median_str(
     stats: &BenchStats,
     input_size_in_bytes: Option<usize>,
-    other: Option<BenchStats>,
+    other: Option<&BenchStats>,
 ) -> (String, String) {
     let avg_ns_diff = compute_diff(stats, input_size_in_bytes, other, |stats| stats.average_ns);
     let median_ns_diff = compute_diff(stats, input_size_in_bytes, other, |stats| stats.median_ns);
@@ -111,24 +117,6 @@ pub(crate) fn min_max_str(stats: &BenchStats, input_size_in_bytes: Option<usize>
     }
 }
 
-pub(crate) fn memory_str(
-    stats: &BenchStats,
-    other: Option<BenchStats>,
-    report_memory: bool,
-) -> String {
-    let mem_diff = compute_diff(stats, None, other, |stats| stats.avg_memory as u64);
-    if !report_memory {
-        return "".to_string();
-    }
-    format!(
-        "Memory: {} {}",
-        bytes_to_string(stats.avg_memory as u64)
-            .bright_cyan()
-            .bold(),
-        mem_diff,
-    )
-}
-
 use std::{
     ops::Deref,
     sync::{Arc, Once},

src/report/plain_reporter.rs

@@ -2,7 +2,7 @@ use std::any::Any;
 
 use yansi::Paint;
 
-use super::{BenchStats, REPORTER_PLUGIN_NAME, avg_median_str, memory_str, min_max_str};
+use super::{BenchStats, REPORTER_PLUGIN_NAME, avg_median_str, min_max_str};
 use crate::{
     plugins::{EventListener, PluginEvents},
     report::{PrintOnce, check_and_print},
@@ -60,11 +60,11 @@ impl EventListener for PlainReporter {
                     let perf_counter = &result.perf_counter;
 
                     let mut stats_columns = self.to_columns(
-                        result.stats,
-                        result.old_stats,
+                        &result.stats,
+                        result.old_stats.as_ref(),
                         result.input_size_in_bytes,
                         &result.output_value,
-                        result.tracked_memory,
+                        &result.formatted_custom_metrics,
                         output_value_column_title,
                     );
                     stats_columns.insert(0, result.bench_id.bench_name.to_string());
@@ -99,34 +99,35 @@ impl PlainReporter {
 
     pub(crate) fn to_columns(
         &self,
-        stats: BenchStats,
-        other: Option<BenchStats>,
+        stats: &BenchStats,
+        other: Option<&BenchStats>,
         input_size_in_bytes: Option<usize>,
         output_value: &Option<String>,
-        report_memory: bool,
+        formatted_custom_metrics: &Vec<(&'static str, String)>,
         output_value_column_title: &'static str,
     ) -> Vec<String> {
-        let (avg_str, median_str) = avg_median_str(&stats, input_size_in_bytes, other);
+        let (avg_str, median_str) = avg_median_str(stats, input_size_in_bytes, other);
         let avg_str = format!("Avg: {}", avg_str);
         let median_str = format!("Median: {}", median_str);
 
-        let min_max = min_max_str(&stats, input_size_in_bytes);
-        let memory_string = memory_str(&stats, other, report_memory);
+        let min_max = min_max_str(stats, input_size_in_bytes);
+
+        let mut cols = Vec::new();
+        for (_, s) in formatted_custom_metrics {
+            cols.push(s.clone());
+        }
+        cols.push(avg_str);
+        cols.push(median_str);
+        cols.push(min_max);
+
         if let Some(output_value) = output_value {
-            vec![
-                memory_string,
-                avg_str,
-                median_str,
-                min_max,
-                format!(
-                    "{}: {}",
-                    output_value_column_title,
-                    output_value.to_string()
-                ),
-            ]
-        } else {
-            vec![memory_string, avg_str, median_str, min_max]
+            cols.push(format!(
+                "{}: {}",
+                output_value_column_title,
+                output_value.to_string()
+            ));
         }
+        cols
     }
 
     fn print_table(&self, table_data: &Vec<Vec<String>>) {

src/report/table_reporter.rs

@@ -2,7 +2,7 @@ use std::any::Any;
 
 use yansi::Paint;
 
-use super::{REPORTER_PLUGIN_NAME, avg_median_str, memory_str, min_max_str};
+use super::{REPORTER_PLUGIN_NAME, avg_median_str, min_max_str};
 use crate::{
     plugins::{EventListener, PluginEvents},
     report::{PrintOnce, check_and_print},
@@ -81,37 +81,41 @@ impl EventListener for TableReporter {
                     .build();
                 table.set_format(format);
 
-                let mut row = prettytable::row!["Name", "Memory", "Avg", "Median", "Min .. Max"];
-                if !results[0].tracked_memory {
-                    row.remove_cell(1);
+                let mut headers = vec![Cell::new("Name")];
+                for (key, _) in &results[0].formatted_custom_metrics {
+                    headers.push(Cell::new(*key));
                 }
+                headers.push(Cell::new("Avg"));
+                headers.push(Cell::new("Median"));
+                headers.push(Cell::new("Min .. Max"));
+
                 let has_output_value = results.iter().any(|r| r.output_value.is_some());
                 if has_output_value {
-                    row.add_cell(Cell::new(output_value_column_title));
+                    headers.push(Cell::new(output_value_column_title));
                 }
-                table.set_titles(row);
+                table.set_titles(Row::new(headers));
                 for result in results {
-                    let (avg_str, median_str) =
-                        avg_median_str(&result.stats, result.input_size_in_bytes, result.old_stats);
+                    let (avg_str, median_str) = avg_median_str(
+                        &result.stats,
+                        result.input_size_in_bytes,
+                        result.old_stats.as_ref(),
+                    );
                     let min_max = min_max_str(&result.stats, result.input_size_in_bytes);
-                    let memory_string =
-                        memory_str(&result.stats, result.old_stats, result.tracked_memory);
-                    let mut row = Row::new(vec![
-                        Cell::new(&result.bench_id.bench_name),
-                        Cell::new(&memory_string),
-                        Cell::new(&avg_str),
-                        Cell::new(&median_str),
-                        Cell::new(&min_max),
-                    ]);
+
+                    let mut row = vec![Cell::new(&result.bench_id.bench_name)];
+                    for (_, formatted) in &result.formatted_custom_metrics {
+                        row.push(Cell::new(formatted));
+                    }
+                    row.push(Cell::new(&avg_str));
+                    row.push(Cell::new(&median_str));
+                    row.push(Cell::new(&min_max));
+
                     if has_output_value {
-                        row.add_cell(Cell::new(
+                        row.push(Cell::new(
                             result.output_value.as_ref().unwrap_or(&"".to_string()),
                         ));
                     }
-                    if !result.tracked_memory {
-                        row.remove_cell(1);
-                    }
-                    table.add_row(row);
+                    table.add_row(Row::new(row));
                 }
                 table.printstd();
             }