Skip to main content

rpi_cli/
timings.rs

1//! Central timing instrumentation for startup profiling.
2//!
3//! Port of native Pi's `packages/coding-agent/src/core/timings.ts`. Enable with
4//! the `PI_TIMING=1` environment variable; disabled builds pay only a single
5//! atomic load per call.
6//!
7//! Timings are grouped by namespace (`main` for the primary startup path,
8//! `extensions` for JS/Rust extension loading) and flushed once at shutdown.
9
10use std::collections::HashMap;
11use std::sync::{Mutex, OnceLock};
12use std::time::Instant;
13
14/// True when `PI_TIMING=1`. Evaluated once.
15pub fn enabled() -> bool {
16    static ENABLED: OnceLock<bool> = OnceLock::new();
17    *ENABLED.get_or_init(|| {
18        std::env::var("PI_TIMING")
19            .map(|v| v == "1")
20            .unwrap_or(false)
21    })
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25pub enum TimingNamespace {
26    Main,
27    Extensions,
28}
29
30impl TimingNamespace {
31    fn label(self) -> &'static str {
32        match self {
33            Self::Main => "main",
34            Self::Extensions => "extensions",
35        }
36    }
37}
38
39struct Group {
40    timings: Vec<(String, i64)>,
41    last: Instant,
42}
43
44fn groups() -> &'static Mutex<HashMap<&'static str, Group>> {
45    static GROUPS: OnceLock<Mutex<HashMap<&'static str, Group>>> = OnceLock::new();
46    GROUPS.get_or_init(|| Mutex::new(HashMap::new()))
47}
48
49/// Reset a namespace's timings and mark "now" as its reference point.
50pub fn reset(namespace: TimingNamespace) {
51    if !enabled() {
52        return;
53    }
54    if let Ok(mut g) = groups().lock() {
55        g.insert(
56            namespace.label(),
57            Group {
58                timings: Vec::new(),
59                last: Instant::now(),
60            },
61        );
62    }
63}
64
65/// Record the elapsed time since the previous `time`/`reset` in this namespace.
66pub fn time(label: impl Into<String>, namespace: TimingNamespace) {
67    if !enabled() {
68        return;
69    }
70    let now = Instant::now();
71    if let Ok(mut g) = groups().lock() {
72        let group = g.entry(namespace.label()).or_insert_with(|| Group {
73            timings: Vec::new(),
74            last: now,
75        });
76        let ms = now.duration_since(group.last).as_millis() as i64;
77        group.timings.push((label.into(), ms));
78        group.last = now;
79    }
80}
81
82/// Print every accumulated namespace to stderr. No-op unless `PI_TIMING=1`.
83pub fn print_timings() {
84    if !enabled() {
85        return;
86    }
87    let Ok(g) = groups().lock() else {
88        return;
89    };
90    // Stable order: main before extensions.
91    for ns in ["main", "extensions"] {
92        let Some(group) = g.get(ns) else { continue };
93        print_group(&format!("Startup Timings: {ns}"), &group.timings);
94    }
95}
96
97fn print_group(title: &str, timings: &[(String, i64)]) {
98    let printable: Vec<_> = timings.iter().filter(|(_, ms)| *ms >= 0).collect();
99    if printable.is_empty() {
100        return;
101    }
102    eprintln!("\n--- {title} ---");
103    for (label, ms) in &printable {
104        eprintln!("  {label}: {ms}ms");
105    }
106    let total: i64 = printable.iter().map(|(_, ms)| *ms).sum();
107    eprintln!("  TOTAL: {total}ms");
108    eprintln!("{}\n", "-".repeat(title.len() + 8));
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn disabled_by_default() {
117        // The env var is unset in the test process unless explicitly set.
118        if std::env::var_os("PI_TIMING").is_none() {
119            assert!(!enabled());
120        }
121    }
122
123    #[test]
124    fn records_and_resets() {
125        // Exercise the recording path directly, independent of the env flag.
126        if let Ok(mut g) = groups().lock() {
127            g.insert(
128                "test",
129                Group {
130                    timings: Vec::new(),
131                    last: Instant::now(),
132                },
133            );
134        }
135        if let Ok(mut g) = groups().lock() {
136            let group = g.get_mut("test").unwrap();
137            group.timings.push(("boot".to_string(), 12));
138        }
139        if let Ok(g) = groups().lock() {
140            assert_eq!(g.get("test").unwrap().timings.len(), 1);
141        }
142    }
143}