Skip to main content

aria2_core/engine/
range_prober.rs

1use 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}