Draft: Download manager

Cargo.toml

@@ -11,6 +11,9 @@ tracing = { workspace = true }
 tonic = { workspace = true }
 tokio = { workspace = true }
 prost = { workspace = true }
+reqwest = "0.12"
+tokio-util = { version = "0.7", features = ["rt"] }
+url = "2.5.4"
 
 [build-dependencies]
 tonic-build = "0.12.3"

src/downloader/config.rs

@@ -0,0 +1,82 @@
+use std::{
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+#[derive(Debug, Clone)]
+pub struct DownloadManagerConfig {
+    max_concurrent: Arc<AtomicUsize>,
+    queue_size: usize,
+}
+
+impl Default for DownloadManagerConfig {
+    fn default() -> Self {
+        Self {
+            max_concurrent: Arc::new(AtomicUsize::new(3)),
+            queue_size: 100,
+        }
+    }
+}
+
+impl DownloadManagerConfig {
+    pub fn queue_size(&self) -> usize {
+        self.queue_size
+    }
+
+    pub fn max_concurrent(&self) -> usize {
+        self.max_concurrent.load(Ordering::Relaxed)
+    }
+
+    pub fn set_max_concurrent(&self, max: usize) {
+        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 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
+    }
+
+    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/handle.rs

@@ -0,0 +1,65 @@
+use super::Status;
+use crate::Error;
+use tokio::{
+    fs::File,
+    sync::{oneshot, watch},
+};
+use tokio_util::sync::CancellationToken;
+
+#[derive(Debug)]
+pub struct DownloadHandle {
+    result: oneshot::Receiver<Result<File, Error>>,
+    status: watch::Receiver<Status>,
+    cancel: CancellationToken,
+}
+
+impl DownloadHandle {
+    pub fn new(
+        result: oneshot::Receiver<Result<File, Error>>,
+        status: watch::Receiver<Status>,
+        cancel: CancellationToken,
+    ) -> Self {
+        Self {
+            result,
+            status,
+            cancel,
+        }
+    }
+
+    pub fn status(&self) -> Status {
+        *self.status.borrow()
+    }
+
+    pub fn is_completed(&self) -> bool {
+        matches!(self.status(), Status::Completed)
+    }
+
+    pub fn is_cancelled(&self) -> bool {
+        matches!(self.status(), Status::Cancelled)
+    }
+
+    pub async fn wait_for_status_update(&mut self) -> Result<(), watch::error::RecvError> {
+        self.status.changed().await
+    }
+
+    pub fn cancel(&self) {
+        self.cancel.cancel();
+    }
+}
+
+impl std::future::Future for DownloadHandle {
+    type Output = Result<tokio::fs::File, Error>;
+
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        use std::pin::Pin;
+        use std::task::Poll;
+        match Pin::new(&mut self.result).poll(cx) {
+            Poll::Ready(Ok(result)) => Poll::Ready(result),
+            Poll::Ready(Err(e)) => Poll::Ready(Err(Error::Oneshot(e))),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}

src/downloader/manager.rs

@@ -0,0 +1,133 @@
+use super::{
+    worker::download_thread, DownloadBuilder, DownloadConfig, DownloadManagerConfig,
+    DownloadRequest,
+};
+use crate::{error::DownloadError, Error};
+use reqwest::Client;
+use std::{path::Path, sync::Arc};
+use tokio::sync::{mpsc, Semaphore};
+use tokio_util::{sync::CancellationToken, task::TaskTracker};
+use url::Url;
+
+#[derive(Debug)]
+pub struct DownloadManager {
+    queue: mpsc::Sender<DownloadRequest>,
+    semaphore: Arc<Semaphore>,
+    cancel: CancellationToken,
+    config: DownloadManagerConfig,
+    tracker: TaskTracker,
+}
+
+impl Default for DownloadManager {
+    fn default() -> Self {
+        Self::with_config(DownloadManagerConfig::default())
+    }
+}
+
+impl DownloadManager {
+    pub fn with_config(config: DownloadManagerConfig) -> Self {
+        let (tx, rx) = mpsc::channel(config.queue_size());
+        let client = Client::new();
+        let semaphore = Arc::new(Semaphore::new(config.max_concurrent()));
+        let tracker = TaskTracker::new();
+        let manager = Self {
+            queue: tx,
+            semaphore: semaphore.clone(),
+            cancel: CancellationToken::new(),
+            config,
+            tracker: tracker.clone(),
+        };
+        // Spawn the dispatcher thread to handle download requests
+        tracker.spawn(dispatcher_thread(client, rx, semaphore, tracker.clone()));
+        manager
+    }
+
+    pub fn download(&self, url: Url, destination: impl AsRef<Path>) -> DownloadBuilder {
+        DownloadBuilder::new(self, url, destination)
+    }
+
+    pub fn download_with_config(
+        &self,
+        url: Url,
+        destination: impl AsRef<Path>,
+        config: DownloadConfig,
+    ) -> DownloadBuilder {
+        self.download(url, destination).with_config(config)
+    }
+
+    pub async fn set_max_parallel_downloads(&self, limit: usize) -> Result<(), Error> {
+        let current = self.config.max_concurrent();
+        if limit > current {
+            self.semaphore.add_permits(limit - current);
+        } else if limit < current {
+            let to_remove = current - limit;
+
+            let permits = self
+                .semaphore
+                .acquire_many(to_remove as u32)
+                .await
+                .map_err(|_| Error::Download(DownloadError::ManagerShutdown))?;
+
+            permits.forget();
+        }
+        self.config.set_max_concurrent(limit);
+
+        Ok(())
+    }
+
+    pub fn cancel_all(&self) {
+        self.cancel.cancel();
+    }
+
+    pub fn queued_downloads(&self) -> usize {
+        self.queue.max_capacity() - self.queue.capacity()
+    }
+
+    pub fn active_downloads(&self) -> usize {
+        // -1 because the dispatcher thread is always running
+        self.tracker.len() - 1
+    }
+
+    pub async fn shutdown(self) -> Result<(), Error> {
+        self.cancel.cancel();
+        self.tracker.close();
+        self.tracker.wait().await;
+        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(
+    client: Client,
+    mut rx: mpsc::Receiver<DownloadRequest>,
+    sem: Arc<Semaphore>,
+    tracker: TaskTracker,
+) {
+    while let Some(request) = rx.recv().await {
+        let permit = match sem.clone().acquire_owned().await {
+            Ok(permit) => permit,
+            Err(_) => break,
+        };
+        let client = client.clone();
+        tracker.spawn(async move {
+            // Move the permit into the worker thread so it's automatically released when the thread finishes
+            let _permit = permit;
+            download_thread(client.clone(), request).await;
+        });
+    }
+}

src/downloader/mod.rs

@@ -0,0 +1,12 @@
+mod config;
+mod handle;
+mod manager;
+mod progress;
+mod request;
+mod worker;
+
+pub use config::{DownloadConfig, DownloadManagerConfig};
+pub use handle::DownloadHandle;
+pub use manager::DownloadManager;
+pub use progress::{DownloadProgress, Status};
+pub use request::{DownloadBuilder, DownloadRequest};