use std::collections::VecDeque;
use std::time::{Duration, Instant};
pub const DEFAULT_WINDOW: Duration = Duration::from_millis(500);
#[derive(Debug, Clone)]
pub struct RateCalculator {
window: Duration,
samples: VecDeque<(Instant, usize)>,
bytes_in_window: usize,
}
impl Default for RateCalculator {
fn default() -> Self {
Self::new(DEFAULT_WINDOW)
}
}
impl RateCalculator {
pub fn new(window: Duration) -> Self {
Self {
window,
samples: VecDeque::new(),
bytes_in_window: 0,
}
}
pub fn add(&mut self, now: Instant, size: usize) {
self.samples.push_back((now, size));
self.bytes_in_window += size;
self.expire(now);
}
pub fn rate_bits_per_second(&mut self, now: Instant) -> Option<f64> {
self.expire(now);
if self.samples.len() < 2 {
return None;
}
let oldest = self.samples.front()?.0;
let span = now.saturating_duration_since(oldest);
if span.is_zero() {
return None;
}
Some((self.bytes_in_window * 8) as f64 / span.as_secs_f64())
}
fn expire(&mut self, now: Instant) {
let cutoff = now.checked_sub(self.window).unwrap_or(now);
while let Some((at, size)) = self.samples.front().copied() {
if at >= cutoff {
break;
}
self.samples.pop_front();
self.bytes_in_window -= size;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_measures_a_steady_stream() {
let epoch = Instant::now();
let mut calculator = RateCalculator::default();
for step in 0..50u64 {
calculator.add(epoch + Duration::from_millis(step * 10), 1500);
}
let rate = calculator
.rate_bits_per_second(epoch + Duration::from_millis(490))
.expect("a full window has a rate");
assert!(
(rate - 1_200_000.0).abs() < 100_000.0,
"expected about 1.2 Mb/s, got {rate}"
);
}
#[test]
fn an_empty_window_has_no_rate() {
let epoch = Instant::now();
let mut calculator = RateCalculator::default();
assert_eq!(None, calculator.rate_bits_per_second(epoch));
calculator.add(epoch, 1500);
assert_eq!(
None,
calculator.rate_bits_per_second(epoch),
"one sample is not a rate"
);
}
#[test]
fn old_samples_expire() {
let epoch = Instant::now();
let mut calculator = RateCalculator::new(Duration::from_millis(200));
for step in 0..20u64 {
calculator.add(epoch + Duration::from_millis(step * 10), 1500);
}
let busy = calculator
.rate_bits_per_second(epoch + Duration::from_millis(190))
.expect("rate");
calculator.add(epoch + Duration::from_millis(1_000), 1500);
calculator.add(epoch + Duration::from_millis(1_100), 1500);
let quiet = calculator
.rate_bits_per_second(epoch + Duration::from_millis(1_100))
.expect("rate");
assert!(
quiet < busy / 2.0,
"the window should have forgotten the busy period: {busy} then {quiet}"
);
}
#[test]
fn it_follows_a_halved_rate() {
let epoch = Instant::now();
let mut fast = RateCalculator::default();
let mut slow = RateCalculator::default();
for step in 0..50u64 {
fast.add(epoch + Duration::from_millis(step * 10), 1500);
}
for step in 0..25u64 {
slow.add(epoch + Duration::from_millis(step * 20), 1500);
}
let at = epoch + Duration::from_millis(490);
let fast_rate = fast.rate_bits_per_second(at).expect("rate");
let slow_rate = slow.rate_bits_per_second(at).expect("rate");
assert!(
(fast_rate / slow_rate - 2.0).abs() < 0.2,
"one should be twice the other: {fast_rate} vs {slow_rate}"
);
}
}