fix: time out stalled upload bodies without timing merge execution

crates/lance-context-server/Cargo.toml

@@ -20,19 +20,21 @@ 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"] }
 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" }

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<Duration>,
+    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::<TimeoutError>() {
+            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::<Bytes, _>(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();
+    }
+}

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,

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,
         ))

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.