1use std::sync::Arc;
2use std::time::Instant;
3
4use futures::StreamExt;
5use futures::stream::FuturesUnordered;
6use reqwest;
7
8use crate::constants;
9use crate::engine::concurrent_segment_manager::ConcurrentSegmentManager;
10use crate::engine::download_cookie::CookieHelper;
11use crate::engine::download_progress::ProgressUpdater;
12use crate::engine::http_segment_downloader::HttpSegmentDownloader;
13use crate::error::{Aria2Error, RecoverableError, Result};
14use crate::filesystem::disk_writer::{CachedDiskWriter, SeekableDiskWriter};
15use crate::filesystem::resume_helper::ResumeState;
16use crate::rate_limiter::{RateLimiter, RateLimiterConfig};
17use crate::request::request_group::RequestGroup;
18
19type SegmentFetchFuture = std::pin::Pin<
20 Box<
21 dyn std::future::Future<
22 Output = (
23 u32,
24 std::result::Result<bytes::Bytes, crate::error::Aria2Error>,
25 ),
26 > + Send,
27 >,
28>;
29
30pub enum ConcurrentDownloadResult {
31 Complete,
32 Fallback { completed_ranges: Vec<(u64, u64)> },
33}
34
35pub struct ConcurrentDownloader {
36 client: Arc<reqwest::Client>,
37 output_path: std::path::PathBuf,
38 headers: Vec<(String, String)>,
39 cookie_helper: CookieHelper,
40 progress_updater: ProgressUpdater,
41 group: Arc<tokio::sync::RwLock<RequestGroup>>,
42 mmap_threshold: u64,
43 file_allocation: String,
44}
45
46impl ConcurrentDownloader {
47 #[allow(clippy::too_many_arguments)]
48 pub fn new(
49 client: Arc<reqwest::Client>,
50 output_path: std::path::PathBuf,
51 headers: Vec<(String, String)>,
52 cookie_helper: CookieHelper,
53 progress_updater: ProgressUpdater,
54 group: Arc<tokio::sync::RwLock<RequestGroup>>,
55 mmap_threshold: u64,
56 file_allocation: String,
57 ) -> Self {
58 Self {
59 client,
60 output_path,
61 headers,
62 cookie_helper,
63 progress_updater,
64 group,
65 mmap_threshold,
66 file_allocation,
67 }
68 }
69
70 pub async fn execute(
71 &mut self,
72 uri: &str,
73 total_length: u64,
74 resume_state: &ResumeState,
75 max_retries_per_segment: u32,
76 ) -> Result<ConcurrentDownloadResult> {
77 {
78 let mut g = self.group.write().await;
79 g.set_total_length(total_length).await;
80 g.set_total_length_atomic(total_length);
81 }
82
83 let options = self.group.read().await.options().clone();
84 let split = options.split.unwrap_or(1) as usize;
85 let max_conn = options
86 .max_connection_per_server
87 .unwrap_or(constants::DEFAULT_MAX_CONNECTION_PER_SERVER as u16)
88 as usize;
89 let seg_size = total_length / split as u64;
90
91 tracing::info!(
92 "Concurrent download started: split={}, max_conn={}, segment_size={} bytes, total={}",
93 split,
94 max_conn,
95 seg_size,
96 total_length
97 );
98
99 let mut manager =
100 ConcurrentSegmentManager::new(total_length, vec![uri.to_string()], Some(seg_size));
101 manager.set_max_connections_per_mirror(max_conn.min(split));
102 manager.set_max_retries(max_retries_per_segment);
103
104 let mut consecutive_416_count = 0u32;
105 let mut total_416_count = 0u32;
106 let fallback_threshold_consecutive = 3u32;
107 let fallback_threshold_ratio = 0.2f64;
108 let mut should_fallback = false;
109
110 if resume_state.should_resume {
111 manager.mark_completed_up_to(resume_state.start_offset, resume_state.existing_length);
112 self.progress_updater.reset(resume_state.start_offset);
113 tracing::debug!(
114 "Resume: marked {} bytes as completed, continuing from offset {}",
115 resume_state.existing_length,
116 resume_state.start_offset
117 );
118 } else {
119 self.progress_updater.reset(0);
120 }
121
122 let cookie_hdr = self.cookie_helper.build_cookie_header(uri);
123
124 let use_mmap = self.file_allocation == "mmap" && total_length >= self.mmap_threshold;
125 let mut writer =
126 CachedDiskWriter::new_with_mmap(&self.output_path, Some(total_length), None, use_mmap);
127
128 let limiter = options
129 .max_download_limit
130 .filter(|&r| r > 0)
131 .map(|r| RateLimiter::new(&RateLimiterConfig::new(Some(r), None)));
132 if let Some(ref limiter) = limiter {
133 let g = self.group.read().await;
134 g.set_rate_limiter(limiter.clone()).await;
135 }
136
137 let mut active: FuturesUnordered<SegmentFetchFuture> = FuturesUnordered::new();
138 let mut active_segs: std::collections::HashMap<u32, u64> = std::collections::HashMap::new();
139 let mut completed_bytes = if resume_state.should_resume {
140 resume_state.start_offset
141 } else {
142 0
143 };
144
145 loop {
146 while active.len() < max_conn {
147 match manager.next_pending_segment_for_mirror(0) {
148 Some((seg_idx, offset, length)) => {
149 let url = uri.to_string();
150 let dl = HttpSegmentDownloader::new(&self.client);
151 let ch = cookie_hdr.clone();
152 let headers = self.headers.clone();
153 active_segs.insert(seg_idx, offset);
154 let fut = Box::pin(async move {
155 let result = dl
156 .download_range(&url, offset, length, ch.as_deref(), &headers)
157 .await;
158 (seg_idx, result)
159 });
160 active.push(fut);
161 tracing::debug!(
162 seg_idx = seg_idx,
163 offset = offset,
164 length = length,
165 "Spawned segment fetch"
166 );
167 }
168 None => break,
169 }
170 }
171
172 if active.is_empty() {
173 if manager.is_complete() {
174 tracing::debug!("All segments complete");
175 break;
176 }
177 if manager.has_failed_segments() && !manager.has_pending_segments() {
178 return Err(Aria2Error::Recoverable(
179 RecoverableError::TemporaryNetworkFailure {
180 message: "Concurrent download: all segments failed".into(),
181 },
182 ));
183 }
184 tracing::warn!(
185 "Concurrent download stuck: no active or pending segments but not complete"
186 );
187 break;
188 }
189
190 if let Some((seg_idx, result)) = active.next().await {
191 let offset = active_segs.remove(&seg_idx).unwrap_or(0);
192 match result {
193 Ok(data) => {
194 let data_len = data.len();
195 if let Some(ref lim) = limiter {
196 lim.acquire_download(data_len as u64).await;
197 }
198 let data_for_manager = data.clone();
199 writer.write_bytes_at(offset, data).await.map_err(|e| {
200 Aria2Error::Fatal(crate::error::FatalError::Config(format!(
201 "Write failed: {}",
202 e
203 )))
204 })?;
205 manager.complete_segment(seg_idx, data_for_manager);
206 completed_bytes += data_len as u64;
207
208 self.progress_updater
209 .update_progress(
210 completed_bytes,
211 constants::PROGRESS_UPDATE_BYTES as u64,
212 constants::HTTP_SPEED_UPDATE_INTERVAL_MS,
213 )
214 .await;
215 }
216 Err(e) => {
217 tracing::warn!(seg_idx = seg_idx, error = %e, "Segment download failed");
218 let is_416 = matches!(
219 &e,
220 Aria2Error::Recoverable(RecoverableError::RangeNotSatisfiable { .. })
221 );
222 if is_416 {
223 consecutive_416_count += 1;
224 total_416_count += 1;
225 tracing::warn!(
226 seg_idx = seg_idx,
227 consecutive_416 = consecutive_416_count,
228 total_416 = total_416_count,
229 "RangeNotSatisfiable (416) detected"
230 );
231 let failure_ratio = total_416_count as f64 / split as f64;
232 let threshold_exceeded = consecutive_416_count
233 >= fallback_threshold_consecutive
234 || failure_ratio >= fallback_threshold_ratio;
235 if threshold_exceeded {
236 tracing::warn!(
237 uri = uri,
238 consecutive_416 = consecutive_416_count,
239 failure_ratio = failure_ratio,
240 "Fallback to sequential mode triggered due to RangeNotSatisfiable errors"
241 );
242 should_fallback = true;
243 break;
244 }
245 } else {
246 consecutive_416_count = 0;
247 }
248 manager.fail_segment(seg_idx);
249 }
250 }
251 }
252 }
253
254 writer.flush().await.map_err(|e| {
255 Aria2Error::Fatal(crate::error::FatalError::Config(format!(
256 "Flush failed: {}",
257 e
258 )))
259 })?;
260
261 if should_fallback {
262 let completed_ranges = manager.completed_ranges();
263 tracing::warn!(
264 "Fallback: {} completed ranges will be preserved",
265 completed_ranges.len()
266 );
267 return Ok(ConcurrentDownloadResult::Fallback { completed_ranges });
268 }
269
270 let final_speed = {
271 let g = self.group.read().await;
272 let elapsed = g.elapsed_time().await;
273 match elapsed {
274 Some(d) if d.as_secs_f64() > 0.0 => {
275 (completed_bytes as f64 / d.as_secs_f64()) as u64
276 }
277 _ => 0,
278 }
279 };
280 {
281 let mut g = self.group.write().await;
282 g.update_progress(completed_bytes).await;
283 g.update_speed(final_speed, 0).await;
284 g.set_completed_length(completed_bytes);
285 g.set_download_speed_cached(final_speed);
286 g.complete().await?;
287 }
288
289 tracing::info!(
290 "Concurrent download complete: {} ({} bytes)",
291 self.output_path.display(),
292 completed_bytes
293 );
294 self.cookie_helper.save_cookies_if_configured();
295 Ok(ConcurrentDownloadResult::Complete)
296 }
297
298 pub async fn execute_with_retry(
299 &mut self,
300 uri: &str,
301 total_length: u64,
302 resume_state: &ResumeState,
303 max_retries_per_segment: u32,
304 ) -> Result<ConcurrentDownloadResult> {
305 tracing::info!(
306 "Using concurrent download mode (split={}, max_retries/segment={})",
307 self.group.read().await.options().split.unwrap_or(1),
308 max_retries_per_segment
309 );
310
311 let all_uris: Vec<String> = {
312 let g = self.group.read().await;
313 g.uris().to_vec()
314 };
315
316 if all_uris.len() > 1 {
317 tracing::info!(
318 "Intelligent multi-mirror selection enabled: {} mirror sources",
319 all_uris.len()
320 );
321 self.execute_with_coordinator(
322 &all_uris,
323 total_length,
324 resume_state,
325 max_retries_per_segment,
326 )
327 .await
328 } else {
329 self.execute(uri, total_length, resume_state, max_retries_per_segment)
330 .await
331 }
332 }
333
334 async fn execute_with_coordinator(
335 &mut self,
336 uris: &[String],
337 total_length: u64,
338 resume_state: &ResumeState,
339 max_retries_per_segment: u32,
340 ) -> Result<ConcurrentDownloadResult> {
341 let split = self.group.read().await.options().split.unwrap_or(1) as u64;
342 let segment_size = total_length.div_ceil(split);
343 let max_conn = self
344 .group
345 .read()
346 .await
347 .options()
348 .max_connection_per_server
349 .unwrap_or(constants::DEFAULT_MAX_CONNECTION_PER_SERVER as u16)
350 as usize;
351
352 let mirror_config = crate::engine::mirror_coordinator::MirrorConfig {
353 max_connections_per_mirror: max_conn.min(split as usize),
354 max_total_connections: max_conn * uris.len(),
355 speed_threshold: constants::MIRROR_SPEED_THRESHOLD,
356 cooldown_secs: constants::MIRROR_COOLDOWN_SECS,
357 max_retries: max_retries_per_segment,
358 };
359
360 let selector = Box::new(
361 crate::selector::adaptive_uri_selector::AdaptiveUriSelector::new_with_uris(
362 Arc::new(crate::selector::server_stat_man::ServerStatMan::new()),
363 uris.to_vec(),
364 ),
365 );
366
367 let segment_manager = ConcurrentSegmentManager::new_with_selector(
368 total_length,
369 uris.to_vec(),
370 Some(segment_size),
371 Arc::new(crate::selector::server_stat_man::ServerStatMan::new()),
372 selector,
373 );
374
375 let mut coordinator =
376 crate::engine::mirror_coordinator::MirrorCoordinator::with_segment_manager(
377 Arc::new(crate::selector::server_stat_man::ServerStatMan::new()),
378 Box::new(crate::selector::uri_selector::InorderUriSelector::new()),
379 segment_manager,
380 mirror_config,
381 uris.to_vec(),
382 );
383
384 if resume_state.should_resume {
385 tracing::debug!(
386 "Resume: existing {} bytes, continuing from offset {}",
387 resume_state.existing_length,
388 resume_state.start_offset
389 );
390 }
391
392 let use_mmap = self.file_allocation == "mmap" && total_length >= self.mmap_threshold;
393 let mut writer =
394 CachedDiskWriter::new_with_mmap(&self.output_path, Some(total_length), None, use_mmap);
395 self.progress_updater.reset(0);
396
397 let mut consecutive_416_count = 0u32;
398 let mut total_416_count = 0u32;
399 let fallback_threshold_consecutive = 3u32;
400 let fallback_threshold_ratio = 0.2f64;
401 let mut should_fallback = false;
402
403 while coordinator.has_pending_segments() || !coordinator.is_complete() {
404 while let Some((mirror_idx, mirror_url, (seg_idx, offset, length))) =
405 coordinator.select_mirror_for_segment()
406 {
407 tracing::info!(
408 "Starting segment {} download: mirror={}, offset={}, size={}",
409 seg_idx,
410 mirror_idx,
411 offset,
412 length
413 );
414
415 let downloader = HttpSegmentDownloader::new(&self.client);
416 let seg_start = Instant::now();
417
418 let cookie_hdr = self.cookie_helper.build_cookie_header(&mirror_url);
419
420 let result = downloader
421 .download_range(
422 &mirror_url,
423 offset,
424 length,
425 cookie_hdr.as_deref(),
426 &self.headers,
427 )
428 .await;
429
430 match result {
431 Ok(data) => {
432 let elapsed = seg_start.elapsed();
433 let speed = if elapsed.as_secs_f64() > 0.0 {
434 (data.len() as f64 / elapsed.as_secs_f64()) as u64
435 } else {
436 0
437 };
438
439 tracing::debug!(
440 "Segment {} complete: {} bytes, speed={} B/s",
441 seg_idx,
442 data.len(),
443 speed
444 );
445
446 let data_for_coordinator = data.clone();
447 writer.write_bytes_at(offset, data).await.map_err(|e| {
448 Aria2Error::Fatal(crate::error::FatalError::Config(format!(
449 "Write failed: {}",
450 e
451 )))
452 })?;
453
454 coordinator.on_segment_complete(
455 mirror_idx,
456 seg_idx,
457 data_for_coordinator,
458 speed,
459 );
460 }
461 Err(e) => {
462 tracing::warn!(
463 "Segment {} download failed (mirror={}): {}",
464 seg_idx,
465 mirror_idx,
466 e
467 );
468
469 let is_416 = matches!(
470 &e,
471 Aria2Error::Recoverable(RecoverableError::RangeNotSatisfiable { .. })
472 );
473 if is_416 {
474 consecutive_416_count += 1;
475 total_416_count += 1;
476 tracing::warn!(
477 seg_idx = seg_idx,
478 consecutive_416 = consecutive_416_count,
479 total_416 = total_416_count,
480 "RangeNotSatisfiable (416) detected"
481 );
482 let failure_ratio = total_416_count as f64 / split as f64;
483 let threshold_exceeded = consecutive_416_count
484 >= fallback_threshold_consecutive
485 || failure_ratio >= fallback_threshold_ratio;
486 if threshold_exceeded {
487 tracing::warn!(
488 uri = mirror_url,
489 consecutive_416 = consecutive_416_count,
490 failure_ratio = failure_ratio,
491 "Fallback to sequential mode triggered due to RangeNotSatisfiable errors"
492 );
493 should_fallback = true;
494 break;
495 }
496 } else {
497 consecutive_416_count = 0;
498 }
499
500 let error_code = constants::HTTP_DEFAULT_ERROR_CODE;
501 coordinator.on_segment_failed(mirror_idx, seg_idx, error_code);
502 }
503 }
504
505 let completed_bytes = {
506 let total = coordinator.num_segments() as u64;
507 let progress_pct = coordinator.progress();
508 if total > 0 {
509 (progress_pct / 100.0 * total as f64) as u64
510 } else {
511 0
512 }
513 };
514
515 self.progress_updater
516 .update_progress(
517 completed_bytes,
518 constants::PROGRESS_UPDATE_BYTES as u64,
519 constants::HTTP_SPEED_UPDATE_INTERVAL_MS,
520 )
521 .await;
522 }
523
524 if coordinator.is_complete() {
525 break;
526 }
527
528 if coordinator.has_failed_segments() {
529 tracing::error!("Permanently failed download segments exist");
530 return Err(Aria2Error::Recoverable(
531 RecoverableError::TemporaryNetworkFailure {
532 message: "Some download segments permanently failed".into(),
533 },
534 ));
535 }
536
537 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
538 }
539
540 writer.flush().await.map_err(|e| {
541 Aria2Error::Fatal(crate::error::FatalError::Config(format!(
542 "Flush failed: {}",
543 e
544 )))
545 })?;
546
547 if should_fallback {
548 let completed_ranges = coordinator.completed_ranges();
549 tracing::warn!(
550 "Fallback: {} completed ranges will be preserved",
551 completed_ranges.len()
552 );
553 return Ok(ConcurrentDownloadResult::Fallback { completed_ranges });
554 }
555
556 let final_speed = {
557 let g = self.group.read().await;
558 let elapsed = g.elapsed_time().await;
559 match elapsed {
560 Some(d) if d.as_secs_f64() > 0.0 => {
561 (self.progress_updater.last_progress_update() as f64 / d.as_secs_f64()) as u64
562 }
563 _ => 0,
564 }
565 };
566
567 {
568 let mut g = self.group.write().await;
569 g.set_total_length(self.progress_updater.last_progress_update())
570 .await;
571 g.set_total_length_atomic(self.progress_updater.last_progress_update());
572 g.set_completed_length(self.progress_updater.last_progress_update());
573 g.update_speed(final_speed, 0).await;
574 g.set_download_speed_cached(final_speed);
575 g.complete().await?;
576 }
577
578 tracing::info!(
579 "Multi-mirror concurrent download complete: {} ({} bytes, {} B/s)",
580 self.output_path.display(),
581 self.progress_updater.last_progress_update(),
582 final_speed
583 );
584 self.cookie_helper.save_cookies_if_configured();
585 Ok(ConcurrentDownloadResult::Complete)
586 }
587}