1use std::sync::OnceLock;
12use std::time::{Duration, Instant};
13
14static TIMING_ENABLED: OnceLock<bool> = OnceLock::new();
16
17#[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
27const MAX_QUERY_ID_LENGTH: usize = 60;
29
30pub 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 #[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 #[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 #[inline]
63 pub fn start_phase(&self) -> Instant {
64 Instant::now()
65 }
66
67 #[inline]
69 pub fn end_phase(&mut self, name: &str, since: Instant) {
70 self.phases.push((name.to_string(), since.elapsed()));
71 }
72
73 #[inline]
76 pub fn add_accumulated_phase(&mut self, name: &str, duration: Duration) {
77 self.phases.push((name.to_string(), duration));
78 }
79
80 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 #[test]
106 fn test_query_timings_struct_fields() {
107 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 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}