diff --git a/apps/rocm/src/main.rs b/apps/rocm/src/main.rs index 75a75ffc3..53f65288d 100644 --- a/apps/rocm/src/main.rs +++ b/apps/rocm/src/main.rs @@ -40,7 +40,7 @@ use rocm_core::{ model_catalog_platforms, model_recipe_featured, model_recipe_target_platform_label, normalize_therock_family, platform_matches_gfx_family, preferred_serve_engine_for_host_gpu_summary, prepend_runtime_path, process_is_running, - read_tcp_stream_to_string, resolve_builtin_model_recipe, resolve_model_recipe, + read_http_response_bounded, resolve_builtin_model_recipe, resolve_model_recipe, runtime_install_root_is_protected, runtime_path_is_same_or_inside, runtime_python_activation_hint, runtime_python_env_bin_dir, runtime_python_executable_in_env, shell_command_for_host, uv_cache_source, write_all_tcp_stream, @@ -18875,6 +18875,7 @@ fn http_get_local_service( endpoint_api_key: Option<&str>, timeout: Duration, ) -> Result<(u16, String)> { + let deadline = std::time::Instant::now() + timeout; let mut stream = connect_tcp_stream(host, port, timeout)?; let host_header = format_host_port(host, port); // Authenticate the probe when the endpoint is protected; loopback endpoints @@ -18888,7 +18889,7 @@ fn http_get_local_service( ); write_all_tcp_stream(&mut stream, request.as_bytes()) .context("failed to write service readiness request")?; - let response = read_tcp_stream_to_string(&mut stream) + let response = read_http_response_bounded(&mut stream, deadline) .context("failed to read service readiness response")?; let (headers, body) = response .split_once("\r\n\r\n") @@ -18909,6 +18910,7 @@ fn http_post_local_service_json( body: &serde_json::Value, timeout: Duration, ) -> Result<(u16, String)> { + let deadline = std::time::Instant::now() + timeout; let mut stream = connect_tcp_stream(host, port, timeout)?; let host_header = format_host_port(host, port); let body = serde_json::to_string(body).context("failed to serialize service request")?; @@ -18919,8 +18921,8 @@ fn http_post_local_service_json( ); write_all_tcp_stream(&mut stream, request.as_bytes()) .context("failed to write service request")?; - let response = - read_tcp_stream_to_string(&mut stream).context("failed to read service response")?; + let response = read_http_response_bounded(&mut stream, deadline) + .context("failed to read service response")?; let (headers, body) = response .split_once("\r\n\r\n") .unwrap_or((response.as_str(), "")); @@ -20228,6 +20230,70 @@ mod tests { Ok(()) } + #[test] + fn lemonade_stop_unload_is_bounded_by_the_request_timeout() -> Result<()> { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::time::Instant; + + // Regression test for the stall this PR fixes: a peer that trickles the + // response one byte at a time, never framing or closing, used to stall + // `read_tcp_stream_to_string`'s read-to-EOF loop indefinitely. That hung + // `unload_lemonade_service_model` past its 5s timeout during scenario + // teardown, showing up as an unexplained multi-minute gap. The unload + // call must now return an error at (not far past) its 5s budget. + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let port = listener.local_addr()?.port(); + thread::spawn(move || { + let Ok((mut stream, _)) = listener.accept() else { + return; + }; + stream.set_read_timeout(Some(Duration::from_secs(2))).ok(); + let mut buffer = [0_u8; 512]; + let _ = stream.read(&mut buffer); + let body = b"{\"status\":\"success\",\"message\":\"ok\"}"; + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + let _ = stream.write_all(header.as_bytes()); + // One byte every 300ms never finishes framing the 35-byte body + // inside the 5s unload timeout below, so the bound under test is + // the deadline firing, not the response completing early. + for byte in body { + if stream.write_all(&[*byte]).is_err() { + return; + } + thread::sleep(Duration::from_millis(300)); + } + }); + + let (_root, paths) = test_paths("lemonade-stop-unload-dribble"); + let record = ManagedServiceRecord::new( + &paths, + "svc-qwen", + "lemonade", + "qwen", + "Qwen3-0.6B-GGUF", + "127.0.0.1", + port, + "managed", + 123, + Some("therock-release".to_owned()), + Some("lemonade-embeddable-10.6.0".to_owned()), + Some("gpu_required".to_owned()), + ); + let started = Instant::now(); + assert!(unload_lemonade_service_model(&record).is_err()); + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(4), + "bounded BY the 5s deadline, not failing early: {elapsed:?}" + ); + assert!(elapsed < Duration::from_secs(8), "{elapsed:?}"); + Ok(()) + } + #[test] fn serve_readiness_wait_withholds_ready_while_the_model_only_lists() -> Result<()> { use std::io::{Read, Write}; @@ -20332,6 +20398,71 @@ mod tests { Ok(()) } + #[test] + fn serve_readiness_ready_verdict_does_not_wait_for_the_peer_to_close() -> Result<()> { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::time::Instant; + + // Regression test for the other half of this PR's fix: a response is + // read to completion by its own framing, not by waiting for the peer + // to close. Before this fix, `read_tcp_stream_to_string` blocked + // until EOF, so a keep-alive engine that answers correctly but never + // closes the socket looked identical to a hung one — the readiness + // probe ran out its timeout and reported not-ready even though the + // answer had already arrived. + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let port = listener.local_addr()?.port(); + let server = thread::spawn(move || { + while let Ok((mut stream, _)) = listener.accept() { + thread::spawn(move || { + stream.set_read_timeout(Some(Duration::from_secs(2))).ok(); + let mut buffer = [0_u8; 1024]; + let Ok(read) = stream.read(&mut buffer) else { + return; + }; + let request = String::from_utf8_lossy(&buffer[..read]).into_owned(); + let body = if request.starts_with("POST /v1/chat/completions ") { + r#"{"choices":[{"message":{"content":"ok"}}]}"# + } else { + r#"{"data":[{"id":"Qwen3-0.6B-GGUF"}]}"# + }; + let _ = write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ); + // Hold the connection open well past the readiness wait's + // timeout below, and deliberately omit `Connection: + // close`. The client must not need EOF to recognize the + // response as complete. + thread::sleep(Duration::from_secs(10)); + }); + } + }); + + let started = Instant::now(); + let readiness = wait_for_service_http_ready( + "vllm", + "127.0.0.1", + port, + "Qwen3-0.6B-GGUF", + None, + Duration::from_secs(5), + ); + + assert_eq!(readiness, EndpointReadiness::Serving); + assert!( + started.elapsed() < Duration::from_secs(3), + "a complete response must be recognized without waiting on the peer to close" + ); + // `server`'s accept loop runs forever; dropping the JoinHandle detaches + // it rather than joining, and the thread dies with the test process. + drop(server); + Ok(()) + } + #[test] fn a_loading_service_keeps_its_status_instead_of_being_demoted() -> Result<()> { use std::io::{Read, Write}; diff --git a/apps/rocm/src/providers.rs b/apps/rocm/src/providers.rs index af6e8609b..d3cde077e 100644 --- a/apps/rocm/src/providers.rs +++ b/apps/rocm/src/providers.rs @@ -6,7 +6,7 @@ use anyhow::{Context, Result, bail}; use rocm_core::{ AppPaths, AuditEventRecord, EndpointReadiness, EndpointReadinessOutcome, ManagedServiceRecord, RocmCliConfig, append_audit_event, connect_tcp_stream, format_host_port, - managed_service_endpoint_readiness, read_tcp_stream_to_string, unix_time_millis, + managed_service_endpoint_readiness, read_http_response_bounded, unix_time_millis, write_all_tcp_stream, }; use std::fs; @@ -596,6 +596,7 @@ fn post_json_to_local_endpoint_body( let (host, port) = parse_http_endpoint(endpoint_url) .with_context(|| format!("unsupported local endpoint URL `{endpoint_url}`"))?; let body = serde_json::to_string(body).context("failed to serialize chat request")?; + let deadline = std::time::Instant::now() + timeout; let mut stream = connect_tcp_stream(&host, port, timeout)?; let host_header = format_host_port(&host, port); let auth_header = local_bearer_header(endpoint_api_key); @@ -607,7 +608,7 @@ fn post_json_to_local_endpoint_body( write_all_tcp_stream(&mut stream, request.as_bytes()) .context("failed to write local provider chat request")?; - let response = read_tcp_stream_to_string(&mut stream) + let response = read_http_response_bounded(&mut stream, deadline) .context("failed to read local provider chat response")?; let (headers, body) = response .split_once("\r\n\r\n") diff --git a/crates/rocm-core/src/lib.rs b/crates/rocm-core/src/lib.rs index ec50c0724..37b69af4a 100644 --- a/crates/rocm-core/src/lib.rs +++ b/crates/rocm-core/src/lib.rs @@ -1009,14 +1009,6 @@ pub fn write_all_tcp_stream(stream: &mut TcpStream, bytes: &[u8]) -> Result<()> .context("failed to write to TCP stream") } -pub fn read_tcp_stream_to_string(stream: &mut TcpStream) -> Result { - let mut response = String::new(); - stream - .read_to_string(&mut response) - .context("failed to read TCP stream")?; - Ok(response) -} - /// Read one HTTP response, bounded by a wall-clock deadline. /// /// Two problems with reading to end-of-stream instead. A response is only @@ -1028,7 +1020,7 @@ pub fn read_tcp_stream_to_string(stream: &mut TcpStream) -> Result { /// slow-drip responder could stretch the total wait to an arbitrary multiple of /// what the caller asked for. This returns as soon as the response is complete by /// its own framing, and never runs past `deadline` in total. -fn read_http_response_bounded(stream: &mut TcpStream, deadline: Instant) -> Result { +pub fn read_http_response_bounded(stream: &mut TcpStream, deadline: Instant) -> Result { let mut response = Vec::new(); let mut chunk = [0_u8; 4096]; while !http_response_is_complete(&response) { @@ -1069,10 +1061,17 @@ fn read_http_response_bounded(stream: &mut TcpStream, deadline: Instant) -> Resu /// those are delimited by the connection closing, so the caller must keep reading /// until EOF. fn http_response_is_complete(response: &[u8]) -> bool { - let text = String::from_utf8_lossy(response); - let Some((headers, body)) = text.split_once("\r\n\r\n") else { + // Headers are ASCII by the HTTP spec, so it is safe to lossy-decode just + // that slice to parse them. The body length check below stays on raw + // bytes: lossy-decoding a body that ends mid multi-byte UTF-8 sequence + // replaces the truncated tail with a 3-byte U+FFFD, which can inflate a + // partial body's *decoded* length past the declared Content-Length and + // report completeness one read early. + let Some(header_end) = response.windows(4).position(|w| w == b"\r\n\r\n") else { return false; }; + let headers = String::from_utf8_lossy(&response[..header_end]); + let body = &response[header_end + 4..]; let header_value = |name: &str| { headers.lines().find_map(|line| { let (key, value) = line.split_once(':')?; @@ -1089,7 +1088,7 @@ fn http_response_is_complete(response: &[u8]) -> bool { if header_value("Transfer-Encoding") .is_some_and(|value| value.to_ascii_lowercase().contains("chunked")) { - return body.ends_with("0\r\n\r\n"); + return body.ends_with(b"0\r\n\r\n"); } false } @@ -8454,6 +8453,25 @@ mod tests { Ok(()) } + #[test] + fn http_response_is_complete_does_not_miscount_a_split_multibyte_char() { + // A body ending in a multi-byte UTF-8 character can arrive one byte + // short of the declared Content-Length. Lossy-decoding the whole + // buffer to check completeness turns that dangling partial sequence + // into a 3-byte U+FFFD replacement, inflating the decoded length past + // the declared one and reporting completeness a read early. + let body = "hi \u{2603}"; // snowman is a 3-byte UTF-8 character + let full = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ) + .into_bytes(); + let truncated = &full[..full.len() - 1]; + + assert!(!http_response_is_complete(truncated)); + assert!(http_response_is_complete(&full)); + } + /// A signal arriving mid-response must not fail the request. /// /// A signal delivered while the client is parked in `read` aborts it with