use std::{
collections::VecDeque,
sync::atomic::{AtomicU64, Ordering},
time::{Duration, Instant},
};
use parking_lot::Mutex;
#[derive(Clone, Copy)]
struct ProgressSnapshot {
progress_bytes: u64,
instant: Instant,
}
pub struct SpeedEstimator {
latest_per_second_snapshots: Mutex<VecDeque<ProgressSnapshot>>,
bytes_per_second: AtomicU64,
time_remaining_millis: AtomicU64,
}
impl Default for SpeedEstimator {
fn default() -> Self {
Self::new(5)
}
}
impl SpeedEstimator {
pub fn new(window_seconds: usize) -> Self {
assert!(window_seconds > 1);
Self {
latest_per_second_snapshots: Mutex::new(VecDeque::with_capacity(window_seconds)),
bytes_per_second: Default::default(),
time_remaining_millis: Default::default(),
}
}
pub fn time_remaining(&self) -> Option<Duration> {
let tr = self.time_remaining_millis.load(Ordering::Relaxed);
if tr == 0 {
return None;
}
Some(Duration::from_millis(tr))
}
pub fn bps(&self) -> u64 {
self.bytes_per_second.load(Ordering::Relaxed)
}
pub fn mbps(&self) -> f64 {
self.bps() as f64 / 1024f64 / 1024f64
}
pub fn add_snapshot(
&self,
progress_bytes: u64,
remaining_bytes: Option<u64>,
instant: Instant,
) {
let first = {
let mut g = self.latest_per_second_snapshots.lock();
let current = ProgressSnapshot {
progress_bytes,
instant,
};
if g.is_empty() {
g.push_back(current);
return;
} else if g.len() < g.capacity() {
g.push_back(current);
g.front().copied().unwrap()
} else {
let first = g.pop_front().unwrap();
g.push_back(current);
first
}
};
let downloaded_bytes_diff = progress_bytes - first.progress_bytes;
let elapsed = instant - first.instant;
let bps = downloaded_bytes_diff as f64 / elapsed.as_secs_f64();
let time_remaining_millis_rounded: u64 = if downloaded_bytes_diff > 0 {
let time_remaining_secs = remaining_bytes.unwrap_or_default() as f64 / bps;
(time_remaining_secs * 1000f64) as u64
} else {
0
};
self.time_remaining_millis
.store(time_remaining_millis_rounded, Ordering::Relaxed);
self.bytes_per_second.store(bps as u64, Ordering::Relaxed);
}
}