downloader: Add DownloadConfig and rate limiting

Author: cyberphantom52

src/downloader/config.rs

@@ -1,6 +1,9 @@
-use std::sync::{
-    atomic::{AtomicUsize, Ordering},
-    Arc,
+use std::{
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
 };
 
 #[derive(Debug, Clone)]
@@ -31,3 +34,34 @@ impl DownloadManagerConfig {
         self.max_concurrent.store(max, Ordering::Relaxed);
     }
 }
+
+#[derive(Debug, Clone)]
+pub struct DownloadConfig {
+    max_retries: usize,
+    user_agent: Option<String>,
+    progress_update_interval: Duration,
+}
+
+impl Default for DownloadConfig {
+    fn default() -> Self {
+        Self {
+            max_retries: 3,
+            user_agent: None,
+            progress_update_interval: Duration::from_millis(1000),
+        }
+    }
+}
+
+impl DownloadConfig {
+    pub fn max_retries(&self) -> usize {
+        self.max_retries
+    }
+
+    pub fn user_agent(&self) -> Option<&str> {
+        self.user_agent.as_deref()
+    }
+
+    pub fn progress_update_interval(&self) -> Duration {
+        self.progress_update_interval
+    }
+}

src/downloader/manager.rs

@@ -1,4 +1,6 @@
-use super::{config::DownloadManagerConfig, download_thread, DownloadHandle, DownloadRequest};
+use super::{
+    config::DownloadManagerConfig, download_thread, DownloadConfig, DownloadHandle, DownloadRequest,
+};
 use crate::{error::DownloadError, Error};
 use reqwest::{Client, Url};
 use std::{path::Path, sync::Arc};
@@ -38,10 +40,11 @@ impl DownloadManager {
         manager
     }
 
-    pub fn download(
+    pub fn download_with_config(
         &self,
         url: Url,
         destination: impl AsRef<Path>,
+        config: DownloadConfig,
     ) -> Result<DownloadHandle, Error> {
         if self.cancel.is_cancelled() {
             return Err(Error::Download(DownloadError::ManagerShutdown));
@@ -55,7 +58,7 @@ impl DownloadManager {
         }
 
         let cancel = self.cancel.child_token();
-        let (req, handle) = DownloadRequest::new_req_handle_pair(url, destination, cancel);
+        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),
@@ -65,6 +68,14 @@ impl DownloadManager {
         Ok(handle)
     }
 
+    pub fn download(
+        &self,
+        url: Url,
+        destination: impl AsRef<Path>,
+    ) -> Result<DownloadHandle, Error> {
+        self.download_with_config(url, destination, DownloadConfig::default())
+    }
+
     pub async fn set_max_parallel_downloads(&self, limit: usize) -> Result<(), Error> {
         let current = self.config.max_concurrent();
         if limit > current {

src/downloader/mod.rs

@@ -5,9 +5,10 @@ mod progress;
 mod types;
 mod worker;
 
-pub use config::DownloadManagerConfig;
+pub use config::{DownloadConfig, DownloadManagerConfig};
 pub use handle::DownloadHandle;
 pub use manager::DownloadManager;
 pub use progress::DownloadProgress;
-pub use types::*;
+pub(self) use types::DownloadRequest;
+pub use types::Status;
 pub(self) use worker::download_thread;

src/downloader/progress.rs

@@ -1,5 +1,7 @@
 use std::time::{Duration, Instant};
 
+const SPEED_UPDATE_INTERVAL: Duration = Duration::from_secs(1);
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct DownloadProgress {
     bytes_downloaded: u64,
@@ -12,7 +14,6 @@ pub struct DownloadProgress {
     last_update: Instant,
     last_speed_update: Instant,
     last_bytes_for_speed: u64,
-    update_interval: Duration,
 }
 
 impl std::fmt::Display for DownloadProgress {
@@ -42,7 +43,7 @@ impl std::fmt::Display for DownloadProgress {
 }
 
 impl DownloadProgress {
-    pub fn new(bytes_downloaded: u64, total_bytes: Option<u64>, update_interval: Duration) -> Self {
+    pub fn new(bytes_downloaded: u64, total_bytes: Option<u64>) -> Self {
         let now = Instant::now();
         Self {
             bytes_downloaded,
@@ -54,53 +55,35 @@ impl DownloadProgress {
             last_update: now,
             last_speed_update: now,
             last_bytes_for_speed: bytes_downloaded,
-            update_interval,
         }
     }
 
-    pub fn update(&self, bytes_downloaded: u64) -> Option<Self> {
-        fn new_update(
-            progress: &DownloadProgress,
-            bytes_downloaded: u64,
-            instant: Instant,
-        ) -> DownloadProgress {
-            DownloadProgress {
-                eta: None,
-                last_update: instant,
-                bytes_downloaded,
-                total_bytes: progress.total_bytes,
-                speed_bps: progress.speed_bps,
-                start_time: progress.start_time,
-                last_speed_update: progress.last_speed_update,
-                last_bytes_for_speed: progress.last_bytes_for_speed,
-                update_interval: progress.update_interval,
-            }
-        }
+    pub fn update(&mut self, bytes_downloaded: u64) {
+        self.bytes_downloaded = bytes_downloaded;
+        self.update_speed(bytes_downloaded);
+        self.update_eta();
+    }
 
+    fn update_speed(&mut self, bytes_downloaded: u64) {
         let now = Instant::now();
 
-        if now.duration_since(self.last_update) < self.update_interval {
-            return None;
-        }
-        let mut new_update = new_update(self, bytes_downloaded, now);
-
-        if now.duration_since(self.last_speed_update) >= Duration::from_secs(1) {
+        if now.duration_since(self.last_speed_update) >= SPEED_UPDATE_INTERVAL {
             let byte_diff = (bytes_downloaded - self.last_bytes_for_speed) as f64;
             let time_diff = now.duration_since(self.last_speed_update).as_secs_f64();
 
-            new_update.last_speed_update = now;
-            new_update.last_bytes_for_speed = bytes_downloaded;
-            new_update.speed_bps = Some((byte_diff / time_diff) as u64);
+            self.last_speed_update = now;
+            self.last_bytes_for_speed = bytes_downloaded;
+            self.speed_bps = Some((byte_diff / time_diff) as u64);
         };
+    }
 
-        if let (Some(speed), Some(total)) = (new_update.speed_bps, self.total_bytes) {
+    fn update_eta(&mut self) {
+        if let (Some(speed), Some(total)) = (self.speed_bps, self.total_bytes) {
             if speed > 0 {
-                let remaining = total.saturating_sub(bytes_downloaded);
-                new_update.eta = Some(Duration::from_secs(remaining / speed));
+                let remaining = total.saturating_sub(self.bytes_downloaded);
+                self.eta = Some(Duration::from_secs(remaining / speed));
             }
-        };
-
-        Some(new_update)
+        }
     }
 
     pub fn percent(&self) -> Option<f64> {

src/downloader/types.rs

@@ -1,7 +1,10 @@
-use super::{DownloadHandle, DownloadProgress};
-use crate::Error;
+use super::{DownloadConfig, DownloadHandle, DownloadProgress};
+use crate::{error::DownloadError, Error};
 use reqwest::Url;
-use std::path::{Path, PathBuf};
+use std::{
+    path::{Path, PathBuf},
+    time::Instant,
+};
 use tokio::{
     fs::File,
     sync::{oneshot, watch},
@@ -10,33 +13,58 @@ use tokio_util::sync::CancellationToken;
 
 #[derive(Debug)]
 pub(crate) struct DownloadRequest {
-    pub url: Url,
-    pub destination: PathBuf,
+    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_req_handle_pair(
+    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, Copy, Clone, PartialEq, Eq)]

src/downloader/worker.rs

@@ -4,8 +4,6 @@ use reqwest::Client;
 use std::time::Duration;
 use tokio::{fs::File, io::AsyncWriteExt};
 
-const MAX_RETRIES: usize = 3;
-
 pub(super) async fn download_thread(client: Client, mut req: DownloadRequest) {
     fn should_retry(e: &Error) -> bool {
         match e {
@@ -24,8 +22,9 @@ pub(super) async fn download_thread(client: Client, mut req: DownloadRequest) {
     }
 
     let mut last_error = None;
-    for attempt in 0..=(MAX_RETRIES + 1) {
-        if attempt > MAX_RETRIES {
+    let max_attempts = req.config().max_retries();
+    for attempt in 0..=(max_attempts + 1) {
+        if attempt > max_attempts {
             req.status.send(Status::Failed).ok();
             req.result
                 .send(Err(Error::Download(DownloadError::RetriesExhausted {
@@ -81,44 +80,42 @@ pub(super) async fn download_thread(client: Client, mut req: DownloadRequest) {
 
 async fn download(client: Client, req: &mut DownloadRequest) -> Result<File, Error> {
     let mut response = client
-        .get(req.url.as_ref())
+        .get(req.url().as_ref())
         .send()
         .await?
         .error_for_status()?;
     let total_bytes = response.content_length();
     let mut bytes_downloaded = 0u64;
 
     // Create the destination directory if it doesn't exist
-    if let Some(parent) = req.destination.parent() {
+    if let Some(parent) = req.destination().parent() {
         tokio::fs::create_dir_all(parent).await?;
     }
-    let mut file = File::create(&req.destination).await?;
+    let mut file = File::create(&req.destination()).await?;
 
-    let update_interval = Duration::from_millis(250);
-    let mut progress = DownloadProgress::new(bytes_downloaded, total_bytes, update_interval);
-    req.status.send(Status::InProgress(progress)).ok();
+    let mut progress = DownloadProgress::new(bytes_downloaded, total_bytes);
+    req.send_progress(progress);
 
     loop {
         tokio::select! {
             _ = req.cancel.cancelled() => {
                 drop(file); // Manually drop the file handle to ensure that deletion doesn't fail
-                tokio::fs::remove_file(&req.destination).await?;
+                tokio::fs::remove_file(&req.destination()).await?;
                 return Err(Error::Download(DownloadError::Cancelled));
             }
             chunk = response.chunk() => {
                 match chunk {
                     Ok(Some(chunk)) => {
                         file.write_all(&chunk).await?;
                         bytes_downloaded += chunk.len() as u64;
-                        if let Some(new_progress) = progress.update(bytes_downloaded) {
-                            progress = new_progress;
-                        }
-                        req.status.send(Status::InProgress(progress)).ok();
+
+                        progress.update(bytes_downloaded);
+                        req.send_progress(progress);
                     }
                     Ok(None) => break,
                     Err(e) => {
                         drop(file); // Manually drop the file handle to ensure that deletion doesn't fail
-                        tokio::fs::remove_file(&req.destination).await?;
+                        tokio::fs::remove_file(&req.destination()).await?;
                         return Err(Error::Reqwest(e))
                     },
                 }
@@ -129,6 +126,6 @@ async fn download(client: Client, req: &mut DownloadRequest) -> Result<File, Err
     // Ensure the data is written to disk
     file.sync_all().await?;
     // Open a new file handle with RO permissions
-    let file = File::options().read(true).open(&req.destination).await?;
+    let file = File::options().read(true).open(&req.destination()).await?;
     Ok(file)
 }