Skip to main content

ite_cli/
profile.rs

1//! Lightweight span profiler. Enabled by setting `ITE_PROFILE=<output-path>`;
2//! when unset, spans cost one branch. The dump is a plain-text table.
3
4use std::path::Path;
5use std::sync::{Mutex, OnceLock};
6use std::time::{Duration, Instant};
7
8/// Collects labeled durations and renders a summary table.
9pub struct Registry {
10    records: Mutex<Vec<(&'static str, Duration)>>,
11}
12
13impl Registry {
14    pub const fn new() -> Self {
15        Self {
16            records: Mutex::new(Vec::new()),
17        }
18    }
19
20    pub fn record(&self, label: &'static str, duration: Duration) {
21        self.records.lock().unwrap().push((label, duration));
22    }
23
24    /// A table of per-label stats, ordered by total time descending.
25    pub fn summary(&self) -> String {
26        let records = self.records.lock().unwrap();
27        if records.is_empty() {
28            return "(no spans recorded)\n".to_string();
29        }
30        let mut by_label: Vec<(&'static str, Vec<Duration>)> = Vec::new();
31        for &(label, duration) in records.iter() {
32            match by_label.iter_mut().find(|(l, _)| *l == label) {
33                Some((_, durations)) => durations.push(duration),
34                None => by_label.push((label, vec![duration])),
35            }
36        }
37        let mut rows: Vec<(&'static str, Stats)> = by_label
38            .into_iter()
39            .map(|(label, durations)| (label, Stats::from_durations(&durations).unwrap()))
40            .collect();
41        rows.sort_by_key(|(_, s)| std::cmp::Reverse(s.total));
42
43        let mut out = format!(
44            "{:<24} {:>7} {:>9} {:>9} {:>9} {:>9} {:>9}\n",
45            "span", "count", "total", "mean", "p50", "p95", "max"
46        );
47        for (label, s) in rows {
48            out.push_str(&format!(
49                "{:<24} {:>7} {:>9} {:>9} {:>9} {:>9} {:>9}\n",
50                label,
51                s.count,
52                format_duration(s.total),
53                format_duration(s.mean),
54                format_duration(s.p50),
55                format_duration(s.p95),
56                format_duration(s.max),
57            ));
58        }
59        out
60    }
61
62    pub fn write_to(&self, path: &Path) -> std::io::Result<()> {
63        std::fs::write(path, self.summary())
64    }
65}
66
67impl Default for Registry {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73pub static GLOBAL: Registry = Registry::new();
74
75/// The output path from `ITE_PROFILE`, if profiling is enabled.
76pub fn output_path() -> Option<&'static str> {
77    static PATH: OnceLock<Option<String>> = OnceLock::new();
78    PATH.get_or_init(|| std::env::var("ITE_PROFILE").ok().filter(|p| !p.is_empty()))
79        .as_deref()
80}
81
82pub fn enabled() -> bool {
83    output_path().is_some()
84}
85
86/// Times a scope and records it in [`GLOBAL`] on drop (no-op when disabled).
87pub struct Span {
88    label: &'static str,
89    start: Option<Instant>,
90}
91
92#[must_use]
93pub fn span(label: &'static str) -> Span {
94    Span {
95        label,
96        start: enabled().then(Instant::now),
97    }
98}
99
100impl Drop for Span {
101    fn drop(&mut self) {
102        if let Some(start) = self.start {
103            GLOBAL.record(self.label, start.elapsed());
104        }
105    }
106}
107
108/// Summary statistics over a set of durations.
109#[derive(Clone, Copy, Debug, PartialEq, Eq)]
110pub struct Stats {
111    pub count: usize,
112    pub total: Duration,
113    pub mean: Duration,
114    pub p50: Duration,
115    pub p95: Duration,
116    pub max: Duration,
117}
118
119impl Stats {
120    /// `None` when `durations` is empty.
121    pub fn from_durations(durations: &[Duration]) -> Option<Self> {
122        if durations.is_empty() {
123            return None;
124        }
125        let mut sorted = durations.to_vec();
126        sorted.sort();
127        let percentile = |q: f64| {
128            let index = ((sorted.len() - 1) as f64 * q).round() as usize;
129            sorted[index]
130        };
131        let total: Duration = sorted.iter().sum();
132        Some(Self {
133            count: sorted.len(),
134            total,
135            mean: total / sorted.len() as u32,
136            p50: percentile(0.50),
137            p95: percentile(0.95),
138            max: *sorted.last().unwrap(),
139        })
140    }
141}
142
143/// Adaptive human-readable duration: ns, µs, ms, or s.
144pub fn format_duration(d: Duration) -> String {
145    let nanos = d.as_nanos();
146    if nanos < 1_000 {
147        format!("{nanos}ns")
148    } else if nanos < 1_000_000 {
149        format!("{:.1}µs", nanos as f64 / 1_000.0)
150    } else if nanos < 1_000_000_000 {
151        format!("{:.2}ms", nanos as f64 / 1_000_000.0)
152    } else {
153        format!("{:.2}s", nanos as f64 / 1_000_000_000.0)
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    fn ms(n: u64) -> Duration {
162        Duration::from_millis(n)
163    }
164
165    #[test]
166    fn stats_of_known_durations() {
167        let durations: Vec<Duration> = (1..=10).map(ms).collect();
168        let stats = Stats::from_durations(&durations).unwrap();
169        assert_eq!(stats.count, 10);
170        assert_eq!(stats.total, ms(55));
171        assert_eq!(stats.mean, Duration::from_micros(5500));
172        assert_eq!(stats.max, ms(10));
173        // Percentile index = round((n-1) * q) over the sorted set.
174        assert_eq!(stats.p50, ms(6));
175        assert_eq!(stats.p95, ms(10));
176    }
177
178    #[test]
179    fn stats_of_single_duration() {
180        let stats = Stats::from_durations(&[ms(10)]).unwrap();
181        assert_eq!(stats.count, 1);
182        assert_eq!(stats.mean, ms(10));
183        assert_eq!(stats.p50, ms(10));
184        assert_eq!(stats.p95, ms(10));
185        assert_eq!(stats.max, ms(10));
186    }
187
188    #[test]
189    fn stats_of_nothing_is_none() {
190        assert_eq!(Stats::from_durations(&[]), None);
191    }
192
193    #[test]
194    fn stats_do_not_require_sorted_input() {
195        let stats = Stats::from_durations(&[ms(9), ms(1), ms(5)]).unwrap();
196        assert_eq!(stats.p50, ms(5));
197        assert_eq!(stats.max, ms(9));
198    }
199
200    #[test]
201    fn summary_orders_labels_by_total_descending() {
202        let reg = Registry::new();
203        reg.record("cheap", ms(2));
204        reg.record("expensive", ms(50));
205        reg.record("cheap", ms(3));
206        let summary = reg.summary();
207        let expensive_at = summary.find("expensive").unwrap();
208        let cheap_at = summary.find("cheap").unwrap();
209        assert!(expensive_at < cheap_at, "summary:\n{summary}");
210        assert!(summary.contains("count"), "has a header:\n{summary}");
211    }
212
213    #[test]
214    fn summary_reports_counts() {
215        let reg = Registry::new();
216        reg.record("draw", ms(1));
217        reg.record("draw", ms(1));
218        reg.record("draw", ms(1));
219        assert!(reg.summary().contains('3'));
220    }
221
222    #[test]
223    fn empty_summary_says_so() {
224        assert_eq!(Registry::new().summary(), "(no spans recorded)\n");
225    }
226
227    #[test]
228    fn write_to_writes_the_summary() {
229        let dir = tempfile::tempdir().unwrap();
230        let path = dir.path().join("profile.txt");
231        let reg = Registry::new();
232        reg.record("scan", ms(7));
233        reg.write_to(&path).unwrap();
234        assert_eq!(std::fs::read_to_string(&path).unwrap(), reg.summary());
235    }
236
237    #[test]
238    fn durations_format_adaptively() {
239        assert_eq!(format_duration(Duration::from_nanos(250)), "250ns");
240        assert_eq!(format_duration(Duration::from_nanos(12_500)), "12.5µs");
241        assert_eq!(format_duration(Duration::from_micros(3_400)), "3.40ms");
242        assert_eq!(format_duration(Duration::from_millis(2_500)), "2.50s");
243    }
244}