use std::collections::VecDeque;
use serde::{Deserialize, Serialize};
pub const MIN_SAMPLES: usize = 3;
pub const MIN_SPAN_MS: u64 = 3_000;
pub const WINDOW_MS: u64 = 30_000;
const MAX_SAMPLES: usize = 256;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Sample {
at_ms: u64,
bytes: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct RateReport {
pub stable: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub bytes_per_second: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub eta_seconds: Option<f64>,
pub samples: usize,
}
impl RateReport {
fn warming(samples: usize) -> Self {
RateReport {
stable: false,
bytes_per_second: None,
eta_seconds: None,
samples,
}
}
}
#[derive(Debug, Default, Clone)]
pub struct RateEstimator {
window: VecDeque<Sample>,
}
impl RateEstimator {
pub fn new() -> Self {
Self::default()
}
pub fn observe(&mut self, at_ms: u64, bytes: u64) {
if let Some(last) = self.window.back() {
if at_ms < last.at_ms {
return;
}
if bytes < last.bytes {
self.window.clear();
}
}
self.window.push_back(Sample { at_ms, bytes });
let cutoff = at_ms.saturating_sub(WINDOW_MS);
while self.window.len() > 1 && self.window.front().is_some_and(|s| s.at_ms < cutoff) {
self.window.pop_front();
}
while self.window.len() > MAX_SAMPLES {
self.window.remove(1);
}
}
pub fn reset(&mut self) {
self.window.clear();
}
pub fn bytes_done(&self) -> Option<u64> {
self.window.back().map(|s| s.bytes)
}
pub fn report(&self, total_bytes: Option<u64>) -> RateReport {
let (Some(first), Some(last)) = (self.window.front(), self.window.back()) else {
return RateReport::warming(0);
};
let span_ms = last.at_ms - first.at_ms;
if self.window.len() < MIN_SAMPLES || span_ms < MIN_SPAN_MS {
return RateReport::warming(self.window.len());
}
let bytes = last.bytes.saturating_sub(first.bytes) as f64;
let rate = bytes / (span_ms as f64 / 1000.0);
let eta = total_bytes.filter(|_| rate > 0.0).map(|total| {
total.saturating_sub(last.bytes) as f64 / rate
});
RateReport {
stable: true,
bytes_per_second: Some(rate),
eta_seconds: eta,
samples: self.window.len(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn first_tick_reports_no_rate_at_all() {
let mut est = RateEstimator::new();
est.observe(0, 0);
est.observe(2, 8 * 1024 * 1024); let report = est.report(Some(1 << 30));
assert!(!report.stable);
assert_eq!(report.bytes_per_second, None);
assert_eq!(report.eta_seconds, None);
}
#[test]
fn three_samples_are_not_enough_without_three_seconds() {
let mut est = RateEstimator::new();
for i in 0..5 {
est.observe(i * 100, i * 1_000_000);
}
assert!(!est.report(None).stable);
}
#[test]
fn three_seconds_are_not_enough_without_three_samples() {
let mut est = RateEstimator::new();
est.observe(0, 0);
est.observe(5_000, 5_000_000);
assert!(!est.report(None).stable);
}
#[test]
fn reports_a_stable_rate_and_eta_once_the_window_qualifies() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
let report = est.report(Some(10_000_000));
assert!(report.stable);
assert_eq!(report.bytes_per_second, Some(1_000_000.0));
assert_eq!(report.eta_seconds, Some(6.0));
}
#[test]
fn a_restarted_transfer_clears_the_window_instead_of_going_negative() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
assert!(est.report(None).stable);
est.observe(5_000, 0); let report = est.report(None);
assert!(!report.stable);
assert_eq!(report.samples, 1);
assert_eq!(est.bytes_done(), Some(0));
}
#[test]
fn eta_is_clamped_at_zero_when_the_total_is_wrong() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
assert_eq!(est.report(Some(1_000_000)).eta_seconds, Some(0.0));
}
#[test]
fn no_total_means_a_rate_but_no_eta() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
let report = est.report(None);
assert!(report.stable);
assert!(report.bytes_per_second.is_some());
assert_eq!(report.eta_seconds, None);
}
#[test]
fn a_stalled_transfer_reports_zero_rather_than_an_eta() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, 1_000_000);
}
let report = est.report(Some(2_000_000));
assert_eq!(report.bytes_per_second, Some(0.0));
assert_eq!(report.eta_seconds, None);
}
#[test]
fn samples_older_than_the_window_are_dropped() {
let mut est = RateEstimator::new();
est.observe(0, 0);
for i in 0..=4u64 {
est.observe(WINDOW_MS + i * 1000, 1_000_000 + i * 1_000_000);
}
assert_eq!(est.report(None).bytes_per_second, Some(1_000_000.0));
}
#[test]
fn a_fast_ticking_job_still_becomes_stable() {
let mut est = RateEstimator::new();
for i in 0..=1000u64 {
est.observe(i * 10, i * 10_000);
}
let report = est.report(None);
assert!(report.stable, "{report:?}");
assert_eq!(report.bytes_per_second, Some(1_000_000.0));
}
#[test]
fn a_backwards_clock_sample_is_ignored() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
est.observe(500, 9_000_000);
assert_eq!(est.bytes_done(), Some(4_000_000));
assert_eq!(est.report(None).bytes_per_second, Some(1_000_000.0));
}
#[test]
fn reset_returns_to_warming() {
let mut est = RateEstimator::new();
for i in 0..=4u64 {
est.observe(i * 1000, i * 1_000_000);
}
est.reset();
let report = est.report(None);
assert!(!report.stable);
assert_eq!(report.samples, 0);
}
#[test]
fn warming_reports_omit_the_absent_numbers_rather_than_nulling_them() {
let json = serde_json::to_string(&RateReport::warming(1)).unwrap();
assert_eq!(json, "{\"stable\":false,\"samples\":1}");
}
}