diff --git a/rust_backend/crates/extensions/src/download/chunked.rs b/rust_backend/crates/extensions/src/download/chunked.rs index 09f15e6c..5bc3103e 100644 --- a/rust_backend/crates/extensions/src/download/chunked.rs +++ b/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; } diff --git a/rust_backend/crates/extensions/src/download/mod.rs b/rust_backend/crates/extensions/src/download/mod.rs index e2d05ca6..188f4ff9 100644 --- a/rust_backend/crates/extensions/src/download/mod.rs +++ b/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); } diff --git a/rust_backend/crates/extensions/src/download/segments.rs b/rust_backend/crates/extensions/src/download/segments.rs index 3993c5fe..361051c3 100644 --- a/rust_backend/crates/extensions/src/download/segments.rs +++ b/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); } diff --git a/rust_backend/crates/extensions/tests/optimization_chunked.rs b/rust_backend/crates/extensions/tests/optimization_chunked.rs new file mode 100644 index 00000000..50cc6424 --- /dev/null +++ b/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, + stop: Arc, + worker: Option>, +} + +/// 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::().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, 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"); +} diff --git a/rust_backend/crates/network/src/client.rs b/rust_backend/crates/network/src/client.rs index 071b1a2d..54a2b39d 100644 --- a/rust_backend/crates/network/src/client.rs +++ b/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, + 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::().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, }, )); diff --git a/rust_backend/crates/network/src/dns.rs b/rust_backend/crates/network/src/dns.rs index 23502c62..1bc613c8 100644 --- a/rust_backend/crates/network/src/dns.rs +++ b/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()