1use std::path::Path;
5use std::sync::{Mutex, OnceLock};
6use std::time::{Duration, Instant};
7
8pub 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 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
75pub 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
86pub 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#[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 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
143pub 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 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}