downloader: Add progress updates

Author: cyberphantom52

crates/download-manager/src/download.rs

@@ -1,12 +1,13 @@
-use crate::{DownloadError, DownloadEvent, DownloadID};
+use crate::{DownloadError, DownloadEvent, DownloadID, Progress};
 use futures_core::Stream;
 use std::path::PathBuf;
-use tokio::sync::{broadcast, oneshot};
-use tokio_stream::wrappers::BroadcastStream;
+use tokio::sync::{broadcast, oneshot, watch};
+use tokio_stream::wrappers::{BroadcastStream, WatchStream};
 use tokio_util::sync::CancellationToken;
 
 pub struct Download {
     id: DownloadID,
+    progress: watch::Receiver<Progress>,
     events: broadcast::Receiver<DownloadEvent>,
     result: oneshot::Receiver<Result<DownloadResult, DownloadError>>,
 
@@ -16,12 +17,14 @@ pub struct Download {
 impl Download {
     pub fn new(
         id: DownloadID,
+        progress: watch::Receiver<Progress>,
         events: broadcast::Receiver<DownloadEvent>,
         result: oneshot::Receiver<Result<DownloadResult, DownloadError>>,
         cancel_token: CancellationToken,
     ) -> Self {
         Download {
             id,
+            progress,
             events,
             result,
             cancel_token,
@@ -36,6 +39,14 @@ impl Download {
         self.cancel_token.cancel();
     }
 
+    pub fn progress_raw(&self) -> watch::Receiver<Progress> {
+        self.progress.clone()
+    }
+
+    pub fn progress(&self) -> impl Stream<Item = Progress> + 'static {
+        WatchStream::new(self.progress_raw())
+    }
+
     pub fn events(&self) -> impl Stream<Item = DownloadEvent> + 'static {
         use tokio_stream::StreamExt as _;
 

crates/download-manager/src/events.rs

@@ -1,7 +1,7 @@
 use crate::DownloadID;
 use reqwest::Url;
 use std::path::PathBuf;
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 #[derive(Debug, Clone)]
 pub enum DownloadEvent {
@@ -39,23 +39,131 @@ pub enum DownloadEvent {
 pub struct Progress {
     pub bytes_downloaded: u64,
     pub total_bytes: Option<u64>,
-    pub instantaneous_bps: f64,
-    pub avg_bps: f64,
 
+    pub instantaneous_bps: f64, // most recent sample
+    pub ema_bps: f64,           // exponential moving average
+
+    // Internal timing / sampling
     started_at: Instant,
-    updated_at: Instant,
+    updated_at: Instant,     // last time we saw any bytes
+    last_sample_at: Instant, // last time we recomputed instantaneous_bps
+    last_sample_bytes: u64,  // bytes_downloaded at last sample
+
+    // Sampling thresholds
+    pub min_sample_interval: Duration,
+    pub min_sample_bytes: u64,
+
+    ema_alpha: f64, // smoothing factor for EMA
 }
 
 impl Progress {
-    pub fn new() -> Self {
+    pub fn new(total_bytes: Option<u64>) -> Self {
         let now = Instant::now();
         Progress {
             bytes_downloaded: 0,
-            total_bytes: None,
+            total_bytes,
+            instantaneous_bps: 0.0,
+            ema_bps: 0.0,
             started_at: now,
             updated_at: now,
-            instantaneous_bps: 0.0,
-            avg_bps: 0.0,
+            last_sample_at: now,
+            last_sample_bytes: 0,
+            ema_alpha: 0.2,
+            min_sample_bytes: 64 * 1024, // 64 KiB
+            min_sample_interval: Duration::from_millis(200),
+        }
+    }
+
+    pub fn with_sample_interval(mut self, min_sample_interval: Duration) -> Self {
+        self.min_sample_interval = min_sample_interval;
+        self
+    }
+
+    pub fn with_sample_bytes(mut self, min_sample_bytes: u64) -> Self {
+        self.min_sample_bytes = min_sample_bytes;
+        self
+    }
+
+    pub fn with_ema_alpha(mut self, ema_alpha: f64) -> Self {
+        self.ema_alpha = ema_alpha;
+        self
+    }
+
+    pub fn bytes_downloaded(&self) -> u64 {
+        self.bytes_downloaded
+    }
+
+    pub fn elapsed(&self) -> Duration {
+        self.started_at.elapsed()
+    }
+
+    pub fn percent(&self) -> Option<f64> {
+        self.total_bytes
+            .filter(|&total| total > 0)
+            .map(|total| (self.bytes_downloaded as f64 / total as f64) * 100.0)
+    }
+
+    pub fn remaining_bytes(&self) -> Option<u64> {
+        self.total_bytes
+            .map(|t| t.saturating_sub(self.bytes_downloaded))
+    }
+
+    pub fn eta(&self) -> Option<Duration> {
+        let remaining = self.remaining_bytes()?;
+        if self.ema_bps > 0.0 {
+            Some(Duration::from_secs_f64(remaining as f64 / self.ema_bps))
+        } else {
+            None
+        }
+    }
+
+    pub fn update(&mut self, chunk_len: u64) -> bool {
+        let now = Instant::now();
+        self.bytes_downloaded += chunk_len;
+        self.updated_at = now;
+
+        let dt = now.duration_since(self.last_sample_at);
+        let byte_delta = self.bytes_downloaded - self.last_sample_bytes;
+
+        if dt >= self.min_sample_interval || byte_delta >= self.min_sample_bytes {
+            let secs = dt.as_secs_f64();
+            if secs > 0.0 && byte_delta > 0 {
+                // instantaneous over sample window
+                let inst = (byte_delta as f64) / secs;
+                self.instantaneous_bps = inst;
+
+                // EMA (if first sample, seed with inst)
+                if self.ema_bps <= 0.0 {
+                    self.ema_bps = inst;
+                } else {
+                    self.ema_bps = self.ema_alpha * inst + (1.0 - self.ema_alpha) * self.ema_bps;
+                }
+            }
+
+            self.last_sample_at = now;
+            self.last_sample_bytes = self.bytes_downloaded;
+            return true;
+        }
+
+        false
+    }
+
+    pub fn force_update(&mut self) {
+        let now = Instant::now();
+        let dt = now.duration_since(self.last_sample_at);
+        let byte_delta = self.bytes_downloaded - self.last_sample_bytes;
+
+        if dt.as_secs_f64() > 0.0 && byte_delta > 0 {
+            let inst = (byte_delta as f64) / dt.as_secs_f64();
+            self.instantaneous_bps = inst;
+            if self.ema_bps <= 0.0 {
+                self.ema_bps = inst;
+            } else {
+                self.ema_bps = self.ema_alpha * inst + (1.0 - self.ema_alpha) * self.ema_bps;
+            }
+            self.last_sample_at = now;
+            self.last_sample_bytes = self.bytes_downloaded;
         }
+        self.updated_at = now;
     }
 }

crates/download-manager/src/request.rs

@@ -1,4 +1,6 @@
-use crate::{Download, DownloadError, DownloadEvent, DownloadID, DownloadManager, DownloadResult};
+use crate::{
+    Download, DownloadError, DownloadEvent, DownloadID, DownloadManager, DownloadResult, Progress,
+};
 use derive_builder::Builder;
 use reqwest::Url;
 use std::{
@@ -14,6 +16,7 @@ pub struct Request {
     destination: PathBuf,
     config: DownloadConfig,
 
+    progress: watch::Sender<Progress>,
     events: broadcast::Sender<DownloadEvent>,
     result: oneshot::Sender<Result<DownloadResult, DownloadError>>,
 
@@ -98,6 +101,11 @@ impl Request {
         let _ = self.result.send(result);
     }
 
+    pub fn update_progress(&self, progress: Progress) {
+        // TODO: Log the error
+        let _ = self.progress.send(progress);
+    }
+
     pub fn start(&self) {
         self.emit(DownloadEvent::Started {
             id: self.id(),
@@ -175,6 +183,7 @@ impl RequestBuilder<'_> {
             .ok_or_else(|| anyhow::anyhow!("Destination must be set"))?;
         let config = self.config.build()?;
 
+        let (progress_tx, progress_rx) = watch::channel(Progress::new(None));
         let (result_tx, result_rx) = oneshot::channel();
         let cancel_token = self.manager.child_token();
 
@@ -186,13 +195,20 @@ impl RequestBuilder<'_> {
             url,
             destination,
             config,
+            progress: progress_tx,
             events: event_tx,
             result: result_tx,
             cancel_token: cancel_token.clone(),
         };
 
         self.manager.queue_request(request)?;
 
-        Ok(Download::new(id, event_rx, result_rx, cancel_token))
+        Ok(Download::new(
+            id,
+            progress_rx,
+            event_rx,
+            result_rx,
+            cancel_token,
+        ))
     }
 }

crates/download-manager/src/worker.rs

@@ -1,4 +1,4 @@
-use crate::{DownloadError, DownloadResult, Request};
+use crate::{DownloadError, DownloadResult, Progress, Request};
 use reqwest::Client;
 use std::time::Duration;
 use tokio::{fs::File, io::AsyncWriteExt};
@@ -84,7 +84,7 @@ async fn attempt_download(
     }
 
     let mut file = File::create(&request.destination()).await?;
-    let mut bytes_downloaded = 0;
+    let mut progress = Progress::new(response.content_length());
     loop {
         tokio::select! {
             _ = request.cancel_token.cancelled() => {
@@ -96,7 +96,9 @@ async fn attempt_download(
                 match chunk {
                     Ok(Some(chunk)) => {
                         file.write_all(&chunk).await?;
-                        bytes_downloaded += chunk.len() as u64;
+                        if progress.update(chunk.len() as u64) {
+                            request.update_progress(progress);
+                        }
                     }
                     Ok(None) => break,
                     Err(e) => {
@@ -108,12 +110,14 @@ async fn attempt_download(
             }
         }
     }
+    progress.force_update();
+    request.update_progress(progress);
 
     // Ensure the data is written to disk
     file.sync_all().await?;
 
     Ok(DownloadResult {
         path: request.destination().to_path_buf(),
-        bytes_downloaded,
+        bytes_downloaded: progress.bytes_downloaded(),
     })
 }