downloader: Add event based updates

Author: cyberphantom52

crates/download-manager/src/context.rs

@@ -3,9 +3,11 @@ use std::sync::{
     Arc,
     atomic::{AtomicU64, AtomicUsize, Ordering},
 };
-use tokio::sync::Semaphore;
+use tokio::sync::{Semaphore, broadcast};
 use tokio_util::sync::CancellationToken;
 
+use crate::DownloadEvent;
+
 pub type DownloadID = u64;
 
 #[derive(Debug)]
@@ -18,17 +20,21 @@ pub(crate) struct Context {
     pub id_counter: AtomicU64,
     pub active: AtomicUsize,
     pub max_concurrent: AtomicUsize,
+
+    pub events: broadcast::Sender<DownloadEvent>,
 }
 
 impl Context {
     pub fn new(max_concurrent: usize) -> Arc<Self> {
+        let (events, _) = broadcast::channel(1024);
         Arc::new(Self {
             semaphore: Arc::new(Semaphore::new(max_concurrent)),
             max_concurrent: AtomicUsize::new(max_concurrent),
             cancel_root: CancellationToken::new(),
             active: AtomicUsize::new(0),
             id_counter: AtomicU64::new(1),
             client: Client::new(),
+            events,
         })
     }
 

crates/download-manager/src/download.rs

@@ -1,12 +1,11 @@
-use crate::{DownloadError, DownloadID};
-use anyhow::anyhow;
+use crate::{DownloadError, DownloadEvent, DownloadID};
 use std::path::PathBuf;
-use tokio::sync::{oneshot, watch};
+use tokio::sync::{broadcast, oneshot};
 use tokio_util::sync::CancellationToken;
 
 pub struct Download {
     id: DownloadID,
-    status: watch::Receiver<Status>,
+    events: broadcast::Receiver<DownloadEvent>,
     result: oneshot::Receiver<Result<DownloadResult, DownloadError>>,
 
     cancel_token: CancellationToken,
@@ -15,13 +14,13 @@ pub struct Download {
 impl Download {
     pub fn new(
         id: DownloadID,
-        status: watch::Receiver<Status>,
+        events: broadcast::Receiver<DownloadEvent>,
         result: oneshot::Receiver<Result<DownloadResult, DownloadError>>,
         cancel_token: CancellationToken,
     ) -> Self {
         Download {
             id,
-            status,
+            events,
             result,
             cancel_token,
         }
@@ -34,10 +33,6 @@ impl Download {
     pub fn cancel(&self) {
         self.cancel_token.cancel();
     }
-
-    pub fn status(&self) -> Status {
-        *self.status.borrow()
-    }
 }
 
 impl std::future::Future for Download {
@@ -58,17 +53,8 @@ impl std::future::Future for Download {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum Status {
-    Queued,
-    Running,
-    Retrying(u32),
-    Completed,
-    Cancelled,
-    Failed,
-}
-
 #[derive(Debug)]
 pub struct DownloadResult {
     pub path: PathBuf,
+    pub bytes_downloaded: u64,
 }

crates/download-manager/src/download_manager.rs

@@ -1,14 +1,16 @@
 mod context;
 mod download;
 mod error;
+mod events;
 mod request;
 mod worker;
 
 use crate::{context::Context, request::RequestBuilder, worker::download_thread};
 pub use crate::{
     context::DownloadID,
-    download::{Download, DownloadResult, Status},
+    download::{Download, DownloadResult},
     error::DownloadError,
+    events::{DownloadEvent, Progress},
     request::Request,
 };
 use reqwest::Url;

crates/download-manager/src/events.rs

@@ -0,0 +1,61 @@
+use crate::DownloadID;
+use reqwest::Url;
+use std::path::PathBuf;
+use std::time::Instant;
+
+#[derive(Debug, Clone)]
+pub enum DownloadEvent {
+    Queued {
+        id: DownloadID,
+        url: Url,
+        destination: PathBuf,
+    },
+    Started {
+        id: DownloadID,
+        url: Url,
+        destination: PathBuf,
+        total_bytes: Option<u64>,
+    },
+    Retrying {
+        id: DownloadID,
+        attempt: u32,
+        next_delay_ms: u64,
+    },
+    Completed {
+        id: DownloadID,
+        path: PathBuf,
+        bytes_downloaded: u64,
+    },
+    Failed {
+        id: DownloadID,
+        error: String,
+    },
+    Cancelled {
+        id: DownloadID,
+    },
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Progress {
+    pub bytes_downloaded: u64,
+    pub total_bytes: Option<u64>,
+    pub instantaneous_bps: f64,
+    pub avg_bps: f64,
+
+    started_at: Instant,
+    updated_at: Instant,
+}
+
+impl Progress {
+    pub fn new() -> Self {
+        let now = Instant::now();
+        Progress {
+            bytes_downloaded: 0,
+            total_bytes: None,
+            started_at: now,
+            updated_at: now,
+            instantaneous_bps: 0.0,
+            avg_bps: 0.0,
+        }
+    }
+}

crates/download-manager/src/request.rs

@@ -1,8 +1,11 @@
-use crate::{Download, DownloadError, DownloadID, DownloadManager, DownloadResult, Status};
+use crate::{Download, DownloadError, DownloadEvent, DownloadID, DownloadManager, DownloadResult};
 use derive_builder::Builder;
 use reqwest::Url;
-use std::path::{Path, PathBuf};
-use tokio::sync::{oneshot, watch};
+use std::{
+    path::{Path, PathBuf},
+    time::Duration,
+};
+use tokio::sync::{broadcast, oneshot, watch};
 use tokio_util::sync::CancellationToken;
 
 pub struct Request {
@@ -11,7 +14,7 @@ pub struct Request {
     destination: PathBuf,
     config: DownloadConfig,
 
-    status: watch::Sender<Status>,
+    events: broadcast::Sender<DownloadEvent>,
     result: oneshot::Sender<Result<DownloadResult, DownloadError>>,
 
     pub cancel_token: CancellationToken,
@@ -62,6 +65,10 @@ impl Request {
         }
     }
 
+    pub fn id(&self) -> DownloadID {
+        self.id
+    }
+
     /* TODO:
      * Add callbacks like `on_update`, `on_progress`, `on_complete`, etc.
      */
@@ -81,39 +88,48 @@ impl Request {
         self.cancel_token.is_cancelled()
     }
 
-    fn mark_status(&self, status: Status) {
+    fn emit(&self, event: DownloadEvent) {
         // TODO: Log the error
-        let _ = self.status.send(status);
+        let _ = self.events.send(event);
     }
 
     fn send_result(self, result: Result<DownloadResult, DownloadError>) {
         // TODO: Log the error
         let _ = self.result.send(result);
     }
 
-    pub fn mark_running(&self) {
-        self.mark_status(Status::Running)
+    pub fn start(&self) {
+        self.emit(DownloadEvent::Started {
+            id: self.id(),
+            url: self.url().clone(),
+            destination: self.destination.clone(),
+            total_bytes: None,
+        });
     }
 
-    pub fn mark_failed(self, error: DownloadError) {
-        match error {
-            DownloadError::Cancelled => self.mark_status(Status::Cancelled),
-            _ => self.mark_status(Status::Failed),
-        }
+    pub fn fail(self, error: DownloadError) {
         self.send_result(Err(error));
     }
 
-    pub fn mark_completed(self, result: DownloadResult) {
-        self.mark_status(Status::Completed);
+    pub fn finish(self, result: DownloadResult) {
+        self.emit(DownloadEvent::Completed {
+            id: self.id(),
+            path: result.path.clone(),
+            bytes_downloaded: result.bytes_downloaded,
+        });
         self.send_result(Ok(result))
     }
 
-    pub fn mark_retrying(&self, retry_count: u32) {
-        self.mark_status(Status::Retrying(retry_count))
+    pub fn retry(&self, attempt: u32, delay: Duration) {
+        self.emit(DownloadEvent::Retrying {
+            id: self.id(),
+            attempt,
+            next_delay_ms: delay.as_millis() as u64,
+        });
     }
 
-    pub fn mark_cancelled(self) {
-        self.mark_status(Status::Cancelled);
+    pub fn cancel(self) {
+        self.emit(DownloadEvent::Cancelled { id: self.id() });
         self.send_result(Err(DownloadError::Cancelled))
     }
 }
@@ -159,23 +175,24 @@ impl RequestBuilder<'_> {
             .ok_or_else(|| anyhow::anyhow!("Destination must be set"))?;
         let config = self.config.build()?;
 
-        let (status_tx, status_rx) = watch::channel(Status::Queued);
         let (result_tx, result_rx) = oneshot::channel();
         let cancel_token = self.manager.child_token();
 
+        let event_tx = self.manager.ctx.events.clone();
+        let event_rx = event_tx.subscribe();
         let id = self.manager.ctx.next_id();
         let request = Request {
             id,
             url,
             destination,
             config,
-            status: status_tx,
+            events: event_tx,
             result: result_tx,
             cancel_token: cancel_token.clone(),
         };
 
         self.manager.queue_request(request)?;
 
-        Ok(Download::new(id, status_rx, result_rx, cancel_token))
+        Ok(Download::new(id, event_rx, result_rx, cancel_token))
     }
 }

crates/download-manager/src/worker.rs

@@ -21,45 +21,43 @@ const BACKOFF_STRATEGY: ExponentialBackoff = ExponentialBackoff {
 };
 
 pub(super) async fn download_thread(client: Client, mut request: Request) {
-    request.mark_running();
+    request.start();
     let mut last_retryable_error: DownloadError =
         DownloadError::Unknown("Unknown Error".to_string());
 
     let retries = request.config().retries();
     for attempt in 0..=retries {
         if attempt > 0 {
-            request.mark_retrying(attempt);
-
-            //TODO: Add proper backoff
             let delay = BACKOFF_STRATEGY.next_delay(attempt);
             let mut interval = tokio::time::interval(delay);
 
+            request.retry(attempt, delay);
             tokio::select! {
                 _ = interval.tick() => {},
                 _ = request.cancel_token.cancelled() => {
-                    request.mark_cancelled();
+                    request.cancel();
                     return;
                 }
             }
         }
 
         match attempt_download(client.clone(), &mut request).await {
             Ok(result) => {
-                request.mark_completed(result);
+                request.finish(result);
                 return;
             }
             Err(error) if error.is_retryable() => {
                 last_retryable_error = error;
                 continue;
             }
             Err(error) => {
-                request.mark_failed(error);
+                request.fail(error);
                 return;
             }
         };
     }
 
-    request.mark_failed(DownloadError::RetriesExhausted {
+    request.fail(DownloadError::RetriesExhausted {
         last_error: Box::new(last_retryable_error),
     });
 }
@@ -86,7 +84,7 @@ async fn attempt_download(
     }
 
     let mut file = File::create(&request.destination()).await?;
-
+    let mut bytes_downloaded = 0;
     loop {
         tokio::select! {
             _ = request.cancel_token.cancelled() => {
@@ -98,6 +96,7 @@ async fn attempt_download(
                 match chunk {
                     Ok(Some(chunk)) => {
                         file.write_all(&chunk).await?;
+                        bytes_downloaded += chunk.len() as u64;
                     }
                     Ok(None) => break,
                     Err(e) => {
@@ -115,5 +114,6 @@ async fn attempt_download(
 
     Ok(DownloadResult {
         path: request.destination().to_path_buf(),
+        bytes_downloaded,
     })
 }