1use std::cell::RefCell;
2use std::collections::BTreeMap;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::time::{Duration, Instant};
5
6use serde_json::json;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct TimingReport {
11 pub operations: Vec<OpTiming>,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct OpTiming {
18 pub name: String,
20 pub p50: Duration,
22 pub p95: Duration,
24 pub p99: Duration,
26 pub count: usize,
28}
29
30#[derive(Debug, Clone)]
31struct Sample {
32 name: String,
33 elapsed: Duration,
34}
35
36thread_local! {
37 static SAMPLES: RefCell<Vec<Sample>> = const { RefCell::new(Vec::new()) };
38}
39
40static TIMING_ENABLED: AtomicBool = AtomicBool::new(false);
41
42#[must_use]
46pub fn timing_enabled_from_env() -> bool {
47 std::env::var("BONES_TIMING")
48 .ok()
49 .is_some_and(|value| is_truthy(value.as_str()))
50}
51
52pub fn set_timing_enabled(enabled: bool) {
54 TIMING_ENABLED.store(enabled, Ordering::Relaxed);
55 if !enabled {
56 clear_timings();
57 }
58}
59
60#[must_use]
62pub fn is_timing_enabled() -> bool {
63 TIMING_ENABLED.load(Ordering::Relaxed)
64}
65
66pub fn clear_timings() {
68 SAMPLES.with(|samples| samples.borrow_mut().clear());
69}
70
71pub fn timed<R>(name: &str, f: impl FnOnce() -> R) -> R {
75 if !is_timing_enabled() {
76 return f();
77 }
78
79 let started = Instant::now();
80 let result = f();
81 record_sample(name, started.elapsed());
82 result
83}
84
85#[must_use]
89pub fn collect_report() -> TimingReport {
90 let samples = SAMPLES.with(|samples| std::mem::take(&mut *samples.borrow_mut()));
91
92 let mut grouped: BTreeMap<String, Vec<Duration>> = BTreeMap::new();
93 for sample in samples {
94 grouped.entry(sample.name).or_default().push(sample.elapsed);
95 }
96
97 let operations = grouped
98 .into_iter()
99 .map(|(name, mut values)| {
100 values.sort_unstable();
101 let count = values.len();
102
103 OpTiming {
104 name,
105 p50: percentile(&values, 50),
106 p95: percentile(&values, 95),
107 p99: percentile(&values, 99),
108 count,
109 }
110 })
111 .collect();
112
113 TimingReport { operations }
114}
115
116impl TimingReport {
117 #[must_use]
119 pub const fn is_empty(&self) -> bool {
120 self.operations.is_empty()
121 }
122
123 #[must_use]
125 pub fn to_json(&self) -> serde_json::Value {
126 let operations = self
127 .operations
128 .iter()
129 .map(|op| {
130 json!({
131 "name": op.name,
132 "count": op.count,
133 "p50_us": op.p50.as_micros(),
134 "p95_us": op.p95.as_micros(),
135 "p99_us": op.p99.as_micros(),
136 })
137 })
138 .collect::<Vec<_>>();
139
140 json!({ "operations": operations })
141 }
142
143 #[must_use]
145 pub fn display_table(&self) -> String {
146 if self.operations.is_empty() {
147 return "No timing samples recorded.".to_string();
148 }
149
150 let mut out = String::new();
151 out.push_str("operation count p50 p95 p99\n");
152 out.push_str("--------------------------------------------------------------\n");
153
154 for op in &self.operations {
155 use std::fmt::Write;
156 let _ = writeln!(
157 out,
158 "{:<28} {:>6} {:>8} {:>8} {:>8}",
159 op.name,
160 op.count,
161 format_duration(op.p50),
162 format_duration(op.p95),
163 format_duration(op.p99)
164 );
165 }
166
167 out
168 }
169}
170
171fn record_sample(name: &str, elapsed: Duration) {
172 SAMPLES.with(|samples| {
173 samples.borrow_mut().push(Sample {
174 name: name.to_string(),
175 elapsed,
176 });
177 });
178}
179
180fn percentile(sorted: &[Duration], pct: u32) -> Duration {
181 if sorted.is_empty() {
182 return Duration::ZERO;
183 }
184
185 let pct_usize = usize::try_from(pct).unwrap_or(100).min(100);
186 let rank = pct_usize.saturating_mul(sorted.len()).saturating_add(99) / 100;
187 let index = rank.saturating_sub(1).min(sorted.len().saturating_sub(1));
188
189 sorted[index]
190}
191
192fn format_duration(duration: Duration) -> String {
193 let micros = duration.as_micros();
194
195 if micros >= 1_000_000 {
196 let secs = micros / 1_000_000;
197 let millis = (micros % 1_000_000) / 1_000;
198 format!("{secs}.{millis:03}s")
199 } else if micros >= 1_000 {
200 let millis = micros / 1_000;
201 let rem = micros % 1_000;
202 format!("{millis}.{rem:03}ms")
203 } else {
204 format!("{micros}µs")
205 }
206}
207
208const fn is_truthy(value: &str) -> bool {
209 value.eq_ignore_ascii_case("1")
210 || value.eq_ignore_ascii_case("true")
211 || value.eq_ignore_ascii_case("yes")
212 || value.eq_ignore_ascii_case("on")
213}
214
215#[cfg(test)]
216mod tests {
217 use super::*;
218
219 static TEST_GUARD: std::sync::Mutex<()> = std::sync::Mutex::new(());
220
221 #[test]
222 fn timed_does_not_record_when_disabled() {
223 let _guard = TEST_GUARD.lock().expect("test guard lock");
224 set_timing_enabled(false);
225 clear_timings();
226
227 let value = timed("disabled", || 7_u8);
228 assert_eq!(value, 7);
229
230 let report = collect_report();
231 assert!(report.is_empty());
232 }
233
234 #[test]
235 fn timed_records_when_enabled() {
236 let _guard = TEST_GUARD.lock().expect("test guard lock");
237 set_timing_enabled(true);
238 clear_timings();
239
240 let value = timed("enabled", || 42_u8);
241 assert_eq!(value, 42);
242
243 let report = collect_report();
244 assert_eq!(report.operations.len(), 1);
245 assert_eq!(report.operations[0].name, "enabled");
246 assert_eq!(report.operations[0].count, 1);
247 assert!(report.operations[0].p50 > Duration::ZERO);
248
249 set_timing_enabled(false);
250 }
251
252 #[test]
253 fn collect_report_groups_and_sorts_operations() {
254 let _guard = TEST_GUARD.lock().expect("test guard lock");
255 clear_timings();
256
257 record_sample("query", Duration::from_micros(3_000));
258 record_sample("query", Duration::from_micros(1_000));
259 record_sample("query", Duration::from_micros(2_000));
260 record_sample("replay", Duration::from_micros(5_000));
261
262 let report = collect_report();
263 assert_eq!(report.operations.len(), 2);
264
265 let query = report
266 .operations
267 .iter()
268 .find(|op| op.name == "query")
269 .expect("query timing should exist");
270 assert_eq!(query.count, 3);
271 assert_eq!(query.p50, Duration::from_micros(2_000));
272 assert_eq!(query.p95, Duration::from_micros(3_000));
273 assert_eq!(query.p99, Duration::from_micros(3_000));
274 }
275
276 #[test]
277 fn truthy_parser_is_case_insensitive() {
278 let _guard = TEST_GUARD.lock().expect("test guard lock");
279
280 assert!(is_truthy("TrUe"));
281 assert!(is_truthy("1"));
282 assert!(is_truthy("YES"));
283 assert!(is_truthy("on"));
284 assert!(!is_truthy("0"));
285 assert!(!is_truthy("false"));
286 }
287
288 #[test]
289 fn display_table_and_json_have_expected_fields() {
290 let _guard = TEST_GUARD.lock().expect("test guard lock");
291 clear_timings();
292
293 record_sample("project", Duration::from_micros(1_500));
294
295 let report = collect_report();
296 let table = report.display_table();
297 assert!(table.contains("operation"));
298 assert!(table.contains("project"));
299
300 let json = report.to_json();
301 let operations = json
302 .get("operations")
303 .and_then(serde_json::Value::as_array)
304 .expect("operations array should exist");
305 assert_eq!(operations.len(), 1);
306
307 let op = &operations[0];
308 assert_eq!(
309 op.get("name"),
310 Some(&serde_json::Value::String("project".to_string()))
311 );
312 assert_eq!(op.get("count"), Some(&serde_json::Value::from(1)));
313 assert!(op.get("p50_us").is_some());
314 }
315}