downloader: Implement a Download Scheduler

- Also Re-Queue requests on retry

Author: cyberphantom52

crates/download-manager/src/download_manager.rs

@@ -3,9 +3,13 @@ mod download;
 mod error;
 mod events;
 mod request;
-mod worker;
+mod scheduler;
 
-use crate::{context::Context, request::RequestBuilder, worker::download_thread};
+use crate::{
+    context::Context,
+    request::RequestBuilder,
+    scheduler::{Scheduler, SchedulerCmd},
+};
 pub use crate::{
     context::DownloadID,
     download::{Download, DownloadResult},
@@ -24,7 +28,7 @@ use tokio_stream::wrappers::BroadcastStream;
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
 
 pub struct DownloadManager {
-    queue: mpsc::Sender<Request>,
+    scheduler_tx: mpsc::Sender<SchedulerCmd>,
     ctx: Arc<Context>,
     tracker: TaskTracker,
 }
@@ -55,22 +59,6 @@ impl DownloadManager {
         Request::builder(self)
     }
 
-    fn queue_request(&self, request: Request) -> Result<(), DownloadError> {
-        self.queue.try_send(request).map_err(|e| match e {
-            mpsc::error::TrySendError::Full(_) => DownloadError::QueueFull,
-            mpsc::error::TrySendError::Closed(_) => DownloadError::ManagerShutdown,
-        })
-    }
-
-    /// Returns the count of pending requests still buffered in the internal mpsc channel.
-    ///
-    /// **Note**: This excludes any request already dequeued by the dispatcher but not yet started.
-    ///
-    /// Consider replacing with an explicit atomic counter.
-    pub fn queued_downloads(&self) -> usize {
-        self.queue.max_capacity() - self.queue.capacity()
-    }
-
     pub fn active_downloads(&self) -> usize {
         self.ctx.active.load(Ordering::Relaxed)
     }
@@ -129,54 +117,16 @@ impl DownloadManagerBuilder {
         let (tx, rx) = mpsc::channel(queue_size);
         let ctx = Context::new(max_concurrent);
         let tracker = TaskTracker::new();
+        let scheduler = Scheduler::new(ctx.clone(), tracker.clone(), rx);
 
         let manager = DownloadManager {
-            queue: tx,
+            scheduler_tx: tx,
             ctx: ctx.clone(),
             tracker: tracker.clone(),
         };
 
-        tracker.spawn(dispatcher_thread(rx, tracker.clone(), ctx));
+        tracker.spawn(async move { scheduler.run().await });
 
         Ok(manager)
     }
 }
-
-async fn dispatcher_thread(
-    mut rx: mpsc::Receiver<Request>,
-    tracker: TaskTracker,
-    ctx: Arc<Context>,
-) {
-    struct ActiveGuard {
-        ctx: Arc<Context>,
-        _permit: tokio::sync::OwnedSemaphorePermit,
-    }
-
-    impl Drop for ActiveGuard {
-        fn drop(&mut self) {
-            self.ctx
-                .active
-                .fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
-        }
-    }
-
-    while let Some(request) = rx.recv().await {
-        let guard = match ctx.semaphore.clone().acquire_owned().await {
-            Ok(p) => {
-                ctx.active.fetch_add(1, Ordering::Relaxed);
-                ActiveGuard {
-                    ctx: ctx.clone(),
-                    _permit: p,
-                }
-            }
-            Err(_) => break,
-        };
-
-        let ctx_clone = ctx.clone();
-        tracker.spawn(async move {
-            // Move the guard into the worker thread so it's automatically released when the thread finishes
-            let _guard = guard;
-            download_thread(request, ctx_clone).await
-        });
-    }
-}

crates/download-manager/src/request.rs

@@ -1,15 +1,16 @@
 use crate::{
-    Download, DownloadError, DownloadEvent, DownloadID, DownloadManager, DownloadResult, Progress,
+    Download, DownloadEvent, DownloadID, DownloadManager, Progress, scheduler::SchedulerCmd,
 };
 use derive_builder::Builder;
-use reqwest::Url;
-use std::{
-    path::{Path, PathBuf},
-    time::Duration,
+use reqwest::{
+    Url,
+    header::{HeaderMap, IntoHeaderName},
 };
+use std::path::{Path, PathBuf};
 use tokio::sync::{broadcast, oneshot, watch};
 use tokio_util::sync::CancellationToken;
 
+#[derive(Debug, Clone)]
 pub struct Request {
     id: DownloadID,
     url: Url,
@@ -18,28 +19,27 @@ pub struct Request {
 
     progress: watch::Sender<Progress>,
     events: broadcast::Sender<DownloadEvent>,
-    result: oneshot::Sender<Result<DownloadResult, DownloadError>>,
 
     pub cancel_token: CancellationToken,
 }
 
-#[derive(Debug, Builder)]
+#[derive(Debug, Builder, Clone)]
 #[builder(pattern = "owned")]
 pub struct DownloadConfig {
     #[builder(default = "3")]
     retries: u32,
-    #[builder(default, setter(strip_option))]
-    user_agent: Option<String>,
     #[builder(default = "false")]
     overwrite: bool,
+    #[builder(default = "HeaderMap::new()", setter(skip))]
+    headers: HeaderMap,
 }
 
 impl Default for DownloadConfig {
     fn default() -> Self {
         DownloadConfig {
             retries: 3,
-            user_agent: None,
             overwrite: false,
+            headers: HeaderMap::new(),
         }
     }
 }
@@ -49,13 +49,13 @@ impl DownloadConfig {
         self.retries
     }
 
-    pub fn user_agent(&self) -> Option<&str> {
-        self.user_agent.as_deref()
-    }
-
     pub fn overwrite(&self) -> bool {
         self.overwrite
     }
+
+    pub fn headers(&self) -> &HeaderMap {
+        &self.headers
+    }
 }
 
 impl Request {
@@ -64,6 +64,7 @@ impl Request {
             url: None,
             destination: None,
             config: DownloadConfigBuilder::default(),
+            headers: HeaderMap::new(),
             manager,
         }
     }
@@ -87,20 +88,11 @@ impl Request {
         &self.config
     }
 
-    pub fn is_cancelled(&self) -> bool {
-        self.cancel_token.is_cancelled()
-    }
-
-    fn emit(&self, event: DownloadEvent) {
+    pub fn emit(&self, event: DownloadEvent) {
         // TODO: Log the error
         let _ = self.events.send(event);
     }
 
-    fn send_result(self, result: Result<DownloadResult, DownloadError>) {
-        // TODO: Log the error
-        let _ = self.result.send(result);
-    }
-
     pub fn update_progress(&self, progress: Progress) {
         // TODO: Log the error
         let _ = self.progress.send(progress);
@@ -110,42 +102,17 @@ impl Request {
         self.emit(DownloadEvent::Started {
             id: self.id(),
             url: self.url().clone(),
-            destination: self.destination.clone(),
+            destination: self.destination().to_path_buf(),
             total_bytes: None,
         });
     }
-
-    pub fn fail(self, error: DownloadError) {
-        self.send_result(Err(error));
-    }
-
-    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 retry(&self, attempt: u32, delay: Duration) {
-        self.emit(DownloadEvent::Retrying {
-            id: self.id(),
-            attempt,
-            next_delay_ms: delay.as_millis() as u64,
-        });
-    }
-
-    pub fn cancel(self) {
-        self.emit(DownloadEvent::Cancelled { id: self.id() });
-        self.send_result(Err(DownloadError::Cancelled))
-    }
 }
 
 pub struct RequestBuilder<'a> {
     url: Option<Url>,
     destination: Option<PathBuf>,
     config: DownloadConfigBuilder,
+    headers: HeaderMap,
 
     manager: &'a DownloadManager,
 }
@@ -166,47 +133,48 @@ impl RequestBuilder<'_> {
         self
     }
 
-    pub fn user_agent(mut self, user_agent: impl AsRef<str>) -> Self {
-        self.config = self.config.user_agent(user_agent.as_ref().into());
-        self
+    pub fn user_agent(self, user_agent: impl AsRef<str>) -> Self {
+        self.header(reqwest::header::USER_AGENT, user_agent)
     }
 
     pub fn overwrite(mut self, overwrite: bool) -> Self {
         self.config = self.config.overwrite(overwrite);
         self
     }
 
+    pub fn header(mut self, header: impl IntoHeaderName, value: impl AsRef<str>) -> Self {
+        self.headers.insert(header, value.as_ref().parse().unwrap());
+        self
+    }
+
     pub fn start(self) -> anyhow::Result<Download> {
         let url = self.url.ok_or_else(|| anyhow::anyhow!("URL must be set"))?;
         let destination = self
             .destination
             .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 id = self.manager.ctx.next_id();
         let (result_tx, result_rx) = oneshot::channel();
+        let (progress_tx, progress_rx) = watch::channel(Progress::new(None));
         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: url.clone(),
             destination: destination.clone(),
             config,
+
+            events: event_tx,
             progress: progress_tx,
-            events: event_tx.clone(),
-            result: result_tx,
             cancel_token: cancel_token.clone(),
         };
 
-        self.manager.queue_request(request)?;
-        event_tx.send(DownloadEvent::Queued {
-            id,
-            url,
-            destination,
-        })?;
+        self.manager
+            .scheduler_tx
+            .try_send(SchedulerCmd::Enqueue { request, result_tx });
 
         Ok(Download::new(
             id,