feat: support X-Dragonfly-Prefetch header to control prefetching (#786)

Signed-off-by: Gaius <[email protected]>

Cargo.toml

@@ -12,7 +12,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.1.112"
+version = "0.1.113"
 authors = ["The Dragonfly Developers"]
 homepage = "https://d7y.io/"
 repository = "https://github.com/dragonflyoss/client.git"
@@ -22,13 +22,13 @@ readme = "README.md"
 edition = "2021"
 
 [workspace.dependencies]
-dragonfly-client = { path = "dragonfly-client", version = "0.1.112" }
-dragonfly-client-core = { path = "dragonfly-client-core", version = "0.1.112" }
-dragonfly-client-config = { path = "dragonfly-client-config", version = "0.1.112" }
-dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.1.112" }
-dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.1.112" }
-dragonfly-client-util = { path = "dragonfly-client-util", version = "0.1.112" }
-dragonfly-client-init = { path = "dragonfly-client-init", version = "0.1.112" }
+dragonfly-client = { path = "dragonfly-client", version = "0.1.113" }
+dragonfly-client-core = { path = "dragonfly-client-core", version = "0.1.113" }
+dragonfly-client-config = { path = "dragonfly-client-config", version = "0.1.113" }
+dragonfly-client-storage = { path = "dragonfly-client-storage", version = "0.1.113" }
+dragonfly-client-backend = { path = "dragonfly-client-backend", version = "0.1.113" }
+dragonfly-client-util = { path = "dragonfly-client-util", version = "0.1.113" }
+dragonfly-client-init = { path = "dragonfly-client-init", version = "0.1.113" }
 thiserror = "1.0"
 dragonfly-api = "=2.0.167"
 reqwest = { version = "0.12.4", features = ["stream", "native-tls", "default-tls", "rustls-tls"] }

dragonfly-client/src/proxy/header.rs

@@ -45,6 +45,12 @@ pub const DRAGONFLY_FILTERED_QUERY_PARAMS_HEADER: &str = "X-Dragonfly-Filtered-Q
 /// The request will also use P2P technology to distribute the content.
 pub const DRAGONFLY_USE_P2P_HEADER: &str = "X-Dragonfly-Use-P2P";
 
+/// DRAGONFLY_PREFETCH_HEADER is the header key of prefetch in http request.
+/// X-Dragonfly-Prefetch priority is higher than prefetch in config.
+/// If the value is "true", the range request will prefetch the entire file.
+/// If the value is "false", the range request will fetch the range content.
+pub const DRAGONFLY_PREFETCH_HEADER: &str = "X-Dragonfly-Prefetch";
+
 /// get_tag gets the tag from http header.
 #[instrument(skip_all)]
 pub fn get_tag(header: &HeaderMap) -> Option<String> {
@@ -144,3 +150,18 @@ pub fn get_use_p2p(header: &HeaderMap) -> bool {
         None => false,
     }
 }
+
+/// get_prefetch gets the prefetch from http header.
+#[instrument(skip_all)]
+pub fn get_prefetch(header: &HeaderMap) -> Option<bool> {
+    match header.get(DRAGONFLY_PREFETCH_HEADER) {
+        Some(value) => match value.to_str() {
+            Ok(value) => Some(value.eq_ignore_ascii_case("true")),
+            Err(err) => {
+                error!("get use p2p from header failed: {}", err);
+                None
+            }
+        },
+        None => None,
+    }
+}

dragonfly-client/src/proxy/mod.rs

@@ -978,15 +978,31 @@ fn make_download_task_request(
             need_back_to_source: false,
             disable_back_to_source: config.proxy.disable_back_to_source,
             certificate_chain: Vec::new(),
-            // Prefetch is set to true if the request header contains the range header and the
-            // prefetch is enabled in the configuration.
-            prefetch: config.proxy.prefetch
-                && reqwest_request_header.contains_key(reqwest::header::RANGE),
+            prefetch: need_prefetch(config.clone(), &reqwest_request_header),
             object_storage: None,
         }),
     })
 }
 
+/// need_prefetch returns whether the prefetch is needed by the configuration and the request
+/// header.
+#[instrument(skip_all)]
+fn need_prefetch(config: Arc<Config>, header: &http::HeaderMap) -> bool {
+    // If the header not contains the range header, the request does not need prefetch.
+    if !header.contains_key(reqwest::header::RANGE) {
+        return false;
+    }
+
+    // If the header contains the X-Dragonfly-Prefetch header, return the value.
+    // Because the X-Dragonfly-Prefetch header has the highest priority.
+    if let Some(prefetch) = header::get_prefetch(header) {
+        return prefetch;
+    }
+
+    // Return the prefetch value from the configuration.
+    return config.proxy.prefetch;
+}
+
 /// make_download_url makes a download url by the given uri.
 #[instrument(skip_all)]
 fn make_download_url(