From 6c8db62079921e0978774c0d1d59fca489ce2c49 Mon Sep 17 00:00:00 2001 From: MinJun Choi Date: Mon, 2 Nov 2026 21:54:21 +0900 Subject: [PATCH] [Bugfix][Rust Frontend] Stop accepting gRPC connections when shutdown starts tonic keeps the listener passed to serve_with_incoming_shutdown until every connection has drained, so during a --shutdown-timeout drain new gRPC clients were accepted and never answered. The incoming stream now drops the listener when shutdown starts, as the HTTP server already does. Signed-off-by: MinJun Choi --- rust/src/server/src/grpc/tests.rs | 93 ++++++++++++++++++++++++++++++- rust/src/server/src/lib.rs | 7 ++- rust/src/server/src/listener.rs | 18 ++++++ 3 files changed, 116 insertions(+), 2 deletions(-) diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 521b3bafa8c8..61736ea1cc92 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -494,7 +494,10 @@ async fn start_grpc_test_server( .add_service(health_service) .add_service(control_service) .add_service(inference_service) - .serve_with_incoming_shutdown(incoming, stop); + .serve_with_incoming_shutdown( + incoming.incoming_until(stop), + std::future::pending(), + ); async move { let result = server.await; server_shutdown.cancel(); @@ -2013,6 +2016,94 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn shutdown_refuses_new_connections_while_streams_drain() { + let (submitted_tx, submitted_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let (inference_service, _control_service, engine_health, _engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-drain-refuses-new-connections".to_vec(), + default_ready_response(), + Arc::new(FakeTextBackend::default()), + None, + move |dealer, push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + submitted_tx.send(()).expect("signal submitted request"); + release_rx.await.expect("release engine output"); + send_outputs( + push, + engine_outputs_for_request( + &request.request_id, + vec![(vec![1, 2], Some(EngineCoreFinishReason::Length))], + ), + ) + .await; + }) + }, + ) + .await; + let (health_reporter, _health_service) = health_reporter(); + let health_monitor = HealthMonitor::register(health_reporter, engine_health).await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener"); + let addr = listener.local_addr().expect("local addr"); + let incoming = MaybeTlsListener::plain(Listener::Tcp(listener)); + let shutdown = tokio_util::sync::CancellationToken::new(); + let server_task = tokio::spawn(health_monitor.serve( + shutdown.clone(), + std::future::pending(), + |stop| { + TonicServer::builder() + .add_service(inference_service) + .serve_with_incoming_shutdown(incoming.incoming_until(stop), std::future::pending()) + }, + )); + let channel = Endpoint::from_shared(format!("http://{addr}")) + .expect("grpc endpoint") + .connect() + .await + .expect("connect grpc channel"); + let mut stream = InferenceClient::new(channel) + .generate_stream(pb::GenerateRequest { + request_id: "test-drain-refuses-new-connections".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + ..Default::default() + }) + .await + .expect("start generation") + .into_inner(); + tokio::time::timeout(Duration::from_secs(5), submitted_rx) + .await + .expect("timed out waiting for engine submission") + .expect("engine submission signal"); + + shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + match TcpStream::connect(addr).await { + Err(error) if error.kind() == io::ErrorKind::ConnectionRefused => break, + _ => tokio::time::sleep(Duration::from_millis(10)).await, + } + } + }) + .await + .expect("new connections must be refused once shutdown starts"); + + release_tx.send(()).expect("release engine output"); + while stream.message().await.expect("open stream must finish").is_some() {} + tokio::time::timeout(Duration::from_secs(5), server_task) + .await + .expect("server must stop once the open stream is done") + .expect("server task") + .expect("grpc server is never forced to abort") + .expect("grpc server"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn dropping_generate_stream_aborts_before_engine_output() { diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index ed556ecc5c24..0aa244bca73c 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -377,7 +377,12 @@ where None => MaybeTlsListener::plain(grpc_listener), }; let server = health_monitor.serve(shutdown, force_shutdown.cancelled(), |stop| { - let server = svc.serve_with_incoming_shutdown(incoming, stop); + // The stream ending starts the drain; the signal never fires and only + // keeps tonic's graceful drain on. + let server = svc.serve_with_incoming_shutdown( + incoming.incoming_until(stop), + std::future::pending(), + ); async move { let result = server.await.context("gRPC server failed"); server_shutdown.cancel(); diff --git a/rust/src/server/src/listener.rs b/rust/src/server/src/listener.rs index 5af274d581b3..0c41910ea69e 100644 --- a/rust/src/server/src/listener.rs +++ b/rust/src/server/src/listener.rs @@ -204,6 +204,24 @@ impl MaybeTlsListener { .listen(listener), ) } + + /// Accept connections until `stop` completes, then drop the listener so new + /// connections are refused while the open ones drain. + pub fn incoming_until( + self, + stop: impl Future, + ) -> impl futures::Stream> { + futures::stream::unfold( + (self, Box::pin(stop)), + |(mut listener, mut stop)| async move { + tokio::select! { + biased; + () = &mut stop => None, + io = futures::StreamExt::next(&mut listener) => io.map(|io| (io, (listener, stop))), + } + }, + ) + } } /// Listener I/O type that may be either a plain TCP/UDS stream or a TLS stream over it.