use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use serde::Serialize;
use tokio::sync::mpsc;
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum Event {
DownloadStarted {
id: String,
filename: String,
url: String,
total_size: Option<u64>,
resumed_bytes: u64,
connections: usize,
parallel: bool,
},
RangeStarted {
index: u64,
start: u64,
end: u64,
},
BytesWritten {
index: u64,
bytes: u64,
},
RangeCompleted {
index: u64,
},
RangeSplit {
index: u64,
new_index: u64,
at: u64,
},
RetryScheduled {
index: Option<u64>,
attempt: u32,
delay_ms: u64,
reason: String,
},
Checkpointed {
durable_bytes: u64,
},
DownloadPaused {
downloaded: u64,
total_size: Option<u64>,
},
VerificationStarted {
algorithm: String,
total_size: u64,
},
VerificationProgress {
bytes: u64,
total_size: u64,
},
VerificationCompleted {
algorithm: String,
ok: bool,
expected: Option<String>,
actual: String,
},
DownloadCompleted {
downloaded: u64,
elapsed_ms: u64,
average_bps: u64,
},
DownloadFailed {
error: String,
},
Note {
level: NoteLevel,
message: String,
},
}
#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum NoteLevel {
Info,
Warn,
Error,
}
#[derive(Debug, Default)]
pub struct Stats {
downloaded: AtomicU64,
durable: AtomicU64,
active_connections: AtomicUsize,
ranges_total: AtomicUsize,
ranges_complete: AtomicUsize,
retries: AtomicU64,
}
impl Stats {
pub fn add_downloaded(&self, n: u64) {
self.downloaded.fetch_add(n, Ordering::Relaxed);
}
pub fn set_downloaded(&self, n: u64) {
self.downloaded.store(n, Ordering::Relaxed);
}
pub fn downloaded(&self) -> u64 {
self.downloaded.load(Ordering::Relaxed)
}
pub fn set_durable(&self, n: u64) {
self.durable.store(n, Ordering::Relaxed);
}
pub fn durable(&self) -> u64 {
self.durable.load(Ordering::Relaxed)
}
pub fn connection_opened(&self) {
self.active_connections.fetch_add(1, Ordering::Relaxed);
}
pub fn connection_closed(&self) {
let _ = self
.active_connections
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
Some(v.saturating_sub(1))
});
}
pub fn active_connections(&self) -> usize {
self.active_connections.load(Ordering::Relaxed)
}
pub fn set_ranges_total(&self, n: usize) {
self.ranges_total.store(n, Ordering::Relaxed);
}
pub fn set_ranges_complete(&self, n: usize) {
self.ranges_complete.store(n, Ordering::Relaxed);
}
pub fn ranges(&self) -> (usize, usize) {
(
self.ranges_complete.load(Ordering::Relaxed),
self.ranges_total.load(Ordering::Relaxed),
)
}
pub fn record_retry(&self) {
self.retries.fetch_add(1, Ordering::Relaxed);
}
pub fn retries(&self) -> u64 {
self.retries.load(Ordering::Relaxed)
}
}
#[derive(Clone)]
pub struct Reporter {
tx: Option<mpsc::UnboundedSender<Event>>,
pub stats: Arc<Stats>,
}
impl Reporter {
pub fn new() -> (Self, mpsc::UnboundedReceiver<Event>) {
let (tx, rx) = mpsc::unbounded_channel();
(
Self {
tx: Some(tx),
stats: Arc::new(Stats::default()),
},
rx,
)
}
pub fn silent() -> Self {
Self {
tx: None,
stats: Arc::new(Stats::default()),
}
}
pub fn emit(&self, event: Event) {
if let Some(tx) = &self.tx {
let _ = tx.send(event);
}
}
pub fn note(&self, level: NoteLevel, message: impl Into<String>) {
self.emit(Event::Note {
level,
message: message.into(),
});
}
pub fn info(&self, message: impl Into<String>) {
self.note(NoteLevel::Info, message);
}
pub fn warn(&self, message: impl Into<String>) {
self.note(NoteLevel::Warn, message);
}
}
pub struct SpeedMeter {
window: Duration,
samples: VecDeque<(Instant, u64)>,
smoothed_bps: f64,
started: Instant,
start_bytes: u64,
}
impl SpeedMeter {
pub fn new(window: Duration, start_bytes: u64) -> Self {
let now = Instant::now();
let mut samples = VecDeque::with_capacity(64);
samples.push_back((now, start_bytes));
Self {
window,
samples,
smoothed_bps: 0.0,
started: now,
start_bytes,
}
}
pub fn record(&mut self, total_downloaded: u64) {
let now = Instant::now();
let instant_bps = self.instant_from(now, total_downloaded);
self.samples.push_back((now, total_downloaded));
while let Some(&(t, _)) = self.samples.front() {
if now.duration_since(t) > self.window && self.samples.len() > 2 {
self.samples.pop_front();
} else {
break;
}
}
const ALPHA: f64 = 0.15;
self.smoothed_bps = if self.smoothed_bps == 0.0 {
instant_bps
} else {
ALPHA * instant_bps + (1.0 - ALPHA) * self.smoothed_bps
};
}
fn instant_from(&self, now: Instant, total: u64) -> f64 {
let Some(&(t0, b0)) = self.samples.front() else {
return 0.0;
};
let dt = now.duration_since(t0).as_secs_f64();
if dt <= 0.0 {
return 0.0;
}
(total.saturating_sub(b0)) as f64 / dt
}
pub fn rolling_bps(&self) -> f64 {
let (&(t0, b0), &(t1, b1)) = match (self.samples.front(), self.samples.back()) {
(Some(a), Some(b)) => (a, b),
_ => return 0.0,
};
let dt = t1.duration_since(t0).as_secs_f64();
if dt <= 0.0 {
return 0.0;
}
(b1.saturating_sub(b0)) as f64 / dt
}
pub fn smoothed_bps(&self) -> f64 {
self.smoothed_bps
}
pub fn average_bps(&self, total_downloaded: u64) -> f64 {
let dt = self.started.elapsed().as_secs_f64();
if dt <= 0.0 {
return 0.0;
}
(total_downloaded.saturating_sub(self.start_bytes)) as f64 / dt
}
pub fn elapsed(&self) -> Duration {
self.started.elapsed()
}
}
#[derive(Debug, Clone, Serialize)]
pub struct Snapshot {
pub filename: String,
pub downloaded: u64,
pub total_size: Option<u64>,
pub bps: f64,
pub smoothed_bps: f64,
pub eta_secs: Option<u64>,
pub elapsed_ms: u64,
pub active_connections: usize,
pub ranges_complete: usize,
pub ranges_total: usize,
pub retries: u64,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn connection_count_never_wraps() {
let s = Stats::default();
s.connection_closed();
assert_eq!(s.active_connections(), 0);
s.connection_opened();
s.connection_opened();
assert_eq!(s.active_connections(), 2);
s.connection_closed();
assert_eq!(s.active_connections(), 1);
}
#[test]
fn silent_reporter_still_counts() {
let r = Reporter::silent();
r.emit(Event::RangeCompleted { index: 1 });
r.stats.add_downloaded(10);
assert_eq!(r.stats.downloaded(), 10);
}
#[test]
fn reporter_survives_dropped_consumer() {
let (r, rx) = Reporter::new();
drop(rx);
r.info("nobody is listening");
assert_eq!(r.stats.downloaded(), 0);
}
#[test]
fn rolling_speed_drops_when_transfer_stalls() {
let mut m = SpeedMeter::new(Duration::from_millis(300), 0);
m.record(1_000_000);
std::thread::sleep(Duration::from_millis(50));
m.record(2_000_000);
let fast = m.rolling_bps();
assert!(fast > 0.0);
for _ in 0..8 {
std::thread::sleep(Duration::from_millis(60));
m.record(2_000_000);
}
assert!(
m.rolling_bps() < fast / 2.0,
"rolling speed should collapse when stalled: {} vs {}",
m.rolling_bps(),
fast
);
}
#[test]
fn average_excludes_resumed_bytes() {
let m = SpeedMeter::new(Duration::from_secs(3), 5_000);
let avg = m.average_bps(6_000);
assert!(avg > 0.0);
assert!(m.average_bps(5_000) == 0.0);
}
}