use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct StageRecord {
pub name: String,
pub duration: Duration,
pub element_count: usize,
}
#[derive(Debug, Default)]
pub struct PipelineTimer {
records: Vec<StageRecord>,
current_start: Option<(String, Instant)>,
pipeline_start: Option<Instant>,
}
impl PipelineTimer {
pub fn new() -> Self {
Self {
records: Vec::new(),
current_start: None,
pipeline_start: Some(Instant::now()),
}
}
pub fn start_stage(&mut self, name: &str) {
if let Some((prev_name, prev_start)) = self.current_start.take() {
self.records.push(StageRecord {
name: prev_name,
duration: prev_start.elapsed(),
element_count: 0,
});
}
self.current_start = Some((name.to_string(), Instant::now()));
}
pub fn end_stage(&mut self, element_count: usize) {
if let Some((name, start)) = self.current_start.take() {
self.records.push(StageRecord {
name,
duration: start.elapsed(),
element_count,
});
}
}
pub fn total_duration(&self) -> Duration {
self.pipeline_start.map(|s| s.elapsed()).unwrap_or_default()
}
pub fn records(&self) -> &[StageRecord] {
&self.records
}
pub fn summary(&self) -> String {
let mut out = String::from("Pipeline Timing Summary\n");
out.push_str(&format!(
"{:<40} {:>10} {:>10}\n",
"Stage", "Time (ms)", "Elements"
));
out.push_str(&"-".repeat(62));
out.push('\n');
for r in &self.records {
out.push_str(&format!(
"{:<40} {:>10.2} {:>10}\n",
r.name,
r.duration.as_secs_f64() * 1000.0,
r.element_count,
));
}
out.push_str(&"-".repeat(62));
out.push('\n');
out.push_str(&format!(
"{:<40} {:>10.2}\n",
"TOTAL",
self.total_duration().as_secs_f64() * 1000.0
));
out
}
pub fn log_summary(&self) {
for line in self.summary().lines() {
log::info!("{}", line);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
#[test]
fn test_stage_recording() {
let mut timer = PipelineTimer::new();
timer.start_stage("Stage A");
thread::sleep(Duration::from_millis(5));
timer.end_stage(100);
timer.start_stage("Stage B");
timer.end_stage(200);
assert_eq!(timer.records().len(), 2);
assert_eq!(timer.records()[0].name, "Stage A");
assert_eq!(timer.records()[0].element_count, 100);
assert_eq!(timer.records()[1].name, "Stage B");
assert_eq!(timer.records()[1].element_count, 200);
assert!(timer.records()[0].duration.as_nanos() > 0);
}
#[test]
fn test_auto_close_previous_stage() {
let mut timer = PipelineTimer::new();
timer.start_stage("Stage 1");
timer.start_stage("Stage 2");
timer.end_stage(50);
assert_eq!(timer.records().len(), 2);
assert_eq!(timer.records()[0].name, "Stage 1");
assert_eq!(timer.records()[0].element_count, 0); assert_eq!(timer.records()[1].name, "Stage 2");
}
#[test]
fn test_summary_format() {
let mut timer = PipelineTimer::new();
timer.start_stage("Test Stage");
timer.end_stage(42);
let summary = timer.summary();
assert!(summary.contains("Test Stage"));
assert!(summary.contains("42"));
assert!(summary.contains("TOTAL"));
}
#[test]
fn test_empty_timer() {
let timer = PipelineTimer::new();
assert!(timer.records().is_empty());
let summary = timer.summary();
assert!(summary.contains("TOTAL"));
}
}