downloader: Use builder pattern to build download requests

Author: cyberphantom52

src/downloader/config.rs

@@ -53,6 +53,21 @@ impl Default for DownloadConfig {
 }
 
 impl DownloadConfig {
+    pub fn with_max_retries(mut self, retries: usize) -> Self {
+        self.max_retries = retries;
+        self
+    }
+
+    pub fn with_user_agent(mut self, user_agent: impl Into<String>) -> Self {
+        self.user_agent = Some(user_agent.into());
+        self
+    }
+
+    pub fn with_progress_interval(mut self, interval: Duration) -> Self {
+        self.progress_update_interval = interval;
+        self
+    }
+
     pub fn max_retries(&self) -> usize {
         self.max_retries
     }

src/downloader/manager.rs

@@ -1,5 +1,5 @@
 use super::{
-    config::DownloadManagerConfig, download_thread, DownloadConfig, DownloadHandle, DownloadRequest,
+    download_thread, DownloadBuilder, DownloadConfig, DownloadManagerConfig, DownloadRequest,
 };
 use crate::{error::DownloadError, Error};
 use reqwest::{Client, Url};
@@ -40,40 +40,25 @@ impl DownloadManager {
         manager
     }
 
-    pub fn download_with_config(
+    pub fn download(
         &self,
-        url: Url,
+        url: impl TryInto<Url>,
         destination: impl AsRef<Path>,
-        config: DownloadConfig,
-    ) -> Result<DownloadHandle, Error> {
-        if self.cancel.is_cancelled() {
-            return Err(Error::Download(DownloadError::ManagerShutdown));
-        }
-
-        let destination = destination.as_ref();
-        if destination.exists() {
-            return Err(Error::Download(DownloadError::FileExists {
-                path: destination.to_path_buf(),
-            }));
-        }
-
-        let cancel = self.cancel.child_token();
-        let (req, handle) = DownloadRequest::new(url, destination, cancel, config);
-
-        self.queue.try_send(req).map_err(|e| match e {
-            mpsc::error::TrySendError::Full(_) => Error::Download(DownloadError::QueueFull),
-            mpsc::error::TrySendError::Closed(_) => Error::Download(DownloadError::ManagerShutdown),
-        })?;
-
-        Ok(handle)
+    ) -> Result<DownloadBuilder, Error> {
+        let url = url
+            .try_into()
+            .map_err(|_| Error::Download(DownloadError::InvalidUrl))?;
+        Ok(DownloadBuilder::new(self, url, destination))
     }
 
-    pub fn download(
+    pub fn download_with_config(
         &self,
         url: Url,
         destination: impl AsRef<Path>,
-    ) -> Result<DownloadHandle, Error> {
-        self.download_with_config(url, destination, DownloadConfig::default())
+        config: DownloadConfig,
+    ) -> Result<DownloadBuilder, Error> {
+        self.download(url, destination)
+            .map(|builder| builder.with_config(config))
     }
 
     pub async fn set_max_parallel_downloads(&self, limit: usize) -> Result<(), Error> {
@@ -116,6 +101,21 @@ impl DownloadManager {
         drop(self.queue);
         Ok(())
     }
+
+    pub fn is_cancelled(&self) -> bool {
+        self.cancel.is_cancelled()
+    }
+
+    pub fn child_token(&self) -> CancellationToken {
+        self.cancel.child_token()
+    }
+
+    pub fn queue_request(&self, req: DownloadRequest) -> Result<(), Error> {
+        self.queue.try_send(req).map_err(|e| match e {
+            mpsc::error::TrySendError::Full(_) => Error::Download(DownloadError::QueueFull),
+            mpsc::error::TrySendError::Closed(_) => Error::Download(DownloadError::ManagerShutdown),
+        })
+    }
 }
 
 async fn dispatcher_thread(

src/downloader/mod.rs

@@ -2,13 +2,12 @@ mod config;
 mod handle;
 mod manager;
 mod progress;
-mod types;
+mod request;
 mod worker;
 
 pub use config::{DownloadConfig, DownloadManagerConfig};
 pub use handle::DownloadHandle;
 pub use manager::DownloadManager;
-pub use progress::DownloadProgress;
-pub(self) use types::DownloadRequest;
-pub use types::Status;
+pub use progress::{DownloadProgress, Status};
+pub use request::{DownloadBuilder, DownloadRequest};
 pub(self) use worker::download_thread;

src/downloader/progress.rs

@@ -112,3 +112,13 @@ impl DownloadProgress {
         self.start_time.elapsed()
     }
 }
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum Status {
+    Queued,
+    InProgress(DownloadProgress),
+    Retrying,
+    Completed,
+    Failed,
+    Cancelled,
+}

src/downloader/request.rs

@@ -0,0 +1,135 @@
+use super::{DownloadConfig, DownloadHandle, DownloadManager, DownloadProgress, Status};
+use crate::{error::DownloadError, Error};
+use reqwest::Url;
+use std::{
+    path::{Path, PathBuf},
+    time::{Duration, Instant},
+};
+use tokio::{
+    fs::File,
+    sync::{oneshot, watch},
+};
+use tokio_util::sync::CancellationToken;
+
+#[derive(Debug)]
+pub struct DownloadRequest {
+    url: Url,
+    destination: PathBuf,
+    pub result: oneshot::Sender<Result<File, Error>>,
+    pub status: watch::Sender<Status>,
+    pub cancel: CancellationToken,
+    config: DownloadConfig,
+
+    // For rate limiting
+    last_progress_update: Instant,
+}
+
+impl DownloadRequest {
+    pub fn new(
+        url: Url,
+        destination: impl AsRef<Path>,
+        cancel: CancellationToken,
+        config: DownloadConfig,
+    ) -> (Self, DownloadHandle) {
+        let (result_tx, result_rx) = oneshot::channel();
+        let (status_tx, status_rx) = watch::channel(Status::Queued);
+        (
+            Self {
+                url,
+                destination: destination.as_ref().to_path_buf(),
+                result: result_tx,
+                status: status_tx,
+                cancel: cancel.clone(),
+                config,
+                last_progress_update: Instant::now(),
+            },
+            DownloadHandle::new(result_rx, status_rx, cancel),
+        )
+    }
+
+    pub fn url(&self) -> &Url {
+        &self.url
+    }
+
+    pub fn destination(&self) -> &Path {
+        &self.destination
+    }
+
+    pub fn config(&self) -> &DownloadConfig {
+        &self.config
+    }
+
+    pub fn send_progress(&mut self, progress: DownloadProgress) {
+        if self.last_progress_update.elapsed() >= self.config.progress_update_interval() {
+            self.last_progress_update = Instant::now();
+            self.status.send(Status::InProgress(progress)).ok();
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct DownloadBuilder<'a> {
+    manager: &'a DownloadManager,
+    url: Url,
+    destination: PathBuf,
+    config: DownloadConfig,
+}
+
+impl<'a> DownloadBuilder<'a> {
+    pub fn new(manager: &'a DownloadManager, url: Url, destination: impl AsRef<Path>) -> Self {
+        Self {
+            manager,
+            url,
+            destination: destination.as_ref().to_path_buf(),
+            config: DownloadConfig::default(),
+        }
+    }
+
+    pub fn with_retries(mut self, retries: usize) -> Self {
+        self.config = self.config.with_max_retries(retries);
+        self
+    }
+
+    pub fn with_user_agent(mut self, user_agent: impl Into<String>) -> Self {
+        self.config = self.config.with_user_agent(user_agent);
+        self
+    }
+
+    /// Set how often progress updates are sent
+    pub fn with_progress_interval(mut self, interval: Duration) -> Self {
+        self.config = self.config.with_progress_interval(interval);
+        self
+    }
+
+    pub fn with_config(mut self, config: DownloadConfig) -> Self {
+        self.config = config;
+        self
+    }
+
+    pub fn url(&self) -> &Url {
+        &self.url
+    }
+
+    pub fn config(&self) -> &DownloadConfig {
+        &self.config
+    }
+
+    pub fn start(self) -> Result<DownloadHandle, Error> {
+        if self.manager.is_cancelled() {
+            return Err(Error::Download(DownloadError::ManagerShutdown));
+        }
+
+        if self.destination.exists() {
+            return Err(Error::Download(DownloadError::FileExists {
+                path: self.destination,
+            }));
+        }
+
+        let cancel = self.manager.child_token();
+        let (req, handle) = DownloadRequest::new(self.url, self.destination, cancel, self.config);
+
+        self.manager.queue_request(req)?;
+
+        Ok(handle)
+    }
+}

src/downloader/types.rs

@@ -27,4 +27,6 @@ pub enum DownloadError {
     ManagerShutdown,
     #[error("File already exists: {path}")]
     FileExists { path: PathBuf },
+    #[error("Invalid URL")]
+    InvalidUrl,
 }