1use std::sync::{LazyLock, Mutex};
2use std::time::Instant;
3
4const DEFAULT_BUDGET_MS: u64 = 5000;
5const WARN_THRESHOLD_PCT: f64 = 80.0;
6const MAX_LATENCY_SAMPLES: usize = 4096;
7
8#[derive(Debug, Clone, Copy)]
9pub struct LatencyBudget {
10 pub total_ms: u64,
11 pub started_at: Instant,
12}
13
14#[derive(Debug, Clone)]
15pub struct LatencyCheckpoint {
16 pub phase: String,
17 pub elapsed_ms: u64,
18 pub budget_remaining_ms: u64,
19 pub over_budget: bool,
20}
21
22#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
23pub struct LatencyReport {
24 pub total_elapsed_ms: u64,
25 pub budget_ms: u64,
26 pub within_budget: bool,
27 pub phases: Vec<PhaseMetric>,
28 pub slowest_phase: Option<String>,
29}
30
31#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
32pub struct PhaseMetric {
33 pub name: String,
34 pub elapsed_ms: u64,
35 pub pct_of_total: f64,
36}
37
38pub struct LatencyTracker {
39 budget: LatencyBudget,
40 phases: Vec<(String, u64)>,
41 last_checkpoint: Instant,
42}
43
44#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
45pub struct LatencySummary {
46 pub sample_count: usize,
47 pub p50_ms: u64,
48 pub p95_ms: u64,
49 pub p99_ms: u64,
50 pub max_ms: u64,
51}
52
53static LATENCY_SAMPLES: LazyLock<Mutex<Vec<u64>>> = LazyLock::new(|| Mutex::new(Vec::new()));
54
55impl LatencyBudget {
56 pub fn new(total_ms: u64) -> Self {
57 Self {
58 total_ms,
59 started_at: Instant::now(),
60 }
61 }
62
63 pub fn default_budget() -> Self {
64 Self::new(DEFAULT_BUDGET_MS)
65 }
66
67 pub fn elapsed_ms(&self) -> u64 {
68 u64::try_from(self.started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
69 }
70
71 pub fn remaining_ms(&self) -> u64 {
72 self.total_ms.saturating_sub(self.elapsed_ms())
73 }
74
75 pub fn is_over_budget(&self) -> bool {
76 self.elapsed_ms() > self.total_ms
77 }
78
79 pub fn is_warning(&self) -> bool {
80 percentage(self.elapsed_ms(), self.total_ms) >= WARN_THRESHOLD_PCT
81 }
82}
83
84impl LatencyTracker {
85 pub fn new(budget_ms: u64) -> Self {
86 let now = Instant::now();
87 Self {
88 budget: LatencyBudget {
89 total_ms: budget_ms,
90 started_at: now,
91 },
92 phases: Vec::new(),
93 last_checkpoint: now,
94 }
95 }
96
97 pub fn checkpoint(&mut self, phase: &str) -> LatencyCheckpoint {
98 let now = Instant::now();
99 let phase_elapsed = duration_ms(self.last_checkpoint, now);
100 self.phases.push((phase.to_owned(), phase_elapsed));
101 self.last_checkpoint = now;
102
103 LatencyCheckpoint {
104 phase: phase.to_owned(),
105 elapsed_ms: phase_elapsed,
106 budget_remaining_ms: self.budget.remaining_ms(),
107 over_budget: self.budget.is_over_budget(),
108 }
109 }
110
111 pub fn report(&self) -> LatencyReport {
112 let total_elapsed_ms = self.budget.elapsed_ms();
113 let phases = self
114 .phases
115 .iter()
116 .map(|(name, elapsed_ms)| PhaseMetric {
117 name: name.clone(),
118 elapsed_ms: *elapsed_ms,
119 pct_of_total: percentage(*elapsed_ms, total_elapsed_ms),
120 })
121 .collect();
122 let slowest_phase = self
123 .phases
124 .iter()
125 .max_by_key(|(_, elapsed_ms)| elapsed_ms)
126 .map(|(name, _)| name.clone());
127
128 LatencyReport {
129 total_elapsed_ms,
130 budget_ms: self.budget.total_ms,
131 within_budget: total_elapsed_ms <= self.budget.total_ms,
132 phases,
133 slowest_phase,
134 }
135 }
136}
137
138pub fn format_latency_report(report: &LatencyReport) -> String {
139 let status = if report.within_budget { "✓" } else { "✗" };
140 let mut output = format!(
141 "Latency: {}ms / {}ms budget ({:.1}%) {status}",
142 report.total_elapsed_ms,
143 report.budget_ms,
144 percentage(report.total_elapsed_ms, report.budget_ms)
145 );
146 let name_width = report
147 .phases
148 .iter()
149 .map(|phase| phase.name.len())
150 .max()
151 .unwrap_or(0);
152
153 for phase in &report.phases {
154 output.push_str(&format!(
155 "\n {:name_width$}: {}ms ({:.1}%)",
156 phase.name, phase.elapsed_ms, phase.pct_of_total
157 ));
158 }
159 output
160}
161
162pub fn record_latency(elapsed_ms: u64) {
163 let mut samples = LATENCY_SAMPLES
164 .lock()
165 .unwrap_or_else(std::sync::PoisonError::into_inner);
166 if samples.len() == MAX_LATENCY_SAMPLES {
167 samples.remove(0);
168 }
169 samples.push(elapsed_ms);
170}
171
172pub fn latency_percentile(percentile: f64) -> Option<u64> {
173 let samples = LATENCY_SAMPLES
174 .lock()
175 .unwrap_or_else(std::sync::PoisonError::into_inner);
176 percentile_from_samples(&samples, percentile)
177}
178
179pub fn latency_summary() -> Option<LatencySummary> {
180 let samples = LATENCY_SAMPLES
181 .lock()
182 .unwrap_or_else(std::sync::PoisonError::into_inner);
183 if samples.is_empty() {
184 return None;
185 }
186
187 Some(LatencySummary {
188 sample_count: samples.len(),
189 p50_ms: percentile_from_samples(&samples, 50.0)?,
190 p95_ms: percentile_from_samples(&samples, 95.0)?,
191 p99_ms: percentile_from_samples(&samples, 99.0)?,
192 max_ms: samples.iter().copied().max()?,
193 })
194}
195
196fn duration_ms(start: Instant, end: Instant) -> u64 {
197 u64::try_from(end.duration_since(start).as_millis()).unwrap_or(u64::MAX)
198}
199
200fn percentage(value: u64, total: u64) -> f64 {
201 if total == 0 {
202 0.0
203 } else {
204 value as f64 * 100.0 / total as f64
205 }
206}
207
208fn percentile_from_samples(samples: &[u64], percentile: f64) -> Option<u64> {
209 if samples.is_empty() {
210 return None;
211 }
212 let mut sorted = samples.to_vec();
213 sorted.sort_unstable();
214 let percentile = if percentile.is_finite() {
215 percentile.clamp(0.0, 100.0)
216 } else {
217 0.0
218 };
219 let rank = (percentile / 100.0 * sorted.len() as f64).ceil() as usize;
220 sorted.get(rank.saturating_sub(1)).copied()
221}
222
223#[cfg(test)]
224mod tests {
225 use super::{
226 LATENCY_SAMPLES, LatencyBudget, LatencyReport, LatencyTracker, PhaseMetric,
227 format_latency_report, latency_percentile, latency_summary, record_latency,
228 };
229 use std::sync::Mutex;
230 use std::thread;
231 use std::time::Duration;
232
233 static SAMPLE_TEST_LOCK: Mutex<()> = Mutex::new(());
234
235 fn clear_samples() {
236 LATENCY_SAMPLES
237 .lock()
238 .unwrap_or_else(std::sync::PoisonError::into_inner)
239 .clear();
240 }
241
242 #[test]
243 fn test_budget_creation() {
244 let budget = LatencyBudget::new(1234);
245 assert_eq!(budget.total_ms, 1234);
246 }
247
248 #[test]
249 fn test_default_budget_creation() {
250 assert_eq!(LatencyBudget::default_budget().total_ms, 5000);
251 }
252
253 #[test]
254 fn test_budget_remaining_decreases() {
255 let budget = LatencyBudget::new(10_000);
256 thread::sleep(Duration::from_millis(10));
257 assert!(budget.remaining_ms() < budget.total_ms);
258 }
259
260 #[test]
261 fn test_budget_over_budget() {
262 let budget = LatencyBudget::new(1);
263 thread::sleep(Duration::from_millis(10));
264 assert!(budget.is_over_budget());
265 }
266
267 #[test]
268 fn test_budget_warning_threshold() {
269 let budget = LatencyBudget::new(12);
270 thread::sleep(Duration::from_millis(10));
271 assert!(budget.is_warning());
272 }
273
274 #[test]
275 fn test_tracker_checkpoint_records_phase() {
276 let mut tracker = LatencyTracker::new(10_000);
277 thread::sleep(Duration::from_millis(10));
278 let checkpoint = tracker.checkpoint("compress");
279 assert_eq!(checkpoint.phase, "compress");
280 assert_eq!(tracker.phases.len(), 1);
281 assert!(checkpoint.elapsed_ms >= 10);
282 }
283
284 #[test]
285 fn test_tracker_multiple_phases() {
286 let mut tracker = LatencyTracker::new(10_000);
287 for phase in ["compress", "upstream", "shape"] {
288 thread::sleep(Duration::from_millis(10));
289 tracker.checkpoint(phase);
290 }
291 assert_eq!(tracker.phases.len(), 3);
292 }
293
294 #[test]
295 fn test_tracker_report_within_budget() {
296 let tracker = LatencyTracker::new(10_000);
297 assert!(tracker.report().within_budget);
298 }
299
300 #[test]
301 fn test_tracker_report_slowest_phase() {
302 let mut tracker = LatencyTracker::new(10_000);
303 tracker.phases = vec![("compress".to_owned(), 5), ("upstream".to_owned(), 20)];
304 assert_eq!(tracker.report().slowest_phase.as_deref(), Some("upstream"));
305 }
306
307 #[test]
308 fn test_format_report_output() {
309 let report = LatencyReport {
310 total_elapsed_ms: 234,
311 budget_ms: 5000,
312 within_budget: true,
313 phases: vec![PhaseMetric {
314 name: "compress".to_owned(),
315 elapsed_ms: 45,
316 pct_of_total: 19.2,
317 }],
318 slowest_phase: Some("compress".to_owned()),
319 };
320 let output = format_latency_report(&report);
321 assert!(output.contains("234ms / 5000ms budget"));
322 assert!(output.contains("compress: 45ms (19.2%)"));
323 }
324
325 #[test]
326 fn test_percentile_calculation() {
327 let _guard = SAMPLE_TEST_LOCK
328 .lock()
329 .unwrap_or_else(std::sync::PoisonError::into_inner);
330 clear_samples();
331 for sample in 1..=100 {
332 record_latency(sample);
333 }
334 assert_eq!(latency_percentile(50.0), Some(50));
335 assert_eq!(latency_percentile(95.0), Some(95));
336 assert_eq!(latency_percentile(99.0), Some(99));
337 clear_samples();
338 }
339
340 #[test]
341 fn test_latency_summary_empty() {
342 let _guard = SAMPLE_TEST_LOCK
343 .lock()
344 .unwrap_or_else(std::sync::PoisonError::into_inner);
345 clear_samples();
346 assert!(latency_summary().is_none());
347 }
348
349 #[test]
350 fn test_record_latency_bounded() {
351 let _guard = SAMPLE_TEST_LOCK
352 .lock()
353 .unwrap_or_else(std::sync::PoisonError::into_inner);
354 clear_samples();
355 for sample in 0..5000 {
356 record_latency(sample);
357 }
358 let summary = latency_summary().expect("samples should produce a summary");
359 assert_eq!(summary.sample_count, 4096);
360 assert_eq!(summary.max_ms, 4999);
361 clear_samples();
362 }
363}