diff --git a/Cargo.toml b/Cargo.toml index 9977d56..7ac9e06 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,7 @@ table_reporter = ["prettytable-rs"] branch_predictor = ["bpu_trasher"] default = ["branch_predictor"] + [[bench]] name = "bench" harness = false diff --git a/src/bench.rs b/src/bench.rs index 4f5dc28..3f24769 100644 --- a/src/bench.rs +++ b/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, /// The size of the output returned by the bench. Enables reporting. pub output_value: Option, - /// 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> = plugins - .downcast_plugin::(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, diff --git a/src/plugins/alloc.rs b/src/plugins/alloc.rs index bfcffaf..b36817a 100644 --- a/src/plugins/alloc.rs +++ b/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)] + } } diff --git a/src/plugins/events.rs b/src/plugins/events.rs index 5bbe78a..2771b82 100644 --- a/src/plugins/events.rs +++ b/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 { diff --git a/src/report/mod.rs b/src/report/mod.rs index 56bc2f2..776723c 100644 --- a/src/report/mod.rs +++ b/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, - other: Option, + 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 } } -pub(crate) fn memory_str( - stats: &BenchStats, - other: Option, - 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}, diff --git a/src/report/plain_reporter.rs b/src/report/plain_reporter.rs index 763f188..0283921 100644 --- a/src/report/plain_reporter.rs +++ b/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, + stats: &BenchStats, + other: Option<&BenchStats>, input_size_in_bytes: Option, output_value: &Option, - report_memory: bool, + formatted_custom_metrics: &Vec<(&'static str, String)>, output_value_column_title: &'static str, ) -> Vec { - 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>) { diff --git a/src/report/table_reporter.rs b/src/report/table_reporter.rs index fa9460e..f2ffaa3 100644 --- a/src/report/table_reporter.rs +++ b/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(); } diff --git a/src/stats.rs b/src/stats.rs index 720058a..20eb31c 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -6,7 +6,7 @@ use yansi::Paint; /// including timing and memory usage. /// /// The data is already aggregated. -#[derive(Serialize, Deserialize, Debug, Clone, Copy)] +#[derive(Serialize, Deserialize, Debug, Clone)] pub struct BenchStats { /// The minimum time taken for an operation, in nanoseconds. pub min_ns: u64, @@ -20,19 +20,18 @@ pub struct BenchStats { /// The median time taken for an operation, in nanoseconds. pub median_ns: u64, - /// The average memory used during the operation, in bytes. - pub avg_memory: usize, + /// Custom metrics collected by plugins, e.g., "memory_consumption". + pub custom_metrics: Vec<(String, f64)>, } /// Compute diff from two values of BenchStats pub fn compute_diff u64>( stats: &BenchStats, input_size_in_bytes: Option, - other: Option, + other: Option<&BenchStats>, f: F, ) -> String { other - .as_ref() .map(|other| { if f(other) == 0 || f(stats) == 0 || f(other) == f(stats) { return "".to_string(); @@ -88,16 +87,8 @@ pub fn format_percentage(diff: f64, smaller_is_better: bool) -> String { } pub fn compute_stats( results: &[RunResult], - memory_consumption: Option<&Vec>, + custom_metrics: Vec<(&'static str, f64)>, ) -> BenchStats { - // Avg memory consumption - let avg_memory = memory_consumption - .map(|memory_consumption| { - let total_memory: usize = memory_consumption.iter().copied().sum(); - total_memory / memory_consumption.len() - }) - .unwrap_or(0); - let mut sorted_results: Vec = results.iter().map(|res| res.duration_ns).collect(); sorted_results.sort(); @@ -123,7 +114,10 @@ pub fn compute_stats( max_ns, average_ns, median_ns, - avg_memory, + custom_metrics: custom_metrics + .into_iter() + .map(|(k, v)| (k.to_string(), v)) + .collect(), } } @@ -141,7 +135,7 @@ mod tests { #[test] fn test_compute_stats_median_odd() { let results = vec![create_res(10), create_res(20), create_res(30)]; - let stats = compute_stats(&results, None); + let stats = compute_stats(&results, Default::default()); assert_eq!( stats.median_ns, 20, "Median should be the middle element for odd count" @@ -156,7 +150,7 @@ mod tests { create_res(30), create_res(40), ]; - let stats = compute_stats(&results, None); + let stats = compute_stats(&results, Default::default()); assert_eq!( stats.median_ns, 25, "Median should be the average of the two middle elements for even count" @@ -170,7 +164,7 @@ mod tests { max_ns: 0, average_ns: 150, median_ns: 0, - avg_memory: 24, + custom_metrics: Default::default(), }; let other_stats = BenchStats { @@ -178,11 +172,11 @@ mod tests { max_ns: 0, average_ns: 100, // different average_ns to see the difference in the output median_ns: 0, - avg_memory: 0, + custom_metrics: Default::default(), }; // Example usage: Using average_ns field for comparison. - let diff = compute_diff(&stats, Some(1000), Some(other_stats), |x| x.average_ns); + let diff = compute_diff(&stats, Some(1000), Some(&other_stats), |x| x.average_ns); // Check the output assert_eq!(diff, "(-33.33%)".red().to_string()); diff --git a/src/write_results.rs b/src/write_results.rs index d5dde70..6aa31ce 100644 --- a/src/write_results.rs +++ b/src/write_results.rs @@ -29,12 +29,15 @@ pub fn fetch_previous_run_and_write_results_to_disk(result: &mut BenchResult) { let filepath = get_bench_file(result); // Check if file exists and deserialize if filepath.exists() { - let content = std::fs::read_to_string(&filepath).unwrap(); - let lines: Vec<_> = content.lines().collect(); - result.old_stats = miniserde::json::from_str(lines[0]).unwrap(); - result.old_perf_counter = lines - .get(1) - .and_then(|line| miniserde::json::from_str(line).ok()); + if let Ok(content) = std::fs::read_to_string(&filepath) { + let lines: Vec<_> = content.lines().collect(); + if !lines.is_empty() { + result.old_stats = miniserde::json::from_str(lines[0]).ok(); + result.old_perf_counter = lines + .get(1) + .and_then(|line| miniserde::json::from_str(line).ok()); + } + } } let perf_counter = &result.perf_counter;