use std::time::{Duration, Instant};
pub struct LatencyTracker {
samples: Vec<Duration>,
sample_rate: usize,
}
impl LatencyTracker {
pub fn new(rate: usize) -> Self {
Self {
samples: Vec::new(),
sample_rate: rate.max(1),
}
}
pub fn record<F, R>(&mut self, iter_index: usize, f: F) -> R
where
F: FnOnce() -> R,
{
if iter_index % self.sample_rate == 0 {
let start = Instant::now();
let r = f();
self.samples.push(start.elapsed());
r
} else {
f()
}
}
pub fn samples_count(&self) -> usize {
self.samples.len()
}
pub fn into_samples(self) -> Vec<Duration> {
self.samples
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LatencyStats {
pub p50: Duration,
pub p95: Duration,
pub p99: Duration,
pub samples_count: usize,
}
impl LatencyStats {
pub fn from_samples(mut samples: Vec<Duration>) -> Self {
let n = samples.len();
if n == 0 {
return Self {
p50: Duration::ZERO,
p95: Duration::ZERO,
p99: Duration::ZERO,
samples_count: 0,
};
}
samples.sort();
let p50 = samples[n / 2];
let p95 = samples[((n as f64 * 0.95).floor() as usize).min(n - 1)];
let p99 = samples[((n as f64 * 0.99).floor() as usize).min(n - 1)];
Self {
p50,
p95,
p99,
samples_count: n,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rate_one_records_every_iter() {
let mut t = LatencyTracker::new(1);
for i in 0..10 {
t.record(i, || std::hint::black_box(i));
}
assert_eq!(t.samples_count(), 10);
}
#[test]
fn rate_n_records_one_in_n() {
let mut t = LatencyTracker::new(5);
for i in 0..50 {
t.record(i, || std::hint::black_box(i));
}
assert_eq!(t.samples_count(), 10);
}
#[test]
fn empty_samples_yield_zero_stats() {
let s = LatencyStats::from_samples(vec![]);
assert_eq!(s.p50, Duration::ZERO);
assert_eq!(s.p95, Duration::ZERO);
assert_eq!(s.p99, Duration::ZERO);
assert_eq!(s.samples_count, 0);
}
#[test]
fn percentiles_are_ordered() {
let samples: Vec<Duration> = (1..=100).map(|i| Duration::from_nanos(i as u64)).collect();
let s = LatencyStats::from_samples(samples);
assert!(s.p50 <= s.p95);
assert!(s.p95 <= s.p99);
}
#[test]
fn into_samples_moves_data() {
let mut t = LatencyTracker::new(1);
for i in 0..5 {
t.record(i, || ());
}
let s = t.into_samples();
assert_eq!(s.len(), 5);
}
#[test]
fn rate_zero_clamps_to_one() {
let mut t = LatencyTracker::new(0);
for i in 0..5 {
t.record(i, || ());
}
assert_eq!(t.samples_count(), 5);
}
}