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

rpi_cli/
timings.rs
