From 0bd0bfedfe9c732f8f67e23118c7a3f65c9a8ca9 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Mon, 5 Oct 2026 22:24:50 +0000 Subject: [PATCH] fix: bound stalled HTTP upload bodies by idle progress --- Cargo.lock | 1 + crates/lance-context-server/Cargo.toml | 4 +- .../lance-context-server/src/body_timeout.rs | 196 ++++++++++++++++++ crates/lance-context-server/src/config.rs | 10 + crates/lance-context-server/src/main.rs | 6 + docs/src/specs/rollout-deployment.md | 2 + 6 files changed, 218 insertions(+), 1 deletion(-) create mode 100644 crates/lance-context-server/src/body_timeout.rs diff --git a/Cargo.lock b/Cargo.lock index 393001de..e44a4c5b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3987,6 +3987,7 @@ dependencies = [ "clap", "etcd-client", "futures", + "http-body-util", "lance", "lance-context-api", "lance-context-client", diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index d101042a..2e478bc9 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -20,6 +20,7 @@ lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" } axum = { version = "0.8", features = ["json", "multipart"] } bytes = "1" futures = "0.3" +http-body-util = "0.1" metrics = "0.24" chrono = { version = "0.4", default-features = false, features = ["clock"] } clap = { version = "4", features = ["derive", "env"] } @@ -27,12 +28,13 @@ lru = "0.12" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } -tower-http = { version = "0.6", features = ["trace"] } +tower-http = { version = "0.6", features = ["trace", "timeout"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } uuid = { version = "1", features = ["v4"] } [dev-dependencies] +tokio = { version = "1", features = ["test-util", "io-util"] } lance = "9.0.0" tower = { version = "0.5", features = ["util"] } lance-context-client = { version = "0.6.7", path = "../lance-context-client" } diff --git a/crates/lance-context-server/src/body_timeout.rs b/crates/lance-context-server/src/body_timeout.rs new file mode 100644 index 00000000..1a9ac3d6 --- /dev/null +++ b/crates/lance-context-server/src/body_timeout.rs @@ -0,0 +1,196 @@ +//! Bound stalled uploads without imposing a deadline on merge or query execution. +use axum::{ + body::{Body, HttpBody}, + extract::{Request, State}, + http::StatusCode, + middleware::Next, + response::{IntoResponse, Response}, +}; +use http_body_util::BodyExt; +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Duration, +}; +use tower_http::timeout::{TimeoutBody, TimeoutError}; + +pub(crate) async fn guard_request_body( + State(idle): State, + request: Request, + next: Next, +) -> Response { + if request.body().is_end_stream() { + return next.run(request).await; + } + let expired = Arc::new(AtomicBool::new(false)); + let body_expired = expired.clone(); + let method = request.method().clone(); + let path = request.uri().path().to_owned(); + let (parts, body) = request.into_parts(); + // TimeoutBody starts its timer when the consumer requests a frame, and + // disarms it after delivery. Handler backpressure between frames therefore + // does not count as upload idleness. + let body = TimeoutBody::new(idle, body).map_err(move |error| { + if error.is::() { + body_expired.store(true, Ordering::Relaxed); + } + error + }); + let response = next.run(Request::from_parts(parts, Body::new(body))).await; + if expired.load(Ordering::Relaxed) { + metrics::counter!("http_request_body_timeouts_total").increment(1); + tracing::warn!(%method, %path, idle_seconds = idle.as_secs(), "request body idle timeout"); + // Extractors normally classify body errors as 400; an upload that made + // no progress is a request timeout, which the client may retry. + (StatusCode::REQUEST_TIMEOUT, "request body idle timeout").into_response() + } else { + response + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::{body::Bytes, middleware, routing::post, Router}; + use futures::stream; + use std::convert::Infallible; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tower::ServiceExt; + + fn app(idle: Duration) -> Router { + Router::new() + .route("/", post(|body: Bytes| async move { body })) + .layer(middleware::from_fn_with_state(idle, guard_request_body)) + } + + #[tokio::test(start_paused = true)] + async fn progressing_upload_may_exceed_total_idle_duration() { + let body = Body::from_stream(stream::unfold(0, |i| async move { + if i == 5 { + None + } else { + tokio::time::sleep(Duration::from_secs(60)).await; + Some((Ok::<_, Infallible>(Bytes::from_static(b"x")), i + 1)) + } + })); + let before = tokio::time::Instant::now(); + let response = app(Duration::from_secs(120)) + .oneshot(Request::post("/").body(body).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert!(before.elapsed() >= Duration::from_secs(300)); + assert_eq!( + axum::body::to_bytes(response.into_body(), 100) + .await + .unwrap(), + "xxxxx" + ); + } + + #[tokio::test(start_paused = true)] + async fn completed_body_does_not_time_out_a_slow_handler() { + let router = Router::new() + .route( + "/", + post(|body: Bytes| async move { + tokio::time::sleep(Duration::from_secs(600)).await; + body + }), + ) + .layer(middleware::from_fn_with_state( + Duration::from_secs(120), + guard_request_body, + )); + let response = router + .oneshot(Request::post("/").body(Body::from("complete")).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test(start_paused = true)] + async fn consumer_backpressure_is_not_upload_idleness() { + let router = Router::new() + .route( + "/", + post(|request: Request| async move { + let mut body = request.into_body(); + let first = body.frame().await.unwrap().unwrap(); + tokio::time::sleep(Duration::from_secs(600)).await; + let second = body.frame().await.unwrap().unwrap(); + assert_eq!(first.data_ref().unwrap(), "a"); + assert_eq!(second.data_ref().unwrap(), "b"); + StatusCode::OK + }), + ) + .layer(middleware::from_fn_with_state( + Duration::from_secs(120), + guard_request_body, + )); + let body = Body::from_stream(stream::iter([ + Ok::<_, Infallible>(Bytes::from_static(b"a")), + Ok(Bytes::from_static(b"b")), + ])); + let response = router + .oneshot(Request::post("/").body(body).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn other_body_errors_keep_the_extractor_status() { + let body = Body::from_stream(stream::once(async { + Err::(std::io::Error::other("upload connection lost")) + })); + let response = app(Duration::from_secs(120)) + .oneshot(Request::post("/").body(body).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn partial_upload_times_out_and_graceful_shutdown_finishes() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let admitted = Arc::new(tokio::sync::Notify::new()); + let notify = admitted.clone(); + let router = app(Duration::from_millis(200)).layer(middleware::from_fn( + move |request: Request, next: Next| { + let notify = notify.clone(); + async move { + notify.notify_one(); + next.run(request).await + } + }, + )); + let (stop, stopped) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + axum::serve(listener, router) + .with_graceful_shutdown(async { stopped.await.unwrap() }) + .await + .unwrap(); + }); + let mut client = tokio::net::TcpStream::connect(address).await.unwrap(); + client + .write_all(b"POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\n\r\nx") + .await + .unwrap(); + admitted.notified().await; + stop.send(()).unwrap(); + let mut response = String::new(); + tokio::time::timeout(Duration::from_secs(5), client.read_to_string(&mut response)) + .await + .unwrap() + .unwrap(); + assert!(response.starts_with("HTTP/1.1 408"), "{response}"); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .unwrap() + .unwrap(); + } +} diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 0fddc20c..434f392d 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -7,6 +7,16 @@ pub struct ServerConfig { #[command(flatten)] pub merge_rollout: lance_context_merge::rollout::MergeRollout, + /// Maximum wait for the next requested upload body frame. Each frame resets + /// this idle timer; completed uploads and handler/merge execution are not timed. + #[arg( + long, + env = "REQUEST_BODY_IDLE_TIMEOUT_SECS", + default_value_t = 120, + value_parser = clap::value_parser!(u64).range(1..) + )] + pub request_body_idle_timeout_secs: u64, + /// Legacy wire field; owned merges now use the real-progress idle timeout. #[arg(long, env = "MERGE_EXECUTION_TIMEOUT_SECS", default_value_t = 3600)] pub merge_execution_timeout_secs: u64, diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index ce0be211..79d2950d 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -1,3 +1,4 @@ +mod body_timeout; mod config; mod error; mod merge_execution; @@ -24,6 +25,7 @@ async fn main() { .init(); let config = ServerConfig::parse(); + let body_idle_timeout = std::time::Duration::from_secs(config.request_body_idle_timeout_secs); let addr = format!("{}:{}", config.host, config.port); if let Err(e) = create_local_dir_if_needed(&config.data_dir) { @@ -73,6 +75,10 @@ async fn main() { let app = routes::router() .with_state(state.clone()) .merge(lance_context_metrics::metrics_router(metrics_handle)) + .layer(axum::middleware::from_fn_with_state( + body_idle_timeout, + body_timeout::guard_request_body, + )) .layer(axum::middleware::from_fn( lance_context_metrics::http_metrics_layer, )) diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index 7d7a70da..db9cefce 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -179,6 +179,8 @@ The caps form an **OR**: whichever binds first ends the pass. Bytes are measured The manifest reclaims whole generations, so a generation is the smallest indivisible merge unit. A pass always finishes the generation that reaches the budget, even if its first generation alone is oversized. Thus the read buffer can exceed the budget by up to one generation; this is not a hard process-RSS limit, and scan, deduplication and commit allocations need additional headroom. +**Upload progress.** `REQUEST_BODY_IDLE_TIMEOUT_SECS` (default `120`, minimum `1`) bounds how long an HTTP handler waits for its next request-body frame. An upload that stops providing frames receives HTTP 408 and increments `http_request_body_timeouts_total`. Each received frame resets the timer. Time spent processing a frame, executing a merge/query after reading the body, or sending a response is outside this timer. This prevents an incomplete upload from holding graceful worker shutdown indefinitely while it is waiting for body data. + **Blob request admission.** `ROLLOUT_MAX_INFLIGHT_BLOB_BYTES` (default `0`, disabled) separately controls concurrent HTTP blob requests. Known-length uploads reserve their declared body size. Unknown-length uploads and blob downloads reserve the whole budget before buffering or loading; they are admitted only when it is idle. Download sizes cannot be trusted from optional, caller-supplied `payload_size` metadata. After loading, a smaller download releases unused capacity for known-length uploads and holds the remaining reservation through the response lifetime. A lone oversized request is still admitted. This conservative policy can reduce concurrency when the budget is enabled. Only merged generation ids and directories are removed. Leftovers stay pending and drain on subsequent passes. The caps apply independently of the count-trigger threshold (`--rollout-merge-after-generations`), including when time-triggered cleanup merges at a threshold of one generation. Server-managed rollout, datagen and generic stores share these settings.