aria2_core/engine/
range_prober.rs1use std::sync::Arc;
2use std::time::Duration;
3
4use reqwest;
5
6use crate::constants;
7
8pub struct RangeProber {
9 client: Arc<reqwest::Client>,
10 headers: Vec<(String, String)>,
11}
12
13impl RangeProber {
14 pub fn new(client: Arc<reqwest::Client>, headers: Vec<(String, String)>) -> Self {
15 Self { client, headers }
16 }
17
18 pub async fn probe_range_support(&self, uri: &str, total_length: u64) -> bool {
19 let probe_stage_1 = self.probe_single_range(uri, "bytes=0-0").await;
20 tracing::debug!(
21 "Range probe stage 1 (bytes=0-0) for {}: {}",
22 uri,
23 probe_stage_1
24 );
25
26 if !probe_stage_1 {
27 tracing::info!(
28 "Range probe stage 1 failed for {}, falling back to sequential",
29 uri
30 );
31 return false;
32 }
33
34 let end_byte = if total_length > 1 {
35 std::cmp::min(999, total_length - 1)
36 } else {
37 0
38 };
39 let range_header = format!("bytes=0-{}", end_byte);
40 let probe_stage_2 = self.probe_single_range(uri, &range_header).await;
41 tracing::debug!(
42 "Range probe stage 2 ({}) for {}: {}",
43 range_header,
44 uri,
45 probe_stage_2
46 );
47
48 if !probe_stage_2 {
49 tracing::info!(
50 "Range probe stage 2 failed for {} ({}), falling back to sequential",
51 uri,
52 range_header
53 );
54 }
55
56 probe_stage_2
57 }
58
59 pub async fn probe_single_range(&self, uri: &str, range_header: &str) -> bool {
60 let mut req =
61 self.client
62 .get(uri)
63 .header("Range", range_header)
64 .timeout(Duration::from_secs(
65 constants::HTTP_DEFAULT_OVERALL_TIMEOUT_SECS,
66 ));
67 for (name, value) in &self.headers {
68 req = req.header(name, value);
69 }
70 match req.send().await {
71 Ok(resp) => {
72 let status = resp.status().as_u16();
73 status == 206
74 }
75 Err(e) => {
76 tracing::warn!("Range probe failed for {} ({}): {}", uri, range_header, e);
77 false
78 }
79 }
80 }
81}