Skip to main content

bones_core/
timing.rs

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/// Aggregated timing report across instrumented operations.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct TimingReport {
11    /// Per-operation timing statistics.
12    pub operations: Vec<OpTiming>,
13}
14
15/// Timing statistics for a single named operation.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct OpTiming {
18    /// Human-readable operation name.
19    pub name: String,
20    /// 50th percentile latency.
21    pub p50: Duration,
22    /// 95th percentile latency.
23    pub p95: Duration,
24    /// 99th percentile latency.
25    pub p99: Duration,
26    /// Number of samples collected for this operation.
27    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/// Returns true when `BONES_TIMING` enables timing collection.
43///
44/// Supported truthy values: `1`, `true`, `yes`, `on` (case-insensitive).
45#[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
52/// Enable or disable timing collection.
53pub fn set_timing_enabled(enabled: bool) {
54    TIMING_ENABLED.store(enabled, Ordering::Relaxed);
55    if !enabled {
56        clear_timings();
57    }
58}
59
60/// Returns true when timing collection is currently enabled.
61#[must_use]
62pub fn is_timing_enabled() -> bool {
63    TIMING_ENABLED.load(Ordering::Relaxed)
64}
65
66/// Clears all recorded timings for the current thread.
67pub fn clear_timings() {
68    SAMPLES.with(|samples| samples.borrow_mut().clear());
69}
70
71/// Execute a closure while recording its duration.
72///
73/// Timing is recorded only when enabled via [`set_timing_enabled`].
74pub 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/// Collect all recorded timings from thread-local storage into a report.
86///
87/// This drains the current thread's sample buffer.
88#[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    /// Returns true when no timing samples were recorded.
118    #[must_use]
119    pub const fn is_empty(&self) -> bool {
120        self.operations.is_empty()
121    }
122
123    /// Render the timing report as JSON.
124    #[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    /// Render the timing report as a simple table for terminal output.
144    #[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}