Skip to main content

ninja/utils/
download.rs

1use bytes::Bytes;
2use futures_util::StreamExt;
3use std::time::Instant;
4use tokio::sync::mpsc;
5
6pub struct Downloader {
7    client: reqwest::Client,
8}
9impl Downloader {
10    pub fn new() -> Self {
11        Self {
12            client: reqwest::Client::new(),
13        }
14    }
15
16    pub async fn download<S: DownloadTarget>(
17        &self,
18        request: DownloadRequest,
19        mut target: S,
20        events: EventSender,
21    ) -> Result<(), DownloadError> {
22        let response = self
23            .client
24            .request(request.method, &request.url)
25            .headers(request.headers)
26            .send()
27            .await?;
28
29        let total = response.content_length();
30        let _ = events.send(DownloadEvent::Started { total });
31
32        let mut stream = response.bytes_stream();
33        let start = Instant::now();
34        let mut downloaded = 0u64;
35
36        while let Some(chunk) = stream.next().await {
37            let chunk = chunk?;
38            downloaded += chunk.len() as u64;
39
40            // STREAM → TARGET (FILE, MEMORY, ETC.)
41            target.write(&chunk).await?;
42
43            // PROGRESS
44            let elapsed = start.elapsed().as_secs_f64().max(0.001);
45            let speed = downloaded as f64 / elapsed;
46            let eta = total.map(|t| (t.saturating_sub(downloaded)) as f64 / speed);
47
48            let _ = events.send(DownloadEvent::Progress {
49                downloaded,
50                total,
51                speed_bps: speed,
52                eta_secs: eta,
53            });
54        }
55
56        target.finish().await?;
57        let _ = events.send(DownloadEvent::Finished);
58
59        Ok(())
60    }
61}
62
63impl Default for Downloader {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69//
70// ===== Request =====
71//
72
73use reqwest::Method;
74use reqwest::header::HeaderMap;
75
76pub struct DownloadRequest {
77    pub url: String,
78    pub method: Method,
79    pub headers: HeaderMap,
80}
81
82impl DownloadRequest {
83    pub fn get(url: impl Into<String>) -> Self {
84        Self {
85            url: url.into(),
86            method: Method::GET,
87            headers: HeaderMap::new(),
88        }
89    }
90}
91
92//
93// ===== Events =====
94//
95
96#[derive(Debug, Clone)]
97pub enum DownloadEvent {
98    Started {
99        total: Option<u64>,
100    },
101    Progress {
102        downloaded: u64,
103        total: Option<u64>,
104        speed_bps: f64,
105        eta_secs: Option<f64>,
106    },
107    Finished,
108    Failed(String),
109}
110
111pub type EventSender = mpsc::UnboundedSender<DownloadEvent>;
112
113//
114// ===== Target (Sink) =====
115//
116
117#[async_trait::async_trait]
118pub trait DownloadTarget: Send {
119    async fn write(&mut self, chunk: &Bytes) -> std::io::Result<()>;
120    async fn finish(&mut self) -> std::io::Result<()>;
121}
122
123//
124// ===== File target (REAL streaming) =====
125//
126
127use tokio::fs::File;
128use tokio::io::AsyncWriteExt;
129
130pub struct FileTarget {
131    file: File,
132}
133
134impl FileTarget {
135    pub async fn create(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
136        Ok(Self {
137            file: File::create(path).await?,
138        })
139    }
140}
141
142#[async_trait::async_trait]
143impl DownloadTarget for FileTarget {
144    async fn write(&mut self, chunk: &Bytes) -> std::io::Result<()> {
145        self.file.write_all(chunk).await
146    }
147
148    async fn finish(&mut self) -> std::io::Result<()> {
149        self.file.flush().await
150    }
151}
152
153//
154// ===== Errors =====
155//
156
157#[derive(Debug)]
158pub enum DownloadError {
159    Http(reqwest::Error),
160    Io(std::io::Error),
161}
162
163impl From<reqwest::Error> for DownloadError {
164    fn from(e: reqwest::Error) -> Self {
165        Self::Http(e)
166    }
167}
168
169impl From<std::io::Error> for DownloadError {
170    fn from(e: std::io::Error) -> Self {
171        Self::Io(e)
172    }
173}