Skip to main content

apiary_query/
timing.rs

1//! Structured timing instrumentation for query execution.
2//!
3//! Enabled via the `APIARY_TIMING=1` environment variable. When unset,
4//! all timing operations are no-ops with zero overhead.
5//!
6//! After each query completes, a single parseable line is emitted to stderr:
7//! ```text
8//! [TIMING] query=Q1.1 total=5044ms parse=2ms plan=15ms ...
9//! ```
10
11use std::sync::OnceLock;
12use std::time::{Duration, Instant};
13
14/// Cached result of the `APIARY_TIMING` environment variable check.
15static TIMING_ENABLED: OnceLock<bool> = OnceLock::new();
16
17/// Returns `true` when `APIARY_TIMING=1` is set.
18#[inline]
19pub fn timing_enabled() -> bool {
20    *TIMING_ENABLED.get_or_init(|| {
21        std::env::var("APIARY_TIMING")
22            .map(|v| v == "1")
23            .unwrap_or(false)
24    })
25}
26
27/// Maximum number of characters from the SQL query used as the timing query ID.
28const MAX_QUERY_ID_LENGTH: usize = 60;
29
30/// Collected timing phases for a single query.
31pub struct QueryTimings {
32    pub query_id: String,
33    pub phases: Vec<(String, Duration)>,
34    pub total: Duration,
35    start: Instant,
36}
37
38impl QueryTimings {
39    /// Create a new `QueryTimings` and start the total wall-clock timer.
40    /// Returns `None` when timing is disabled (zero overhead path).
41    #[inline]
42    pub fn begin(query_id: impl Into<String>) -> Option<Self> {
43        if !timing_enabled() {
44            return None;
45        }
46        Some(Self {
47            query_id: query_id.into(),
48            phases: Vec::new(),
49            total: Duration::ZERO,
50            start: Instant::now(),
51        })
52    }
53
54    /// Create a `QueryTimings` using the first [`MAX_QUERY_ID_LENGTH`] chars
55    /// of the SQL string as the query identifier.
56    #[inline]
57    pub fn begin_from_sql(sql: &str) -> Option<Self> {
58        Self::begin(sql.chars().take(MAX_QUERY_ID_LENGTH).collect::<String>())
59    }
60
61    /// Start timing a named phase. Returns the `Instant` to pass to [`end_phase`].
62    #[inline]
63    pub fn start_phase(&self) -> Instant {
64        Instant::now()
65    }
66
67    /// Record the elapsed time for a phase started with [`start_phase`].
68    #[inline]
69    pub fn end_phase(&mut self, name: &str, since: Instant) {
70        self.phases.push((name.to_string(), since.elapsed()));
71    }
72
73    /// Record a pre-computed duration for a phase that was accumulated
74    /// across multiple iterations (e.g., per-table I/O totals).
75    #[inline]
76    pub fn add_accumulated_phase(&mut self, name: &str, duration: Duration) {
77        self.phases.push((name.to_string(), duration));
78    }
79
80    /// Finish timing and emit the results to stderr.
81    pub fn finish(mut self) {
82        self.total = self.start.elapsed();
83
84        let mut parts: Vec<String> = Vec::with_capacity(self.phases.len() + 2);
85        parts.push(format!("query={}", self.query_id));
86        parts.push(format!("total={}ms", self.total.as_millis()));
87        for (name, dur) in &self.phases {
88            parts.push(format!("{}={}ms", name, dur.as_millis()));
89        }
90
91        eprintln!("[TIMING] {}", parts.join(" "));
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    // Note: Because OnceLock caches the value once, the tests below that
100    // toggle the env var operate on the *constructor* returning None/Some
101    // based on the first call in the process. In a test binary the init
102    // happens once; to exercise both paths we call `timing_enabled()`
103    // directly and test the struct API independently.
104
105    #[test]
106    fn test_query_timings_struct_fields() {
107        // Exercise the struct API directly (no env dependency).
108        let mut t = QueryTimings {
109            query_id: "Q1.1".to_string(),
110            phases: Vec::new(),
111            total: Duration::ZERO,
112            start: Instant::now(),
113        };
114
115        let s1 = t.start_phase();
116        t.end_phase("parse", s1);
117        let s2 = t.start_phase();
118        t.end_phase("plan", s2);
119        t.add_accumulated_phase("data_read", Duration::from_millis(42));
120
121        assert_eq!(t.phases.len(), 3);
122        assert_eq!(t.phases[0].0, "parse");
123        assert_eq!(t.phases[1].0, "plan");
124        assert_eq!(t.phases[2].0, "data_read");
125        assert_eq!(t.phases[2].1, Duration::from_millis(42));
126
127        t.finish();
128    }
129
130    #[test]
131    fn test_begin_from_sql_truncates() {
132        // Test the truncation logic directly on a long string.
133        let long_sql = "A".repeat(200);
134        let truncated: String = long_sql.chars().take(MAX_QUERY_ID_LENGTH).collect();
135        assert_eq!(truncated.len(), MAX_QUERY_ID_LENGTH);
136    }
137
138    #[test]
139    fn test_max_query_id_length_constant() {
140        assert_eq!(MAX_QUERY_ID_LENGTH, 60);
141    }
142}