perf(download): reuse connections across probes and retries

rust_backend/crates/extensions/src/download/chunked.rs

@@ -27,7 +27,11 @@ pub(super) fn download(
         check().map_err(|error| Failure::new("cancelled", error, attempts))?;
         let response = open(network, url, &options, "bytes=0-1", "", false, check);
         let failure = match response {
-            Ok(stream) if matches!(stream.response.status, 200 | 206) => break stream.response,
+            Ok(mut stream) if matches!(stream.response.status, 200 | 206) => {
+                // Finish the two-byte range so the first chunk reuses this connection.
+                stream.discard(check);
+                break stream.response;
+            }
             Ok(stream) => {
                 let mut failure = status_failure(&stream.response, &options.policy, attempts);
                 failure.error = format!("chunked probe HTTP {}", stream.response.status);
@@ -123,7 +127,7 @@ pub(super) fn download(
             if !matches!(status, 200 | 206) {
                 failure = status_failure(&stream.response, &options.policy, attempts);
                 failure.error = format!("chunked HTTP {status} at offset {start}");
-                drop(stream);
+                stream.discard(check);
                 if !retryable(status) || attempt == options.policy.max_attempts {
                     break;
                 }

rust_backend/crates/extensions/src/download/mod.rs

@@ -244,7 +244,7 @@ pub(crate) fn download(
             } else {
                 retry_delay
             };
-            drop(stream);
+            stream.discard(check);
             if !retryable(status) || attempt == options.policy.max_attempts {
                 return Err(failure);
             }

rust_backend/crates/extensions/src/download/segments.rs

@@ -299,6 +299,7 @@ fn fetch(
                 if !(200..300).contains(&status) {
                     let mut failure = status_failure(&stream.response, &options.policy, attempt);
                     failure.error = format!("segment {index} HTTP error: {status}");
+                    stream.discard(check);
                     if !retryable(status) || attempt == options.policy.max_attempts {
                         return Err(failure);
                     }

rust_backend/crates/extensions/tests/optimization_chunked.rs

@@ -0,0 +1,202 @@
+#![cfg(unix)]
+use serde_json::{Value, json};
+use spotiflac_core::cancellation::{CancellationDomain, CancellationRegistry};
+use spotiflac_extensions::RuntimeLimits;
+use spotiflac_extensions::environment::ExtensionEnvironment;
+use std::fs;
+use std::io::{Read, Write};
+use std::net::{TcpListener, TcpStream};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::thread;
+use std::time::{Duration, Instant};
+
+const BODY: &[u8] = b"0123456789abcdefghij";
+static SERIAL: Mutex<()> = Mutex::new(());
+
+struct Server {
+    base: String,
+    connections: Arc<AtomicUsize>,
+    stop: Arc<AtomicBool>,
+    worker: Option<thread::JoinHandle<()>>,
+}
+
+/// Writes headers and body separately, as proxies and CDNs commonly flush
+/// them. The body is delayed well inside the 20 ms small-body drain bound.
+fn respond(stream: &mut TcpStream, head: &str, body: &[u8]) -> std::io::Result<()> {
+    stream.write_all(head.as_bytes())?;
+    stream.flush()?;
+    thread::sleep(Duration::from_millis(3));
+    stream.write_all(body)
+}
+
+/// Serves byte ranges over keep-alive HTTP/1.1. `/busy` answers the first
+/// ranged chunk (after the two-byte probe) with a small retryable 503 body.
+fn connection(mut stream: TcpStream, stop: &AtomicBool, busy: &AtomicBool) {
+    stream
+        .set_read_timeout(Some(Duration::from_millis(25)))
+        .unwrap();
+    let deadline = Instant::now() + Duration::from_secs(10);
+    let mut request = Vec::new();
+    while Instant::now() < deadline && !stop.load(Ordering::Acquire) {
+        let mut bytes = [0; 1024];
+        match stream.read(&mut bytes) {
+            Ok(0) => break,
+            Ok(count) => request.extend_from_slice(&bytes[..count]),
+            Err(error)
+                if matches!(
+                    error.kind(),
+                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
+                ) =>
+            {
+                continue;
+            }
+            Err(_) => break,
+        }
+        assert!(request.len() <= 8192);
+        if !request.windows(4).any(|part| part == b"\r\n\r\n") {
+            continue;
+        }
+        let text = String::from_utf8(std::mem::take(&mut request)).unwrap();
+        let path = text.split_whitespace().nth(1).unwrap().to_owned();
+        let range = text
+            .lines()
+            .find_map(|line| {
+                let (name, value) = line.split_once(':')?;
+                name.eq_ignore_ascii_case("range")
+                    .then(|| value.trim().trim_start_matches("bytes=").to_owned())
+            })
+            .expect("chunked requests carry a range");
+        let (start, end) = range.split_once('-').unwrap();
+        let start: usize = start.parse().unwrap();
+        let end = end.parse::<usize>().unwrap().min(BODY.len() - 1);
+        let result = if path == "/busy"
+            && start == 0
+            && end > 1
+            && busy.swap(false, Ordering::AcqRel)
+        {
+            respond(
+                &mut stream,
+                "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 4\r\nConnection: keep-alive\r\n\r\n",
+                b"busy",
+            )
+        } else {
+            let part = &BODY[start..=end];
+            respond(
+                &mut stream,
+                &format!(
+                    "HTTP/1.1 206 Partial Content\r\nContent-Range: bytes {start}-{end}/{}\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n",
+                    BODY.len(),
+                    part.len()
+                ),
+                part,
+            )
+        };
+        if result.is_err() {
+            break;
+        }
+    }
+}
+
+impl Server {
+    fn new() -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let base = format!("http://{}", listener.local_addr().unwrap());
+        listener.set_nonblocking(true).unwrap();
+        let connections = Arc::new(AtomicUsize::new(0));
+        let count = connections.clone();
+        let stop = Arc::new(AtomicBool::new(false));
+        let stopped = stop.clone();
+        let busy = Arc::new(AtomicBool::new(true));
+        let worker = thread::spawn(move || {
+            let deadline = Instant::now() + Duration::from_secs(10);
+            let mut workers = Vec::new();
+            while Instant::now() < deadline && !stopped.load(Ordering::Acquire) {
+                match listener.accept() {
+                    Ok((stream, _)) => {
+                        count.fetch_add(1, Ordering::AcqRel);
+                        let (stopped, busy) = (stopped.clone(), busy.clone());
+                        workers.push(thread::spawn(move || connection(stream, &stopped, &busy)));
+                    }
+                    Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
+                        thread::sleep(Duration::from_millis(1));
+                    }
+                    Err(error) => panic!("chunked fixture accept: {error}"),
+                }
+            }
+            stopped.store(true, Ordering::Release);
+            for worker in workers {
+                worker.join().unwrap();
+            }
+        });
+        Self {
+            base,
+            connections,
+            stop,
+            worker: Some(worker),
+        }
+    }
+}
+
+impl Drop for Server {
+    fn drop(&mut self) {
+        self.stop.store(true, Ordering::Release);
+        if let Some(worker) = self.worker.take() {
+            let _ = worker.join();
+        }
+    }
+}
+
+fn download(route: &str) -> (Value, Vec<u8>, usize) {
+    let _serial = SERIAL.lock().unwrap_or_else(|error| error.into_inner());
+    let server = Server::new();
+    let directory = tempfile::tempdir().unwrap();
+    let environment = ExtensionEnvironment::new(
+        directory.path(),
+        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=",
+        "1",
+    )
+    .unwrap();
+    environment.set_allow_private_network(true).unwrap();
+    let manifest = json!({"name":"example.chunked","version":"1","description":"Generic chunked fixture",
+        "type":["download_provider"],"permissions":{"file":true,"network":["127.0.0.1"],"allowHttp":true}}).to_string();
+    let runtime = environment
+        .load(
+            &manifest,
+            r#"registerExtension({run(url){
+        return file.download(url, 'result.bin', {chunked: 8, maxAttempts: 2});
+    }});"#,
+            RuntimeLimits::default(),
+        )
+        .unwrap();
+    let registry = CancellationRegistry::new(CancellationDomain::Download);
+    let lease = Arc::new(registry.acquire("chunked").unwrap());
+    let arguments = json!([format!("{}/{route}", server.base)]).to_string();
+    let outcome = runtime
+        .call_download("run", &arguments, Some(lease), 10_000)
+        .unwrap();
+    let written = fs::read(directory.path().join("example.chunked/result.bin")).unwrap_or_default();
+    let connections = server.connections.load(Ordering::Acquire);
+    environment.shutdown();
+    (
+        serde_json::from_str(&outcome).unwrap(),
+        written,
+        connections,
+    )
+}
+
+#[test]
+fn chunked_probe_and_ranges_share_one_keep_alive_connection() {
+    let (value, written, connections) = download("plain");
+    assert_eq!(value["success"], true, "{value}");
+    assert_eq!(written, BODY);
+    assert_eq!(connections, 1, "probe body must not discard its connection");
+}
+
+#[test]
+fn small_retryable_error_body_keeps_connection_for_retry() {
+    let (value, written, connections) = download("busy");
+    assert_eq!(value["success"], true, "{value}");
+    assert_eq!(written, BODY);
+    assert_eq!(connections, 1, "503 body must not discard its connection");
+}

rust_backend/crates/network/src/client.rs

@@ -74,6 +74,9 @@ impl Pool {
         http.set_connect_timeout(Some(Duration::from_secs(10)));
         http.set_happy_eyeballs_timeout(Some(Duration::from_millis(300)));
         http.set_keepalive(Some(Duration::from_secs(30)));
+        // Go's dialer default: small request/control writes are not held back
+        // by Nagle while waiting for a delayed ACK.
+        http.set_nodelay(true);
         let connector = HttpsConnectorBuilder::new()
             .with_tls_config(tls)
             .https_or_http()
@@ -85,6 +88,10 @@ impl Pool {
             .pool_idle_timeout(Duration::from_secs(60))
             .pool_max_idle_per_host(10)
             .http1_max_buf_size(10 << 20)
+            // Go's 4 MiB per-stream window doubles one stream's ceiling per
+            // RTT. The 5 MiB connection window, and so peak buffering per
+            // connection, stays hyper's default.
+            .http2_initial_stream_window_size(4 << 20)
             .build(Connector(connector));
         Self {
             client,
@@ -184,7 +191,11 @@ impl NetworkService {
                 tokio::pin!(operation);
                 let timeout = tokio::time::sleep(timeout);
                 tokio::pin!(timeout);
-                let mut heartbeat = tokio::time::interval(Duration::from_millis(10));
+                // The first check runs below; an immediate first tick would
+                // only repeat it before the operation is polled.
+                let period = Duration::from_millis(10);
+                let mut heartbeat =
+                    tokio::time::interval_at(tokio::time::Instant::now() + period, period);
                 loop {
                     // Create before checking: notify_waiters also reaches a Notified
                     // future that has not been polled yet, closing the lost-wake gap.
@@ -377,6 +388,7 @@ pub struct HttpResponse {
 
 struct ResponseBody {
     reader: Box<dyn AsyncRead + Unpin + Send>,
+    small_http1: bool,
     _permit: OwnedSemaphorePermit,
 }
 
@@ -461,6 +473,35 @@ impl HttpStream {
         result
     }
 
+    /// Release an unwanted body (a range probe or retryable error). A small
+    /// HTTP/1.1 body is finished within 20 ms so its connection can return to
+    /// the pool; HTTP/2, compressed, large or slow bodies are dropped at once.
+    pub fn discard(&mut self, check: impl Fn() -> Result<(), String>) {
+        let Some(mut body) = self.body.take() else {
+            return;
+        };
+        if !body.small_http1 || self.failure.is_some() {
+            return;
+        }
+        let _ = self
+            .service
+            .run(self.generation, Duration::from_millis(20), &check, async {
+                let mut buffer = [0; 2048];
+                let mut received = 0;
+                loop {
+                    let count = body
+                        .reader
+                        .read(&mut buffer)
+                        .await
+                        .map_err(|error| error.to_string())?;
+                    received += count;
+                    if count == 0 || received > 2048 {
+                        return Ok(());
+                    }
+                }
+            });
+    }
+
     fn stall_message(&self) -> String {
         format!(
             "download stalled: no data received for {}s (network timeout)",
@@ -873,6 +914,15 @@ impl NetworkSession {
                 && response_headers
                     .get(header::CONTENT_ENCODING)
                     .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"gzip"));
+            // Same bound as redirect bodies: only an uncompressed HTTP/1.1
+            // body of at most 2 KiB may be finished by `HttpStream::discard`.
+            let small_http1 = !decompress
+                && response.version() == http::Version::HTTP_11
+                && response_headers
+                    .get(header::CONTENT_LENGTH)
+                    .and_then(|value| value.to_str().ok())
+                    .and_then(|value| value.parse::<usize>().ok())
+                    .is_some_and(|length| length <= 2048);
             let stream = response
                 .into_body()
                 .into_data_stream()
@@ -908,6 +958,7 @@ impl NetworkSession {
                 },
                 ResponseBody {
                     reader,
+                    small_http1,
                     _permit: permit,
                 },
             ));

rust_backend/crates/network/src/dns.rs

@@ -85,6 +85,7 @@ impl Dns {
         let mut http = HttpConnector::new();
         http.enforce_http(false);
         http.set_connect_timeout(Some(Duration::from_secs(5)));
+        http.set_nodelay(true);
         let connector = HttpsConnectorBuilder::new()
             .with_tls_config(tls)
             .https_only()