From 11e80f4bce6524fe740490eed1879ffa113ece1e Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Tue, 6 Oct 2026 02:08:28 +0000 Subject: [PATCH] fix: join HTTP handlers before retiring stalled shutdown responses --- crates/lance-context-server/src/config.rs | 11 + crates/lance-context-server/src/drain.rs | 295 +++++++++++++++ .../lance-context-server/src/drain/tests.rs | 348 ++++++++++++++++++ crates/lance-context-server/src/main.rs | 12 +- docs/src/specs/rollout-deployment.md | 3 + 5 files changed, 663 insertions(+), 6 deletions(-) create mode 100644 crates/lance-context-server/src/drain.rs create mode 100644 crates/lance-context-server/src/drain/tests.rs diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 434f392..91622e8 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -17,6 +17,17 @@ pub struct ServerConfig { )] pub request_body_idle_timeout_secs: u64, + /// During graceful shutdown, retire response connections whose socket writes + /// are blocked for this long, after all admitted HTTP handlers have finished. + /// This does not impose a total deadline on queries, merges, or downloads. + #[arg( + long, + env = "HTTP_SHUTDOWN_RESPONSE_IDLE_SECS", + default_value_t = 120, + value_parser = clap::value_parser!(u64).range(1..) + )] + pub http_shutdown_response_idle_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/drain.rs b/crates/lance-context-server/src/drain.rs new file mode 100644 index 0000000..b9981e9 --- /dev/null +++ b/crates/lance-context-server/src/drain.rs @@ -0,0 +1,295 @@ +//! Join accepted handlers before retiring stalled response connections. A slow +//! reader must not pin shutdown, and a disconnected client must not cancel a +//! handler halfway through publishing a storage mutation. Response bodies remain +//! cancellable: blob bodies own their data, and staged-append streams only write +//! unreachable immutable files (their separate coordinator publishes them). +use tracing::Instrument; + +use axum::{ + extract::{Request, State}, + http::StatusCode, + middleware::Next, + response::{IntoResponse, Response}, + serve::Listener, + Router, +}; +use std::{ + future::Future, + io, + net::{Shutdown, SocketAddr}, + pin::Pin, + sync::{Arc, Mutex, Weak}, + task::{Context, Poll}, + time::Duration, +}; +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::{TcpListener, TcpStream}, + sync::Notify, + time::Instant, +}; + +/// Serve until shutdown, joining storage handlers before retiring blocked +/// response sockets. Returning also joins handlers detached by client disconnect. +pub(crate) async fn serve( + listener: TcpListener, + app: Router, + response_idle: Duration, + shutdown: impl Future + Send + 'static, +) -> io::Result<()> { + let drain = Arc::new(Drain::default()); + let app = app.layer(axum::middleware::from_fn_with_state( + drain.clone(), + track_handler, + )); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let response_drain = { + let drain = drain.clone(); + tokio::spawn(async move { + if shutdown_rx.await.is_ok() { + drain.retire_idle_responses(response_idle).await; + } + }) + }; + let admission = drain.clone(); + let result = axum::serve(DrainListener::new(listener, drain.clone()), app) + .with_graceful_shutdown(async move { + shutdown.await; + admission.close_admission(); + let _ = shutdown_tx.send(()); + }) + .await; + response_drain.abort(); + let _ = response_drain.await; + drain.close_admission(); + drain.joined().await; + result +} + +#[derive(Default)] +struct Admission { + closed: bool, + active: usize, +} +#[derive(Default)] +pub(crate) struct Drain { + admission: Mutex, + changed: Notify, + connections: Mutex>>, +} +struct Handler(Arc); +impl Drop for Handler { + fn drop(&mut self) { + let mut admission = self.0.admission.lock().unwrap(); + admission.active -= 1; + self.0.changed.notify_waiters(); + } +} +impl Drain { + fn admit(self: &Arc) -> Option { + let mut admission = self.admission.lock().unwrap(); + if admission.closed { + return None; + } + admission.active += 1; + Some(Handler(self.clone())) + } + pub(crate) fn close_admission(&self) { + let mut admission = self.admission.lock().unwrap(); + if !admission.closed { + tracing::info!( + active_http_handlers = admission.active, + "closed HTTP admission; joining accepted work" + ); + admission.closed = true; + self.changed.notify_waiters(); + } + } + pub(crate) async fn joined(&self) { + loop { + let changed = self.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.admission.lock().unwrap().active == 0 { + return; + } + changed.await; + } + } + /// Invoke only after admission closes. No handler time limit: all accepted + /// writes finish before any response connection can be forcibly closed. + pub(crate) async fn retire_idle_responses(&self, idle: Duration) { + self.joined().await; + tracing::info!("accepted HTTP handlers joined; draining response connections"); + loop { + { + let mut connections = self.connections.lock().unwrap(); + connections.retain(|weak| { + let Some(connection) = weak.upgrade() else { + return false; + }; + // Serialize the idle check/close with actual socket writes, + // so a concurrent successful write cannot race retirement. + let blocked = connection.blocked_since.lock().unwrap(); + if blocked.is_some_and(|since| since.elapsed() >= idle) { + // The duplicate descriptor refers to this exact socket; + // no FD-number reuse or Pod/network administration needed. + match connection.control.shutdown(Shutdown::Both) { + Ok(()) => { + metrics::counter!("http_shutdown_idle_responses_total") + .increment(1); + tracing::warn!( + idle_seconds = idle.as_secs(), + "closed stalled HTTP response after handlers joined" + ); + return false; + } + Err(error) => { + tracing::warn!(%error, "failed to close stalled response") + } + } + } + true + }); + } + tokio::time::sleep(idle.min(Duration::from_secs(1))).await; + } + } +} + +pub(crate) async fn track_handler( + State(drain): State>, + request: Request, + next: Next, +) -> Response { + let Some(handler) = drain.admit() else { + return (StatusCode::SERVICE_UNAVAILABLE, "worker is draining").into_response(); + }; + // Dropping a JoinHandle detaches it. A client disconnect therefore cannot + // drop an accepted mutation halfway through a storage commit. The guard + // remains in the task until success, error or panic. + match tokio::spawn( + async move { + let _handler = handler; + next.run(request).await + } + .in_current_span(), + ) + .await + { + Ok(response) => response, + Err(error) => { + tracing::error!(%error, "HTTP handler task failed"); + StatusCode::INTERNAL_SERVER_ERROR.into_response() + } + } +} + +struct Connection { + control: std::net::TcpStream, + blocked_since: Mutex>, +} +impl Connection { + fn wrote(blocked: &mut Option, result: &Poll>) { + match result { + Poll::Pending => { + blocked.get_or_insert_with(Instant::now); + } + Poll::Ready(Ok(n)) if *n > 0 => *blocked = None, + _ => {} + } + } +} +pub(crate) struct DrainListener { + listener: TcpListener, + drain: Arc, +} +impl DrainListener { + pub(crate) fn new(listener: TcpListener, drain: Arc) -> Self { + Self { listener, drain } + } +} +pub(crate) struct ConnectionIo { + stream: TcpStream, + connection: Arc, +} +impl Listener for DrainListener { + type Io = ConnectionIo; + type Addr = SocketAddr; + async fn accept(&mut self) -> (Self::Io, Self::Addr) { + loop { + let (stream, address) = Listener::accept(&mut self.listener).await; + let wrapped = (|| { + let stream = stream.into_std()?; + let connection = Arc::new(Connection { + control: stream.try_clone()?, + blocked_since: Mutex::new(None), + }); + let stream = TcpStream::from_std(stream)?; + let mut connections = self.drain.connections.lock().unwrap(); + // Amortize pruning instead of scanning every live connection + // for each accepted request. Dead weak entries retain no socket. + if connections.len().is_multiple_of(256) { + connections.retain(|c| c.strong_count() != 0); + } + connections.push(Arc::downgrade(&connection)); + Ok::<_, io::Error>(ConnectionIo { stream, connection }) + })(); + match wrapped { + Ok(io) => return (io, address), + Err(error) => { + tracing::warn!(%error, "failed to track accepted socket"); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } + } + } + fn local_addr(&self) -> io::Result { + self.listener.local_addr() + } +} +impl AsyncRead for ConnectionIo { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(cx, buf) + } +} +impl AsyncWrite for ConnectionIo { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = &mut *self; + let mut blocked = this.connection.blocked_since.lock().unwrap(); + let result = Pin::new(&mut this.stream).poll_write(cx, buf); + Connection::wrote(&mut blocked, &result); + result + } + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + let this = &mut *self; + let mut blocked = this.connection.blocked_since.lock().unwrap(); + let result = Pin::new(&mut this.stream).poll_write_vectored(cx, bufs); + Connection::wrote(&mut blocked, &result); + result + } + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_flush(cx) + } + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(cx) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/lance-context-server/src/drain/tests.rs b/crates/lance-context-server/src/drain/tests.rs new file mode 100644 index 0000000..ae75356 --- /dev/null +++ b/crates/lance-context-server/src/drain/tests.rs @@ -0,0 +1,348 @@ +use super::*; +use axum::{ + body::{Body, Bytes}, + middleware, + routing::{get, post}, +}; +use std::{ + convert::Infallible, + sync::atomic::{AtomicBool, Ordering}, +}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::oneshot, + task::JoinHandle, +}; +use tower::ServiceExt; + +async fn start( + app: Router, + idle: Duration, +) -> (SocketAddr, oneshot::Sender<()>, JoinHandle>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, rx) = oneshot::channel(); + let server = tokio::spawn(serve(listener, app, idle, async { + let _ = rx.await; + })); + (addr, tx, server) +} + +async fn request(addr: SocketAddr, path: &str) -> TcpStream { + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + // Keep the receiver's window small enough to exercise real backpressure + // without allocating a large response in either process. + socket.set_recv_buffer_size(4096).unwrap(); + let mut stream = socket.connect(addr).await.unwrap(); + stream + .write_all( + format!("GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n") + .as_bytes(), + ) + .await + .unwrap(); + stream +} + +#[tokio::test] +async fn disconnect_does_not_cancel_accepted_mutation_or_skip_shutdown_join() { + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let committed = Arc::new(AtomicBool::new(false)); + let app = Router::new().route( + "/", + post({ + let (entered, release, committed) = + (entered.clone(), release.clone(), committed.clone()); + move || { + let (entered, release, committed) = + (entered.clone(), release.clone(), committed.clone()); + async move { + entered.notify_one(); + release.notified().await; + committed.store(true, Ordering::SeqCst); + StatusCode::OK + } + } + }), + ); + let (addr, shutdown, mut server) = start(app, Duration::from_millis(30)).await; + let mut client = TcpStream::connect(addr).await.unwrap(); + client + .write_all(b"POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 0\r\n\r\n") + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + drop(client); + shutdown.send(()).unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(150), &mut server) + .await + .is_err() + ); + assert!(!committed.load(Ordering::SeqCst)); + release.notify_one(); + tokio::time::timeout(Duration::from_secs(2), server) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(committed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn closed_admission_rejects_without_running_handler() { + let drain = Arc::new(Drain::default()); + let app = Router::new() + .route( + "/", + get(|| async { + panic!("must not run"); + #[allow(unreachable_code)] + StatusCode::OK + }), + ) + .layer(middleware::from_fn_with_state(drain.clone(), track_handler)); + drain.close_admission(); + let response = app + .oneshot(Request::get("/").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + drain.joined().await; +} + +#[tokio::test] +async fn panicking_and_failed_handlers_release_all_join_waiters() { + let drain = Arc::new(Drain::default()); + let app = Router::new() + .route( + "/panic", + get(|| async { + panic!("injected handler panic"); + #[allow(unreachable_code)] + StatusCode::OK + }), + ) + .route("/error", get(|| async { StatusCode::BAD_REQUEST })) + .layer(middleware::from_fn_with_state(drain.clone(), track_handler)); + for (path, status) in [ + ("/panic", StatusCode::INTERNAL_SERVER_ERROR), + ("/error", StatusCode::BAD_REQUEST), + ] { + let response = app + .clone() + .oneshot(Request::get(path).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), status); + } + // Both observers must wake, including a waiter registered just before drop. + let guard = drain.admit().unwrap(); + drain.close_admission(); + let a = { + let d = drain.clone(); + tokio::spawn(async move { d.joined().await }) + }; + let b = { + let d = drain.clone(); + tokio::spawn(async move { d.joined().await }) + }; + tokio::task::yield_now().await; + drop(guard); + tokio::time::timeout(Duration::from_secs(2), async { + a.await.unwrap(); + b.await.unwrap(); + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn stalled_response_is_retired_only_after_all_handlers_join() { + let release = Arc::new(Notify::new()); + let entered = Arc::new(Notify::new()); + let produced = Arc::new(Notify::new()); + struct Dropped(Arc); + impl Drop for Dropped { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let dropped = Arc::new(AtomicBool::new(false)); + let app = Router::new() + .route( + "/download", + get({ + let produced = produced.clone(); + let dropped = dropped.clone(); + move || { + let produced = produced.clone(); + let held = Dropped(dropped.clone()); + async move { + Body::from_stream(futures::stream::unfold((0, held), move |(i, held)| { + let produced = produced.clone(); + async move { + if i == 1024 { + None + } else { + produced.notify_one(); + Some(( + Ok::<_, Infallible>(Bytes::from(vec![b'x'; 64 * 1024])), + (i + 1, held), + )) + } + } + })) + } + } + }), + ) + .route( + "/write", + get({ + let (release, entered) = (release.clone(), entered.clone()); + move || { + let (release, entered) = (release.clone(), entered.clone()); + async move { + entered.notify_one(); + release.notified().await; + StatusCode::OK + } + } + }), + ); + let (addr, shutdown, mut server) = start(app, Duration::from_millis(50)).await; + let download = request(addr, "/download").await; + tokio::time::timeout(Duration::from_secs(2), produced.notified()) + .await + .unwrap(); + let _write = request(addr, "/write").await; + tokio::time::timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + shutdown.send(()).unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(250), &mut server) + .await + .is_err() + ); + // The response must not have been cancelled while another admitted storage + // future is still running, even after several response-idle intervals. + assert!(!dropped.load(Ordering::SeqCst)); + release.notify_one(); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .unwrap() + .unwrap() + .unwrap(); + // Server completed while the non-reading client still owns its connection. + assert!(dropped.load(Ordering::SeqCst)); + drop(download); +} + +#[tokio::test] +async fn progressing_download_survives_beyond_total_idle_duration() { + let app = Router::new().route( + "/", + get(|| async { + Body::from_stream(futures::stream::unfold(0, |i| async move { + if i == 16 { + None + } else { + tokio::time::sleep(Duration::from_millis(20)).await; + Some((Ok::<_, Infallible>(Bytes::from_static(b"progress")), i + 1)) + } + })) + }), + ); + let idle = Duration::from_millis(100); + let (addr, shutdown, server) = start(app, idle).await; + let mut client = request(addr, "/").await; + let mut first = [0; 1]; + client.read_exact(&mut first).await.unwrap(); + let since = Instant::now(); + shutdown.send(()).unwrap(); + let mut rest = Vec::new(); + tokio::time::timeout(Duration::from_secs(5), client.read_to_end(&mut rest)) + .await + .unwrap() + .unwrap(); + assert!(since.elapsed() > idle * 2); + assert_eq!( + rest.windows(b"progress".len()) + .filter(|w| *w == b"progress") + .count(), + 16 + ); + server.await.unwrap().unwrap(); +} + +#[test] +fn response_idle_timeout_must_be_positive() { + use clap::Parser; + assert!(crate::config::ServerConfig::try_parse_from([ + "server", + "--http-shutdown-response-idle-secs", + "0" + ]) + .is_err()); +} + +#[tokio::test] +async fn resumed_socket_writes_reset_idle_even_when_download_exceeds_it() { + let idle = Duration::from_secs(1); + let drain = Arc::new(Drain::default()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut client = TcpStream::connect(listener.local_addr().unwrap()) + .await + .unwrap(); + let mut listener = DrainListener::new(listener, drain.clone()); + let (mut output, _) = listener.accept().await; + let connection = output.connection.clone(); + let total = 32 * 1024 * 1024; + let sender = tokio::spawn(async move { + let chunk = [b'x'; 64 * 1024]; + for _ in 0..total / chunk.len() { + output.write_all(&chunk).await?; + } + output.shutdown().await + }); + // Unlike the incremental-body test, force a real socket Pending first. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if connection.blocked_since.lock().unwrap().is_some() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + drain.close_admission(); + let retirement = tokio::spawn(async move { drain.retire_idle_responses(idle).await }); + let started = Instant::now(); + let received = tokio::time::timeout(Duration::from_secs(30), async { + let mut buffer = [0; 64 * 1024]; + let mut received = 0; + loop { + let n = client.read(&mut buffer).await.unwrap(); + if n == 0 { + break; + } + assert!(buffer[..n].iter().all(|b| *b == b'x')); + received += n; + tokio::time::sleep(Duration::from_millis(8)).await; + } + received + }) + .await + .unwrap(); + retirement.abort(); + let _ = retirement.await; + sender.await.unwrap().unwrap(); + assert_eq!(received, total); + assert!(started.elapsed() > idle * 2); +} diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 79d2950..b00c4f7 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -1,5 +1,6 @@ mod body_timeout; mod config; +mod drain; mod error; mod merge_execution; mod routes; @@ -26,6 +27,7 @@ async fn main() { let config = ServerConfig::parse(); let body_idle_timeout = std::time::Duration::from_secs(config.request_body_idle_timeout_secs); + let response_idle = std::time::Duration::from_secs(config.http_shutdown_response_idle_secs); let addr = format!("{}:{}", config.host, config.port); if let Err(e) = create_local_dir_if_needed(&config.data_dir) { @@ -87,15 +89,13 @@ async fn main() { tracing::info!("Starting lance-context-server on {}", addr); let listener = TcpListener::bind(&addr).await.unwrap(); - axum::serve(listener, app) - .with_graceful_shutdown(wait_for_shutdown_signal()) - .await - .unwrap(); + let result = drain::serve(listener, app, response_idle, wait_for_shutdown_signal()).await; - // Connections have drained. Deterministically close every resident writer - // before the runtime tears down. + // Accepted HTTP handlers, including those whose clients disconnected, have + // joined before we close resident writers. state.merge_executions.shutdown().await; state.shutdown().await; + result.expect("HTTP server failed"); } #[cfg(unix)] diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index db9cefc..f1fcdfb 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -181,6 +181,9 @@ The manifest reclaims whole generations, so a generation is the smallest indivis **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. +**Shutdown response progress.** On shutdown the worker closes HTTP admission and joins accepted handler futures, including handlers whose clients disconnected. It then retires sockets that have remained blocked while writing responses for `HTTP_SHUTDOWN_RESPONSE_IDLE_SECS` (default `120`, minimum `1`). Successful socket writes reset this idle interval; it is not a total query, merge, or download deadline. Retirement increments `http_shutdown_idle_responses_total` and logs a warning. Resident writers close after HTTP draining and handler joining. Response bodies remain cancellable: blob responses hold already-read data, while staged append responses write only unreachable immutable files that a separate coordinator must publish. The listener holds one additional file descriptor per open connection for exact-socket shutdown. + + **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.