@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1940c8441e

[Rust Frontend][gRPC] Add engine-aware health reporting (#48992)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Connor Carpenter <connorc@nvidia.com>

Connor Carpenter committed Jul 20, 2026 at 19:54 UTC 1940c8441eb39d7db88764b8e700d37e50573a84

9 files changed +333 -40

rust/Cargo.lock

+23 -8

@@ -4876,6 +4876,7 @@ dependencies = [

4876 "futures-core",

4877 "pin-project-lite",

4878 "tokio",

4879 + "tokio-util",

4880 ]

4881

4882 [[package]]

@@ -4937,9 +4938,9 @@ dependencies = [

4938

4939 [[package]]

4940 name = "tonic"

4940 -version = "0.14.5"

4941 +version = "0.14.6"

4942 source = "registry+https://github.com/rust-lang/crates.io-index"

4942 -checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"

4943 +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"

4944 dependencies = [

4945 "async-trait",

4946 "axum",

@@ -4966,9 +4967,9 @@ dependencies = [

4967

4968 [[package]]

4969 name = "tonic-build"

4969 -version = "0.14.5"

4970 +version = "0.14.6"

4971 source = "registry+https://github.com/rust-lang/crates.io-index"

4971 -checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"

4972 +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"

4973 dependencies = [

4974 "prettyplease",

4975 "proc-macro2",

@@ -4976,11 +4977,24 @@ dependencies = [

4977 "syn 2.0.117",

4978 ]

4979

4980 +[[package]]

4981 +name = "tonic-health"

4982 +version = "0.14.6"

4983 +source = "registry+https://github.com/rust-lang/crates.io-index"

4984 +checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"

4985 +dependencies = [

4986 + "prost",

4987 + "tokio",

4988 + "tokio-stream",

4989 + "tonic",

4990 + "tonic-prost",

4991 +]

4992 +

4993 [[package]]

4994 name = "tonic-prost"

4981 -version = "0.14.5"

4995 +version = "0.14.6"

4996 source = "registry+https://github.com/rust-lang/crates.io-index"

4983 -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"

4997 +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"

4998 dependencies = [

4999 "bytes",

5000 "prost",

@@ -4989,9 +5003,9 @@ dependencies = [

5003

5004 [[package]]

5005 name = "tonic-prost-build"

4992 -version = "0.14.5"

5006 +version = "0.14.6"

5007 source = "registry+https://github.com/rust-lang/crates.io-index"

4994 -checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"

5008 +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"

5009 dependencies = [

5010 "prettyplease",

5011 "proc-macro2",

@@ -5730,6 +5744,7 @@ dependencies = [

5744 "tokio-stream",

5745 "tokio-util",

5746 "tonic",

5747 + "tonic-health",

5748 "tonic-prost",

5749 "tonic-prost-build",

5750 "tower",

rust/Cargo.toml

+5 -4

@@ -118,10 +118,11 @@ tokio = { version = "1.47.1", features = [

118 tokio-openssl = "0.6"

119 tokio-stream = "0.1"

120 tokio-util = { version = "0.7.18", features = ["rt"] }

121 -tonic = "0.14.5"

122 -tonic-build = "0.14.5"

123 -tonic-prost = "0.14.5"

124 -tonic-prost-build = "0.14.5"

121 +tonic = "0.14.6"

122 +tonic-build = "0.14.6"

123 +tonic-health = "0.14.6"

124 +tonic-prost = "0.14.6"

125 +tonic-prost-build = "0.14.6"

126 tool-parser = "1.2.0"

127 tower = { version = "0.5.3", features = ["util"] }

128 tower-http = { version = "0.6.8", features = ["cors", "trace"] }

rust/src/engine-core-client/src/client.rs

@@ -460,6 +460,12 @@ impl EngineCoreClient {

460 self.inner.is_healthy()

461 }

462

463 + /// Subscribe to engine health changes. The current value is `true` while

464 + /// the client is healthy and changes permanently to `false` on failure.

465 + pub fn subscribe_health(&self) -> tokio::sync::watch::Receiver<bool> {

466 + self.inner.subscribe_health()

467 + }

468 +

469 /// Return the first persistent health error observed by the client, if any.

470 pub fn health_error(&self) -> Option<Arc<Error>> {

471 self.inner.health_error()

rust/src/engine-core-client/src/client/imp.rs

+20 -1

@@ -9,7 +9,7 @@ use arc_swap::ArcSwapOption;

9 use parking_lot::Mutex;

10 use thiserror_ext::AsReport as _;

11 use tokio::runtime::Handle;

12 -use tokio::sync::mpsc;

12 +use tokio::sync::{mpsc, watch};

13 use tracing::{debug, info, trace, warn};

14 use vllm_metrics::METRICS;

15 use zeromq::RouterSendHalf;

@@ -36,6 +36,7 @@ pub(crate) struct ClientInner {

36 request_reg: Mutex<RequestRegistry>,

37 utility_reg: Mutex<UtilityRegistry>,

38 health_error: ArcSwapOption<Error>,

39 + health_tx: watch::Sender<bool>,

40 }

42 impl ClientInner {

@@ -57,6 +58,7 @@ impl ClientInner {

58 request_reg: Mutex::new(RequestRegistry::new(engines)),

59 utility_reg: Mutex::new(UtilityRegistry::default()),

60 health_error: ArcSwapOption::empty(),

61 + health_tx: watch::Sender::new(true),

62 }

63 }

@@ -169,6 +171,7 @@ impl ClientInner {

171 /// persistent health error.

172 pub fn close_registries(&self, error: Arc<Error>) {

173 let persistent_error = self.record_health_error(error);

174 + self.publish_unhealthy();

175 let request_senders = self.request_reg.lock().close();

176 let utility_senders = self.utility_reg.lock().close();

177

@@ -191,6 +194,12 @@ impl ClientInner {

194 self.health_error.load().is_none()

195 }

196

197 + /// Subscribe to engine health changes. The current value is `true` while

198 + /// the client is healthy and changes permanently to `false` on failure.

199 + pub fn subscribe_health(&self) -> watch::Receiver<bool> {

200 + self.health_tx.subscribe()

201 + }

202 +

203 /// Resolve one utility output to the waiting caller. Returns `true` if a

204 /// waiting caller existed.

205 pub fn resolve_utility_output(&self, output: UtilityOutput) -> bool {

@@ -280,6 +289,11 @@ impl ClientInner {

289 .expect("health error must be recorded before registries close")

290 }

291

292 + /// Publish the sticky healthy-to-unhealthy transition.

293 + fn publish_unhealthy(&self) {

294 + self.health_tx.send_if_modified(|healthy| std::mem::replace(healthy, false));

295 + }

296 +

297 /// Assert there is a recorded health error and return a `Shared` variant

298 /// wrapping it for error returns when the client is already closed.

299 fn closed_error(&self) -> Error {

@@ -461,13 +475,18 @@ mod tests {

475 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

476 async fn close_registries_records_first_health_error_only() {

477 let inner = test_inner().await;

478 + let mut health = inner.subscribe_health();

479 + assert!(*health.borrow());

480

481 inner.close_registries(Arc::new(Error::EngineCoreDead));

482 + health.changed().await.expect("health sender remains open");

483 assert!(!inner.is_healthy());

484 + assert!(!*health.borrow());

485 assert!(matches!(

486 inner.health_error().as_deref(),

487 Some(Error::EngineCoreDead)

488 ));

489 + assert!(!*inner.subscribe_health().borrow());

490

491 inner.close_registries(Arc::new(client_closed!("shutdown")));

492 assert!(matches!(

rust/src/server/Cargo.toml

@@ -35,6 +35,7 @@ tokio-openssl.workspace = true

35 tokio-stream.workspace = true

36 tokio-util.workspace = true

37 tonic.workspace = true

38 +tonic-health.workspace = true

39 tonic-prost.workspace = true

40 tower.workspace = true

41 tower-http.workspace = true

rust/src/server/src/grpc/health.rs new

+70

@@ -0,0 +1,70 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +use tokio::sync::watch;

5 +use tokio_util::sync::CancellationToken;

6 +use tonic::server::NamedService;

7 +use tonic_health::ServingStatus;

8 +use tonic_health::server::HealthReporter;

9 +use tracing::{info, warn};

10 +

11 +use super::GenerateGrpcService;

12 +

13 +pub(crate) async fn monitor_health(

14 + mut health_reporter: HealthReporter,

15 + mut engine_health: watch::Receiver<bool>,

16 + shutdown: CancellationToken,

17 +) {

18 + let generate_service = GenerateGrpcService::NAME;

19 + let status = ServingStatus::NotServing;

20 + let health_event_first = tokio::select! {

21 + result = engine_health.wait_for(|healthy| !*healthy) => {

22 + match result {

23 + Ok(_) => warn!(

24 + generate_service,

25 + overall_service = true,

26 + status = ?status,

27 + reason = "engine_unhealthy",

28 + "marking gRPC health services as not serving"

29 + ),

30 + Err(error) => warn!(

31 + %error,

32 + generate_service,

33 + overall_service = true,

34 + status = ?status,

35 + reason = "health_channel_closed",

36 + "engine health channel closed; marking gRPC health services as not serving"

37 + ),

38 + }

39 + true

40 + }

41 + _ = shutdown.cancelled() => {

42 + info!(

43 + generate_service,

44 + overall_service = true,

45 + status = ?status,

46 + reason = "server_shutdown",

47 + "server shutting down; marking gRPC health services as not serving"

48 + );

49 + false

50 + }

51 + };

52 +

53 + health_reporter.set_not_serving::<GenerateGrpcService>().await;

54 + // Generate is currently the only engine-backed gRPC service, so overall

55 + // server health intentionally mirrors it.

56 + health_reporter.set_service_status("", status).await;

57 +

58 + if health_event_first {

59 + shutdown.cancelled().await;

60 + info!(

61 + generate_service,

62 + overall_service = true,

63 + reason = "server_shutdown",

64 + "server shutting down; closing gRPC health watches"

65 + );

66 + }

67 +

68 + health_reporter.clear_service_status(generate_service).await;

69 + health_reporter.clear_service_status("").await;

70 +}

rust/src/server/src/grpc/mod.rs

@@ -4,6 +4,7 @@

4 //! gRPC Generate service backed by the shared [`vllm_text::TextLlm`] facade.

6 mod convert;

7 +mod health;

9 use std::pin::Pin;

10 use std::sync::Arc;

@@ -24,8 +25,11 @@ pub mod pb {

25 tonic::include_proto!("vllm");

26 }

28 +pub(crate) use health::monitor_health;

29 pub use pb::generate_server::GenerateServer;

31 +pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;

32 +

33 #[cfg(test)]

34 mod tests;

rust/src/server/src/grpc/tests.rs

+175 -12

@@ -16,6 +16,10 @@ use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};

16 use tokio::net::TcpStream;

17 use tokio_openssl::SslStream;

18 use tonic::transport::{Channel, Endpoint, Server as TonicServer, Uri};

19 +use tonic_health::pb::HealthCheckRequest;

20 +use tonic_health::pb::health_check_response::ServingStatus as HealthServingStatus;

21 +use tonic_health::pb::health_client::HealthClient;

22 +use tonic_health::server::health_reporter;

23 use tower::service_fn;

24 use vllm_chat::{

25 ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,

@@ -200,7 +204,11 @@ impl ChatRenderer for FakeTextBackend {

204 async fn setup_grpc_service(

205 engine_id: impl Into<EngineId>,

206 output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,

203 -) -> (GenerateServer<GenerateServiceImpl>, MockEngineTask) {

207 +) -> (

208 + GenerateServer<GenerateServiceImpl>,

209 + tokio::sync::watch::Receiver<bool>,

210 + MockEngineTask,

211 +) {

212 let ipc = IpcNamespace::new().expect("create ipc namespace");

213 let handshake_address = ipc.handshake_endpoint();

214 let engine_id = engine_id.into();

@@ -232,6 +240,7 @@ async fn setup_grpc_service(

240 )

241 .await

242 .expect("connect client");

243 + let engine_health = client.subscribe_health();

244

245 let chat = ChatLlm::from_shared_backend(

246 test_llm(client),

@@ -240,6 +249,7 @@ async fn setup_grpc_service(

249 let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));

250 (

251 GenerateServer::new(GenerateServiceImpl::new(state)),

252 + engine_health,

253 engine_task,

254 )

255 }

@@ -254,25 +264,51 @@ async fn grpc_test_server(

264 tokio::task::JoinHandle<()>,

265 MockEngineTask,

266 ) {

257 - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;

267 + let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;

268 + let (channel, server_task) = start_grpc_test_server(

269 + svc,

270 + engine_health,

271 + tokio_util::sync::CancellationToken::new(),

272 + )

273 + .await;

274 + (GenerateClient::new(channel), server_task, engine_task)

275 +}

276 +

277 +async fn start_grpc_test_server(

278 + generate_service: GenerateServer<GenerateServiceImpl>,

279 + engine_health: tokio::sync::watch::Receiver<bool>,

280 + shutdown: tokio_util::sync::CancellationToken,

281 +) -> (Channel, tokio::task::JoinHandle<()>) {

282 + let (health_reporter, health_service) = health_reporter();

283 + health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;

284

285 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");

286 let addr = listener.local_addr().expect("local addr");

287

288 let server_task = tokio::spawn(async move {

289 let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));

264 - TonicServer::builder()

265 - .add_service(svc)

266 - .serve_with_incoming(incoming)

267 - .await

268 - .expect("grpc server");

290 + let server = TonicServer::builder()

291 + .add_service(health_service)

292 + .add_service(generate_service)

293 + .serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());

294 + let health_monitor =

295 + super::monitor_health(health_reporter, engine_health, shutdown.clone());

296 + let server = async move {

297 + let result = server.await;

298 + shutdown.cancel();

299 + result

300 + };

301 + let (server_result, ()) = tokio::join!(server, health_monitor);

302 + server_result.expect("grpc server");

303 });

304

271 - let grpc_client = GenerateClient::connect(format!("http://{addr}"))

305 + let channel = Endpoint::from_shared(format!("http://{addr}"))

306 + .expect("grpc endpoint")

307 + .connect()

308 .await

273 - .expect("connect grpc client");

309 + .expect("connect grpc channel");

310

275 - (grpc_client, server_task, engine_task)

311 + (channel, server_task)

312 }

313

314 /// Spin up a TLS gRPC server (server cert from `certs`, `cert_reqs` mTLS mode).

@@ -283,7 +319,7 @@ async fn grpc_tls_test_server(

319 certs: &TestCerts,

320 cert_reqs: i32,

321 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {

286 - let (svc, engine_task) = setup_grpc_service(engine_id, output_specs).await;

322 + let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;

323 let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))

324 .expect("build grpc tls config");

325

@@ -373,7 +409,8 @@ async fn grpc_server_with_keepalive(

409 engine_id: impl Into<EngineId>,

410 keepalive: Option<Duration>,

411 ) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {

376 - let (svc, engine_task) = setup_grpc_service(engine_id, default_stream_output_specs()).await;

412 + let (svc, _engine_health, engine_task) =

413 + setup_grpc_service(engine_id, default_stream_output_specs()).await;

414

415 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");

416 let addr = listener.local_addr().expect("local addr").to_string();

@@ -1035,3 +1072,129 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {

1072

1073 server_task.abort();

1074 }

1075 +

1076 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

1077 +#[serial]

1078 +async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {

1079 + let (generate_service, _connected_engine_health, _engine_task) =

1080 + setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;

1081 + let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);

1082 + let (channel, server_task) = start_grpc_test_server(

1083 + generate_service,

1084 + engine_health,

1085 + tokio_util::sync::CancellationToken::new(),

1086 + )

1087 + .await;

1088 + let mut health_client = HealthClient::new(channel);

1089 +

1090 + let mut health_streams = Vec::new();

1091 + for service in ["vllm.Generate", ""] {

1092 + let service_label = if service.is_empty() {

1093 + "overall"

1094 + } else {

1095 + service

1096 + };

1097 + let mut stream = health_client

1098 + .watch(HealthCheckRequest {

1099 + service: service.to_string(),

1100 + })

1101 + .await

1102 + .unwrap_or_else(|error| {

1103 + panic!("failed to start health watch for {service_label}: {error}")

1104 + })

1105 + .into_inner();

1106 + let initial = stream

1107 + .message()

1108 + .await

1109 + .unwrap_or_else(|error| {

1110 + panic!("failed to read initial health status for {service_label}: {error}")

1111 + })

1112 + .unwrap_or_else(|| {

1113 + panic!("health watch for {service_label} ended before its initial status")

1114 + });

1115 + assert_eq!(

1116 + initial.status,

1117 + HealthServingStatus::Serving as i32,

1118 + "unexpected initial health status for {service_label}"

1119 + );

1120 + health_streams.push((service_label, stream));

1121 + }

1122 +

1123 + engine_health_tx.send(false).expect("publish unhealthy engine state");

1124 +

1125 + for (service_label, mut stream) in health_streams {

1126 + let update = tokio::time::timeout(Duration::from_secs(2), stream.message())

1127 + .await

1128 + .unwrap_or_else(|_| panic!("timed out waiting for health update for {service_label}"))

1129 + .unwrap_or_else(|error| {

1130 + panic!("failed to read health update for {service_label}: {error}")

1131 + })

1132 + .unwrap_or_else(|| panic!("health watch for {service_label} ended before its update"));

1133 + assert_eq!(

1134 + update.status,

1135 + HealthServingStatus::NotServing as i32,

1136 + "unexpected health status for {service_label}"

1137 + );

1138 + }

1139 +

1140 + server_task.abort();

1141 +}

1142 +

1143 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

1144 +#[serial]

1145 +async fn grpc_health_watch_closes_on_graceful_shutdown() {

1146 + let (generate_service, engine_health, _engine_task) = setup_grpc_service(

1147 + b"engine-grpc-health-shutdown",

1148 + default_stream_output_specs(),

1149 + )

1150 + .await;

1151 + let shutdown = tokio_util::sync::CancellationToken::new();

1152 + let (channel, server_task) =

1153 + start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;

1154 + let mut health_client = HealthClient::new(channel);

1155 + let mut stream = health_client

1156 + .watch(HealthCheckRequest {

1157 + service: "vllm.Generate".to_string(),

1158 + })

1159 + .await

1160 + .expect("start health watch for vllm.Generate")

1161 + .into_inner();

1162 +

1163 + let initial = stream

1164 + .message()

1165 + .await

1166 + .expect("read initial health status for vllm.Generate")

1167 + .expect("health watch ended before its initial status");

1168 + assert_eq!(

1169 + initial.status,

1170 + HealthServingStatus::Serving as i32,

1171 + "unexpected initial health status for vllm.Generate"

1172 + );

1173 +

1174 + shutdown.cancel();

1175 +

1176 + let update = tokio::time::timeout(Duration::from_secs(2), stream.message())

1177 + .await

1178 + .expect("timed out waiting for shutdown health update for vllm.Generate")

1179 + .expect("failed to read shutdown health update for vllm.Generate")

1180 + .expect("health watch ended before its shutdown update");

1181 + assert_eq!(

1182 + update.status,

1183 + HealthServingStatus::NotServing as i32,

1184 + "unexpected shutdown health status for vllm.Generate"

1185 + );

1186 +

1187 + let stream_end = tokio::time::timeout(Duration::from_secs(2), stream.message())

1188 + .await

1189 + .expect("timed out waiting for vllm.Generate health watch to close")

1190 + .expect("failed while closing vllm.Generate health watch");

1191 + assert!(

1192 + stream_end.is_none(),

1193 + "vllm.Generate health watch remained open"

1194 + );

1195 +

1196 + tokio::time::timeout(Duration::from_secs(2), server_task)

1197 + .await

1198 + .expect("timed out waiting for gRPC server shutdown")

1199 + .expect("gRPC server task failed");

1200 +}

rust/src/server/src/lib.rs

+29 -15

@@ -39,6 +39,7 @@ use tokio::net::TcpListener;

39 use tokio::time::{Instant, sleep_until};

40 use tokio_util::sync::CancellationToken;

41 use tonic::transport::Server as TonicServer;

42 +use tonic_health::server::health_reporter;

43 use tower::ServiceExt as _;

44 use tracing::{info, trace, warn};

45 use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends};

@@ -203,14 +204,19 @@ where

204 .map(tls::build_grpc_server_config)

205 .transpose()

206 .context("invalid gRPC TLS configuration")?;

206 - let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));

207 + let (health_reporter, health_service) = health_reporter();

208 + let engine_health = state.engine_core_client().subscribe_health();

209 + health_reporter.set_serving::<grpc::GenerateGrpcService>().await;

210 + let generate_service =

211 + grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));

212 let svc = TonicServer::builder()

213 .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL))

214 .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))

215 .layer(middleware::request_runtime_layer(state.clone()))

211 - .add_service(svc);

216 + .add_service(health_service)

217 + .add_service(generate_service);

218 info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");

213 - Some((grpc_listener, svc, grpc_tls))

219 + Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))

220 } else {

221 None

222 };

@@ -294,7 +300,8 @@ where

300 let server_shutdown = server_shutdown.clone();

301 let force_shutdown = force_shutdown.clone();

302 async move {

297 - let Some((grpc_listener, svc, grpc_tls)) = grpc_setup else {

303 + let Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = grpc_setup

304 + else {

305 // No gRPC configured: just wait for shutdown so we do not race the

306 // join! by resolving early and tripping the cancellation token.

307 shutdown.cancelled().await;

@@ -304,19 +311,26 @@ where

311 Some(context) => MaybeTlsListener::tls(grpc_listener, context),

312 None => MaybeTlsListener::plain(grpc_listener),

313 };

307 - let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned());

308 -

309 - let result = tokio::select! {

310 - result = server => {

311 - result.context("gRPC server failed")

312 - }

313 - _ = force_shutdown.cancelled() => {

314 - warn!("gRPC graceful shutdown deadline elapsed; aborting server");

315 - Ok(())

316 - }

314 + let server =

315 + svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());

316 + let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown);

317 +

318 + let server = async move {

319 + let result = tokio::select! {

320 + result = server => {

321 + result.context("gRPC server failed")

322 + }

323 + _ = force_shutdown.cancelled() => {

324 + warn!("gRPC graceful shutdown deadline elapsed; aborting server");

325 + Ok(())

326 + }

327 + };

328 +

329 + server_shutdown.cancel();

330 + result

331 };

332

319 - server_shutdown.cancel();

333 + let (result, ()) = tokio::join!(server, health_monitor);

334 result

335 }

336 };