jellytau_lib/download/
worker.rs1use log::warn;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::time::Duration;
6
7use futures_util::StreamExt;
8use tokio::fs;
9use tokio::io::AsyncWriteExt;
10
11use super::DownloadTask;
12
13pub struct DownloadWorker {
15 client: reqwest::Client,
17 max_retries: u32,
19}
20
21impl DownloadWorker {
22 pub fn new() -> Self {
23 let client = reqwest::Client::builder()
24 .timeout(Duration::from_secs(300)) .https_only(true)
26 .build()
27 .expect("Failed to create HTTP client");
28
29 Self {
30 client,
31 max_retries: 3,
32 }
33 }
34
35 pub async fn download<F>(
44 &self,
45 task: &DownloadTask,
46 stop: &AtomicBool,
47 on_progress: F,
48 ) -> Result<DownloadResult, DownloadError>
49 where
50 F: Fn(u64, Option<u64>) + Send + Sync,
51 {
52 let mut retries = 0;
53
54 loop {
55 if stop.load(Ordering::SeqCst) {
56 return Err(DownloadError::Stopped);
57 }
58 match self.try_download(task, stop, &on_progress).await {
59 Ok(result) => return Ok(result),
60 Err(e) if retries < self.max_retries && e.is_retryable() => {
61 retries += 1;
62 let delay = Self::exponential_backoff(retries);
63 warn!(
64 "Download failed (attempt {}/{}), retrying in {:?}: {}",
65 retries, self.max_retries, delay, e
66 );
67 tokio::time::sleep(delay).await;
68 }
69 Err(e) => return Err(e),
70 }
71 }
72 }
73
74 async fn try_download<F>(
76 &self,
77 task: &DownloadTask,
78 stop: &AtomicBool,
79 on_progress: &F,
80 ) -> Result<DownloadResult, DownloadError>
81 where
82 F: Fn(u64, Option<u64>) + Send + Sync,
83 {
84 if let Some(parent) = task.target_path.parent() {
86 fs::create_dir_all(parent)
87 .await
88 .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
89 }
90
91 let temp_path = partial_path(&task.target_path);
93 let existing_bytes = if temp_path.exists() {
94 fs::metadata(&temp_path).await.map(|m| m.len()).unwrap_or(0)
95 } else {
96 0
97 };
98
99 let mut request = self.client.get(&task.url);
101 if existing_bytes > 0 {
102 request = request.header("Range", format!("bytes={}-", existing_bytes));
103 }
104
105 let response = request
107 .send()
108 .await
109 .map_err(|e| DownloadError::Network(e.to_string()))?;
110
111 if !response.status().is_success() && response.status().as_u16() != 206 {
113 return Err(DownloadError::Http(response.status().as_u16()));
114 }
115
116 let resume_from = resume_offset(existing_bytes, response.status().as_u16());
120 if existing_bytes > 0 && resume_from == 0 {
121 warn!(
122 "Server ignored the Range request (HTTP {}) — restarting {} from the beginning \
123 instead of appending to {} existing bytes",
124 response.status().as_u16(),
125 task.target_path.display(),
126 existing_bytes
127 );
128 }
129
130 let _total_bytes = response
133 .headers()
134 .get(reqwest::header::CONTENT_LENGTH)
135 .and_then(|v| v.to_str().ok())
136 .and_then(|v| v.parse::<u64>().ok())
137 .map(|len| len + resume_from);
138
139 let mut file = if resume_from > 0 {
142 fs::OpenOptions::new().append(true).open(&temp_path).await
143 } else {
144 fs::File::create(&temp_path).await
145 }
146 .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
147
148 let mut downloaded = resume_from;
150 let mut stream = response.bytes_stream();
151 let mut last_progress_emit = std::time::Instant::now();
152
153 while let Some(chunk) = stream.next().await {
154 if stop.load(Ordering::SeqCst) {
160 let _ = file.flush().await;
161 let _ = file.sync_all().await;
162 return Err(DownloadError::Stopped);
163 }
164
165 let chunk = chunk.map_err(|e| DownloadError::Network(e.to_string()))?;
166
167 file.write_all(&chunk)
168 .await
169 .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
170
171 downloaded += chunk.len() as u64;
172
173 if last_progress_emit.elapsed() > Duration::from_millis(500)
175 || downloaded.is_multiple_of(1024 * 1024)
176 {
177 last_progress_emit = std::time::Instant::now();
178 on_progress(downloaded, _total_bytes);
179 }
180 }
181
182 file.sync_all()
183 .await
184 .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
185
186 fs::rename(&temp_path, &task.target_path)
188 .await
189 .map_err(|e| DownloadError::FileSystem(e.to_string()))?;
190
191 Ok(DownloadResult {
192 bytes_downloaded: downloaded,
193 })
194 }
195
196 fn exponential_backoff(retry_count: u32) -> Duration {
198 let base_delay = 5; let delay_secs = base_delay * 3u64.pow(retry_count - 1); Duration::from_secs(delay_secs)
201 }
202}
203
204pub fn resume_offset(existing_bytes: u64, status: u16) -> u64 {
221 if existing_bytes == 0 {
222 return 0;
223 }
224 if status == 206 {
226 existing_bytes
227 } else {
228 0
229 }
230}
231
232pub fn partial_path(target: &std::path::Path) -> std::path::PathBuf {
249 let mut s = target.as_os_str().to_os_string();
250 s.push(".part");
251 std::path::PathBuf::from(s)
252}
253
254#[derive(Debug)]
256pub struct DownloadResult {
257 pub bytes_downloaded: u64,
258}
259
260#[derive(Debug)]
262pub enum DownloadError {
263 Network(String),
264 Http(u16),
265 FileSystem(String),
266 Stopped,
270}
271
272impl DownloadError {
273 fn is_retryable(&self) -> bool {
275 match self {
276 DownloadError::Network(_) => true,
277 DownloadError::Http(status) => *status >= 500, DownloadError::FileSystem(_) => false,
279 DownloadError::Stopped => false,
281 }
282 }
283
284 pub fn is_stopped(&self) -> bool {
287 matches!(self, DownloadError::Stopped)
288 }
289}
290
291impl std::fmt::Display for DownloadError {
292 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293 match self {
294 DownloadError::Network(msg) => write!(f, "Network error: {}", msg),
295 DownloadError::Http(status) => write!(f, "HTTP error {}", status),
296 DownloadError::FileSystem(msg) => write!(f, "File system error: {}", msg),
297 DownloadError::Stopped => write!(f, "Download stopped by request"),
298 }
299 }
300}
301
302impl std::error::Error for DownloadError {}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307
308 #[test]
315 fn test_resume_offset_only_appends_when_the_server_honoured_the_range() {
316 assert_eq!(resume_offset(0, 200), 0);
318 assert_eq!(resume_offset(0, 206), 0);
319
320 assert_eq!(resume_offset(5_000, 206), 5_000);
322
323 assert_eq!(
326 resume_offset(5_000, 200),
327 0,
328 "a 200 carries the whole stream; appending it corrupts the file"
329 );
330 }
331
332 #[test]
338 fn test_partial_path_appends_rather_than_replacing_the_extension() {
339 use std::path::Path;
340
341 assert_eq!(
342 partial_path(Path::new("/media/movie.mp4")),
343 Path::new("/media/movie.mp4.part"),
344 "the cleanup paths delete \"{{file_path}}.part\"; this must produce it"
345 );
346
347 assert_ne!(
349 partial_path(Path::new("/media/movie.mp4")),
350 partial_path(Path::new("/media/movie.mkv")),
351 );
352
353 assert_eq!(
355 partial_path(Path::new("/media/track")),
356 Path::new("/media/track.part"),
357 );
358
359 assert_eq!(
361 partial_path(Path::new("/media/S01.E02.episode.mkv")),
362 Path::new("/media/S01.E02.episode.mkv.part"),
363 );
364 }
365
366 #[test]
367 fn test_exponential_backoff() {
368 assert_eq!(
369 DownloadWorker::exponential_backoff(1),
370 Duration::from_secs(5)
371 );
372 assert_eq!(
373 DownloadWorker::exponential_backoff(2),
374 Duration::from_secs(15)
375 );
376 assert_eq!(
377 DownloadWorker::exponential_backoff(3),
378 Duration::from_secs(45)
379 );
380 }
381
382 #[test]
383 fn test_error_retryable() {
384 assert!(DownloadError::Network("timeout".to_string()).is_retryable());
385 assert!(DownloadError::Http(500).is_retryable());
386 assert!(DownloadError::Http(503).is_retryable());
387 assert!(!DownloadError::Http(404).is_retryable());
388 assert!(!DownloadError::FileSystem("disk full".to_string()).is_retryable());
389 assert!(!DownloadError::Stopped.is_retryable());
391 assert!(DownloadError::Stopped.is_stopped());
392 assert!(!DownloadError::Network("timeout".to_string()).is_stopped());
393 }
394}