From d6b9fedbc729a93b86b5991154e7c601d010b4b9 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 26 Jul 2026 23:27:21 +0800 Subject: [PATCH] fix(host): bound graceful shutdown and stop rejoining completed workers WorkerSupervisor::shutdown re-awaited every handle after a timeout, including those already joined inside the timeout block, which panics with 'JoinHandle polled after completion' whenever a cooperative worker exits before an uncooperative one under a shared timeout. Join each handle at most once against a shared absolute deadline and abort only the handles that miss it. The existing hung-task test never tripped this because the hung worker happened to sort first; add a regression test that pins the hung worker at the end of the join order. serve_router relied on axum's graceful shutdown, which waits for in-flight connections without bound while SSE bodies never observe cancellation, so one stalled client kept the process alive until SIGKILL. Race the server against a drain deadline that starts at the cancellation signal; past the deadline the remaining connections are abandoned and die with process exit. Co-authored-by: Claude --- backend/crates/gateway-host/src/serve.rs | 34 +++++++++++--- .../crates/gateway-host/src/workers/mod.rs | 20 +++----- .../crates/gateway-host/tests/workers/mod.rs | 47 +++++++++++++++++++ 3 files changed, 81 insertions(+), 20 deletions(-) diff --git a/backend/crates/gateway-host/src/serve.rs b/backend/crates/gateway-host/src/serve.rs index 8dae21d75..fce785e1b 100644 --- a/backend/crates/gateway-host/src/serve.rs +++ b/backend/crates/gateway-host/src/serve.rs @@ -121,16 +121,36 @@ pub(crate) async fn serve_router( wait_for_shutdown(&shutdown_cancellation).await; shutdown_connections.begin_draining(); }; - axum::serve( + // axum 的优雅关闭会无限等待存量连接结束，而 SSE 等长响应体不观察 + // 取消信号：慢客户端可以把关闭拖住直到被 SIGKILL。drain 截止点从 + // 取消信号起算，逾期即放弃等待，存量连接随进程退出终止。 + let drain_deadline = { + let cancellation = cancellation.clone(); + async move { + cancellation.cancelled().await; + tokio::time::sleep(drain_timeout).await; + } + }; + let serve = axum::serve( listener, router.into_make_service_with_connect_info::(), ) - .with_graceful_shutdown(shutdown) - .await - .map_err(ServeError::Serve)?; - - connections.begin_draining(); - connections.wait_until_idle(drain_timeout).await; + .with_graceful_shutdown(shutdown); + tokio::select! { + biased; + result = serve => { + result.map_err(ServeError::Serve)?; + connections.begin_draining(); + connections.wait_until_idle(drain_timeout).await; + } + () = drain_deadline => { + tracing::warn!( + target: "gateway_shutdown", + timeout_secs = drain_timeout.as_secs(), + "优雅 drain 超时，放弃等待存量连接" + ); + } + } Ok(()) } diff --git a/backend/crates/gateway-host/src/workers/mod.rs b/backend/crates/gateway-host/src/workers/mod.rs index f537be3f6..078c7230d 100644 --- a/backend/crates/gateway-host/src/workers/mod.rs +++ b/backend/crates/gateway-host/src/workers/mod.rs @@ -108,22 +108,16 @@ impl WorkerSupervisor { Ok(()) } - /// 先协作取消；超时后 abort 不合作任务，lease 依赖 TTL 兜底。 + /// 先协作取消；所有任务共享同一绝对截止点，逾期任务被 abort， + /// lease 依赖 TTL 兜底。每个句柄恰好 join 一次：已在截止点前退出的 + /// 不会被二次 poll，逾期的在 abort 后等待其真正终止。 pub async fn shutdown(&self, timeout: Duration) { self.cancellation.cancel(); - let mut handles = std::mem::take(&mut *lock_unpoisoned(&self.handles)); - let joined = tokio::time::timeout(timeout, async { - for (_, handle) in &mut handles { - let _ = handle.await; - } - }) - .await - .is_ok(); - if !joined { - for (_, handle) in &handles { + let handles = std::mem::take(&mut *lock_unpoisoned(&self.handles)); + let deadline = tokio::time::Instant::now() + timeout; + for (_, mut handle) in handles { + if tokio::time::timeout_at(deadline, &mut handle).await.is_err() { handle.abort(); - } - for (_, handle) in handles { let _ = handle.await; } } diff --git a/backend/crates/gateway-host/tests/workers/mod.rs b/backend/crates/gateway-host/tests/workers/mod.rs index 81a33f6ac..e717c3c14 100644 --- a/backend/crates/gateway-host/tests/workers/mod.rs +++ b/backend/crates/gateway-host/tests/workers/mod.rs @@ -626,6 +626,53 @@ async fn shutdown_timeout_aborts_uncooperative_task_and_drops_guard() { ); } +#[tokio::test] +async fn shutdown_timeout_must_not_rejoin_workers_that_already_exited() { + // 挂死任务放在 join 顺序最末的 kind（Retention），保证超时触发时 + // 排在它前面的合作任务句柄都已被 join 过一次。 + let hung = HungTask { + entered: Arc::new(AtomicBool::new(false)), + notification: Arc::new(Notify::new()), + }; + let cooperative = CancellationTask { + entered: Arc::new(AtomicUsize::new(0)), + observed: Arc::new(AtomicUsize::new(0)), + }; + let mut plan = ACTIVE_KINDS + .into_iter() + .map(|kind| { + if kind == WorkerKind::Retention { + registration(kind, kind.as_str(), hung.clone(), test_schedule(), true) + } else { + registration( + kind, + kind.as_str(), + cooperative.clone(), + test_schedule(), + true, + ) + } + }) + .collect::>(); + plan.push(WorkerContribution::Disabled { + kind: WorkerKind::OpsFlush, + reason: WorkerDisabledReason::NoBufferedOpsEvents, + }); + let leases = Arc::new(FakeLeasePort::default()); + let supervisor = supervisor(); + supervisor.start(plan, leases.clone()).expect("complete plan"); + hung.notification.notified().await; + wait_until(|| cooperative.entered.load(Ordering::SeqCst) == ACTIVE_KINDS.len() - 1).await; + + supervisor.shutdown(Duration::from_millis(100)).await; + + assert_eq!( + cooperative.observed.load(Ordering::SeqCst), + ACTIVE_KINDS.len() - 1 + ); + assert!(lock_unpoisoned(&leases.active_guards).is_empty()); +} + #[tokio::test] async fn running_worker_becomes_unhealthy_when_last_success_is_stale() { let task = SucceedsThenHangsTask {