Reduce overhead of chunk timeouts in wasix package downloads

lib/wasix/Cargo.toml

@@ -103,6 +103,7 @@ ahash = "0.8.11"
 hyper-util = { version = "0.1.5", features = ["server", "server-graceful", "tokio", "service", "client"], optional = true }
 http-body-util = { version="0.1.1", optional = true }
 toml = "0.8"
+pin-utils = "0.1.0"
 
 [target.'cfg(not(target_arch = "riscv64"))'.dependencies.reqwest]
 workspace = true

lib/wasix/src/http/reqwest.rs

@@ -69,12 +69,59 @@ impl ReqwestHttpClient {
         let headers = std::mem::take(response.headers_mut());
 
         let status = response.status();
+
+        // Download the body.
         let data = if let Some(timeout) = self.response_body_chunk_timeout {
+            // Download the body with a chunk timeout.
+            // The timeout prevents long stalls.
+
             let mut stream = response.bytes_stream();
             let mut buf = Vec::new();
-            while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? {
-                buf.extend_from_slice(&chunk);
+
+            // Creating tokio timeouts has overhead, so instead of a fresh
+            // timeout per chunk a shared timeout is used, and a chunk counter
+            // is kept. Only if no chunk was downloaded within the timeout a
+            // timeout error is raised.
+            'OUTER: loop {
+                let timeout = tokio::time::sleep(timeout);
+                pin_utils::pin_mut!(timeout);
+
+                let mut chunk_count = 0;
+
+                loop {
+                    tokio::select! {
+                        // Biased because the timeout is secondary,
+                        // and chunks should always have priority.
+                        biased;
+
+                        res = stream.try_next() => {
+                            match res {
+                                Ok(Some(chunk)) => {
+                                    buf.extend_from_slice(&chunk);
+                                    chunk_count += 1;
+                                }
+                                Ok(None) => {
+                                    break 'OUTER;
+                                }
+                                Err(e) => {
+                                    return Err(e.into());
+                                }
+                            }
+                        }
+
+                        _ = &mut timeout => {
+                            if chunk_count == 0 {
+                                return Err(anyhow::anyhow!("Timeout while downloading response body"));
+                            } else {
+                                // Timeout, but chunks were still downloaded, so
+                                // just continue with a fresh timeout.
+                                continue 'OUTER;
+                            }
+                        }
+                    }
+                }
             }
+
             buf
         } else {
             response.bytes().await?.to_vec()