From 86cf3f9ee6df69a819769a1e41aebc86f446e20c Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Thu, 7 May 2026 10:55:32 -0700 Subject: [PATCH] fix(server): unblock shutdown with active SSE streams Signal server cancellation before worker cleanup, wire long-lived SSE streams to the shutdown token, and backstop HTTP drain after five seconds so open browser streams cannot keep the listener alive indefinitely. --- lib/crates/fabro-server/src/serve.rs | 288 +++++++++++++++--- lib/crates/fabro-server/src/server.rs | 32 +- .../src/server/handler/completions.rs | 4 + .../fabro-server/src/server/handler/events.rs | 132 +++++++- lib/crates/fabro-server/src/server/tests.rs | 2 + lib/crates/fabro-server/src/test_support.rs | 2 + 6 files changed, 389 insertions(+), 71 deletions(-) diff --git a/lib/crates/fabro-server/src/serve.rs b/lib/crates/fabro-server/src/serve.rs index 4f0f26ea4..eb234c7db 100644 --- a/lib/crates/fabro-server/src/serve.rs +++ b/lib/crates/fabro-server/src/serve.rs @@ -1,3 +1,4 @@ +use std::future::{Future, IntoFuture}; use std::net::SocketAddr; use std::path::{Path, PathBuf}; use std::sync::{Arc, RwLock}; @@ -26,8 +27,9 @@ use object_store::local::LocalFileSystem; use object_store::memory::InMemory; use object_store::{ClientOptions, ObjectStore, RetryConfig}; use tokio::net::{TcpListener, UnixListener}; -use tokio::sync::watch; -use tokio::time::interval; +use tokio::task::JoinHandle; +use tokio::time::{interval, sleep}; +use tokio_util::sync::CancellationToken; use tracing::{error, info, warn}; use crate::canonical_origin::resolve_canonical_origin; @@ -45,6 +47,67 @@ use crate::static_files; pub const DEFAULT_TCP_PORT: u16 = 32276; type EnvLookup = Arc Option + Send + Sync>; +const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(5); + +async fn force_exit_after_shutdown(token: CancellationToken, grace: Duration) { + token.cancelled().await; + sleep(grace).await; +} + +async fn serve_until_shutdown( + serve_fut: F, + shutdown: CancellationToken, + grace: Duration, +) -> std::io::Result<()> +where + F: IntoFuture>, +{ + let fut = serve_fut.into_future(); + tokio::pin!(fut); + tokio::select! { + res = &mut fut => res, + () = force_exit_after_shutdown(shutdown, grace) => { + warn!( + grace_ms = grace.as_millis(), + "Graceful shutdown timed out; abandoning open connections" + ); + Ok(()) + } + } +} + +fn spawn_shutdown_orchestrator_inner( + shutdown: CancellationToken, + signal: S, + cleanup: C, +) -> JoinHandle<()> +where + S: Future + Send + 'static, + C: Future + Send + 'static, +{ + tokio::spawn(async move { + signal.await; + shutdown.cancel(); + cleanup.await; + }) +} + +fn spawn_shutdown_orchestrator( + shutdown: CancellationToken, + state: Arc, +) -> JoinHandle<()> { + let signal = async { + shutdown_signal().await; + set_server_title(ServerTitlePhase::Stopping, None); + }; + let cleanup = async move { + if let Err(err) = shutdown_active_workers(&state).await { + error!(error = %err, "Failed to stop active workers during shutdown"); + } + }; + spawn_shutdown_orchestrator_inner(shutdown, signal, cleanup) +} + #[derive(Debug, Clone)] pub(crate) struct ObjectStoreBuildOptions { pub client_options: ClientOptions, @@ -714,6 +777,7 @@ where let artifact_store = fabro_store::ArtifactStore::new(artifact_object_store, artifact_prefix); let env_lookup: EnvLookup = Arc::new(process_env_var); resolve_canonical_origin(&resolved_server_settings, &env_lookup).map_err(anyhow::Error::msg)?; + let shutdown = CancellationToken::new(); let state = build_app_state(AppStateConfig { resolved_settings: resolved_app_settings, registry_factory_override: None, @@ -725,6 +789,7 @@ where env_lookup, github_api_base_url: None, http_client: None, + shutdown: shutdown.clone(), })?; let reconciled = reconcile_incomplete_runs_on_startup(&state).await?; if reconciled > 0 { @@ -772,21 +837,10 @@ where ) .await?; - let (shutdown_tx, shutdown_rx) = watch::channel(false); - let shutdown_state = Arc::clone(&state); - tokio::spawn(async move { - shutdown_signal().await; - set_server_title(ServerTitlePhase::Stopping, None); - if let Err(err) = shutdown_active_workers(&shutdown_state).await { - error!(error = %err, "Failed to stop active workers during shutdown"); - } - let _ = shutdown_tx.send(true); - }); - spawn_auth_store_reapers( Arc::clone(&auth_code_store), Arc::clone(&auth_token_store), - shutdown_rx.clone(), + shutdown.clone(), ); // Spawn config polling task @@ -796,11 +850,15 @@ where let run_overrides_for_poll = run_overrides.clone(); let server_overrides_for_poll = server_overrides.clone(); let data_dir_for_poll = data_dir.clone(); + let shutdown_for_poll = shutdown.clone(); tokio::spawn(async move { let mut interval = interval(Duration::from_secs(5)); interval.tick().await; // skip first immediate tick loop { - interval.tick().await; + tokio::select! { + () = shutdown_for_poll.cancelled() => break, + _ = interval.tick() => {} + } match load_config_file::(config_path_for_poll.as_deref(), "settings.toml") { Ok(new_disk_settings) => { @@ -890,23 +948,38 @@ where None }; - match bound_listener.listener { + let cleanup_handle = spawn_shutdown_orchestrator(shutdown.clone(), Arc::clone(&state)); + + let serve_result = match bound_listener.listener { BoundListener::Unix(listener) => { announce_server_ready(&bind_addr, styles); - axum::serve(listener, router) - .with_graceful_shutdown(wait_for_shutdown(shutdown_rx.clone())) - .await?; + serve_until_shutdown( + axum::serve(listener, router).with_graceful_shutdown({ + let token = shutdown.clone(); + async move { token.cancelled().await } + }), + shutdown.clone(), + SHUTDOWN_GRACE_PERIOD, + ) + .await } BoundListener::Tcp(listener) => { announce_server_ready(&bind_addr, styles); - axum::serve( - listener, - router.into_make_service_with_connect_info::(), + serve_until_shutdown( + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown({ + let token = shutdown.clone(); + async move { token.cancelled().await } + }), + shutdown.clone(), + SHUTDOWN_GRACE_PERIOD, ) - .with_graceful_shutdown(wait_for_shutdown(shutdown_rx.clone())) - .await?; + .await } - } + }; #[cfg(debug_assertions)] if let Some(ref mut child) = watch_web_child { @@ -914,6 +987,16 @@ where let _ = child.wait(); } + if shutdown.is_cancelled() { + if let Err(join_err) = cleanup_handle.await { + warn!(error = %join_err, "Shutdown orchestrator task panicked"); + } + } else { + cleanup_handle.abort(); + } + + serve_result?; + if let Some(manager) = webhook_manager { manager.shutdown().await; } @@ -1016,25 +1099,18 @@ async fn shutdown_signal() { info!("Shutdown signal received, stopping server"); } -async fn wait_for_shutdown(mut shutdown_rx: watch::Receiver) { - if *shutdown_rx.borrow() { - return; - } - let _ = shutdown_rx.changed().await; -} - fn spawn_auth_store_reapers( auth_codes: Arc, auth_tokens: Arc, - shutdown_rx: watch::Receiver, + shutdown: CancellationToken, ) { - spawn_auth_code_reaper(auth_codes, shutdown_rx.clone()); - spawn_refresh_token_reaper(auth_tokens, shutdown_rx); + spawn_auth_code_reaper(auth_codes, shutdown.clone()); + spawn_refresh_token_reaper(auth_tokens, shutdown); } fn spawn_auth_code_reaper( auth_codes: Arc, - mut shutdown_rx: watch::Receiver, + shutdown: CancellationToken, ) { tokio::spawn(async move { let mut interval = interval(Duration::from_secs(30)); @@ -1042,7 +1118,7 @@ fn spawn_auth_code_reaper( loop { tokio::select! { - _ = shutdown_rx.changed() => break, + () = shutdown.cancelled() => break, _ = interval.tick() => { if let Err(err) = auth_codes.gc_expired(chrono::Utc::now()).await { warn!(error = %err, "Failed to garbage collect expired auth codes"); @@ -1055,7 +1131,7 @@ fn spawn_auth_code_reaper( fn spawn_refresh_token_reaper( auth_tokens: Arc, - mut shutdown_rx: watch::Receiver, + shutdown: CancellationToken, ) { tokio::spawn(async move { let mut interval = interval(Duration::from_hours(6)); @@ -1063,7 +1139,7 @@ fn spawn_refresh_token_reaper( loop { tokio::select! { - _ = shutdown_rx.changed() => break, + () = shutdown.cancelled() => break, _ = interval.tick() => { let cutoff = chrono::Utc::now() - chrono::Duration::days(7); if let Err(err) = auth_tokens.gc_expired(cutoff).await { @@ -1121,7 +1197,11 @@ fn server_bind_title(bind: &Bind) -> String { test uses tokio::net::TcpListener separately" )] mod tests { + use std::io; use std::path::PathBuf; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::Poll; use std::time::Duration; use fabro_config::bind::{Bind, BindRequest}; @@ -1130,14 +1210,17 @@ mod tests { use fabro_types::settings::interp::InterpString; use fabro_types::settings::server::{LogDestination, ObjectStoreSettings}; use fabro_util::Home; + use tokio::time::sleep; + use tokio_util::sync::CancellationToken; use super::{ - GitHubMetaResolver, ServeArgs, ServerTitlePhase, apply_effective_log_destination, - bind_tcp_host_with_fallback, build_local_object_store_with_preference, - build_object_store_from_settings_with_lookup, build_slatedb_store, - resolve_bind_request_from_server_settings, resolve_github_webhook_ip_allowlist, - resolve_startup_github_webhook_ip_allowlist, serve_overrides, server_bind_title, - server_title, + GitHubMetaResolver, SHUTDOWN_GRACE_PERIOD, ServeArgs, ServerTitlePhase, + apply_effective_log_destination, bind_tcp_host_with_fallback, + build_local_object_store_with_preference, build_object_store_from_settings_with_lookup, + build_slatedb_store, force_exit_after_shutdown, resolve_bind_request_from_server_settings, + resolve_github_webhook_ip_allowlist, resolve_startup_github_webhook_ip_allowlist, + serve_overrides, serve_until_shutdown, server_bind_title, server_title, + spawn_shutdown_orchestrator_inner, }; use crate::server::ResolvedAppStateSettings; @@ -1182,6 +1265,121 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn force_exit_after_shutdown_does_not_resolve_before_cancellation() { + let token = CancellationToken::new(); + let future = force_exit_after_shutdown(token, Duration::from_secs(5)); + tokio::pin!(future); + + tokio::time::advance(Duration::from_hours(1)).await; + + assert!(matches!(futures_util::poll!(&mut future), Poll::Pending)); + } + + #[tokio::test(start_paused = true)] + async fn force_exit_after_shutdown_resolves_grace_after_cancellation() { + let token = CancellationToken::new(); + let grace = Duration::from_secs(5); + let future = force_exit_after_shutdown(token.clone(), grace); + tokio::pin!(future); + + token.cancel(); + assert!(matches!(futures_util::poll!(&mut future), Poll::Pending)); + + tokio::time::advance( + grace + .checked_sub(Duration::from_millis(1)) + .expect("test grace should be longer than one millisecond"), + ) + .await; + assert!(matches!(futures_util::poll!(&mut future), Poll::Pending)); + + tokio::time::advance(Duration::from_millis(1)).await; + future.await; + } + + #[tokio::test(start_paused = true)] + async fn shutdown_orchestration_backstops_http_independent_of_cleanup() { + let shutdown = CancellationToken::new(); + let (signal_tx, signal_rx) = tokio::sync::oneshot::channel::<()>(); + let cleanup_started = Arc::new(AtomicBool::new(false)); + let cleanup_finished = Arc::new(AtomicBool::new(false)); + let cleanup_started_for_task = Arc::clone(&cleanup_started); + let cleanup_finished_for_task = Arc::clone(&cleanup_finished); + + let cleanup_handle = spawn_shutdown_orchestrator_inner( + shutdown.clone(), + async move { + signal_rx.await.expect("synthetic signal should be sent"); + }, + async move { + cleanup_started_for_task.store(true, Ordering::SeqCst); + sleep(Duration::from_mins(1)).await; + cleanup_finished_for_task.store(true, Ordering::SeqCst); + }, + ); + let serve_handle = tokio::spawn(serve_until_shutdown( + std::future::pending::>(), + shutdown.clone(), + SHUTDOWN_GRACE_PERIOD, + )); + + tokio::task::yield_now().await; + signal_tx + .send(()) + .expect("synthetic signal receiver should still be alive"); + tokio::task::yield_now().await; + + assert!(shutdown.is_cancelled()); + assert!(cleanup_started.load(Ordering::SeqCst)); + assert!(!serve_handle.is_finished()); + + tokio::time::advance(SHUTDOWN_GRACE_PERIOD).await; + tokio::task::yield_now().await; + + assert!(serve_handle.is_finished()); + serve_handle + .await + .expect("serve task should not panic") + .expect("serve timeout should be reported as a graceful shutdown"); + assert!(!cleanup_handle.is_finished()); + assert!(!cleanup_finished.load(Ordering::SeqCst)); + + tokio::time::advance(Duration::from_mins(1)).await; + cleanup_handle + .await + .expect("cleanup task should finish after its own work completes"); + assert!(cleanup_finished.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn serve_call_site_aborts_cleanup_on_early_serve_error() { + let shutdown = CancellationToken::new(); + let cleanup_handle = tokio::spawn(async { + std::future::pending::<()>().await; + }); + let serve_result = serve_until_shutdown( + async { Err(io::Error::new(io::ErrorKind::AddrInUse, "listener failed")) }, + shutdown.clone(), + SHUTDOWN_GRACE_PERIOD, + ) + .await; + + assert!(serve_result.is_err()); + assert!(!shutdown.is_cancelled()); + + if shutdown.is_cancelled() { + panic!("early serve error should not mark shutdown as cancelled"); + } else { + cleanup_handle.abort(); + } + + let join_err = cleanup_handle + .await + .expect_err("cleanup task should be aborted after early serve error"); + assert!(join_err.is_cancelled()); + } + #[test] fn runtime_server_settings_preserve_storage_dir_override() { let mut resolved = resolved_runtime_settings("_version = 1\n"); diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index c91c623d8..f08a63c7c 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -520,19 +520,20 @@ pub struct AppState { pub(crate) files_in_flight: FilesInFlight, pull_request_create_locks: PullRequestCreateLocks, - pub(crate) vault: Arc>, - pub(super) server_secrets: ServerSecrets, - pub(crate) llm_source: Arc, - manifest_run_defaults: RwLock>, - manifest_run_settings: RwLock>, - pub(crate) server_settings: RwLock>, - pub(crate) env_lookup: EnvLookup, + pub(crate) vault: Arc>, + pub(super) server_secrets: ServerSecrets, + pub(crate) llm_source: Arc, + manifest_run_defaults: RwLock>, + manifest_run_settings: RwLock>, + pub(crate) server_settings: RwLock>, + pub(crate) env_lookup: EnvLookup, pub(crate) github_api_base_url: String, - http_client: Option, - shutting_down: AtomicBool, - registry_factory_override: Option>, - slack_service: Option>, - slack_started: AtomicBool, + http_client: Option, + shutdown: CancellationToken, + shutting_down: AtomicBool, + registry_factory_override: Option>, + slack_service: Option>, + slack_started: AtomicBool, } type PullRequestCreateLocks = Arc>>>>; @@ -592,6 +593,7 @@ pub(crate) struct AppStateConfig { pub(crate) env_lookup: EnvLookup, pub(crate) github_api_base_url: Option, pub(crate) http_client: Option, + pub(crate) shutdown: CancellationToken, } #[derive(Clone)] @@ -817,6 +819,10 @@ impl AppState { self.scheduler_notify.notify_waiters(); } + pub(crate) fn shutdown_token(&self) -> CancellationToken { + self.shutdown.clone() + } + fn is_shutting_down(&self) -> bool { self.shutting_down.load(Ordering::Relaxed) } @@ -1511,6 +1517,7 @@ pub(crate) fn build_app_state(config: AppStateConfig) -> anyhow::Result anyhow::Result) }); + let stream = + futures_util::StreamExt::take_until(stream, state.shutdown_token().cancelled_owned()); Sse::new(stream) .keep_alive(KeepAlive::default()) @@ -265,6 +267,7 @@ async fn attach_run_events( }, }; let (sender, receiver) = mpsc::unbounded_channel(); + let shutdown = state.shutdown_token(); tokio::spawn(async move { let mut next_seq = start_seq; @@ -340,22 +343,31 @@ async fn attach_run_events( return; }; - while let Some(result) = live_stream.next().await { - let Ok(event) = result else { - return; - }; - let terminal = attach_event_is_terminal(&event); - if let Some(sse_event) = sse_event_from_store(&event) { - if sender - .send(Ok::(sse_event)) - .is_err() - { - return; + loop { + tokio::select! { + biased; + () = shutdown.cancelled() => break, + next = live_stream.next() => { + let Some(result) = next else { + return; + }; + let Ok(event) = result else { + return; + }; + let terminal = attach_event_is_terminal(&event); + if let Some(sse_event) = sse_event_from_store(&event) { + if sender + .send(Ok::(sse_event)) + .is_err() + { + return; + } + } + if terminal { + return; + } } } - if terminal { - return; - } } }); @@ -383,11 +395,16 @@ fn denied_lifecycle_event_name(body: &EventBody) -> Option<&'static str> { #[cfg(test)] mod stage_events_tests { + use std::time::Duration; + use axum::body::{Body, to_bytes}; use axum::http::{Request, StatusCode, header}; use fabro_store::EventPayload; use fabro_types::RunId; + use fabro_workflow::event as workflow_event; + use http_body_util::BodyExt; use serde_json::json; + use tokio::time::timeout; use tower::ServiceExt; use crate::test_support::{build_test_router, test_app_state}; @@ -442,6 +459,93 @@ mod stage_events_tests { serde_json::from_slice(&bytes).expect("response body should be valid JSON") } + fn assert_event_stream_response(response: &axum::response::Response) { + assert_eq!(response.status(), StatusCode::OK); + let content_type = response + .headers() + .get(header::CONTENT_TYPE) + .expect("SSE response should set content-type") + .to_str() + .expect("content-type should be valid UTF-8"); + assert!( + content_type.contains("text/event-stream"), + "expected text/event-stream content-type, got {content_type:?}" + ); + } + + async fn assert_sse_body_is_live(body: &mut Body) { + let result = timeout(Duration::from_millis(100), body.frame()).await; + assert!( + result.is_err(), + "SSE body should remain open before shutdown cancellation" + ); + } + + async fn assert_sse_body_completes_after_shutdown(mut body: Body) { + timeout(Duration::from_secs(1), async { + while let Some(frame) = body.frame().await { + frame.expect("SSE body frame should be readable"); + } + }) + .await + .expect("SSE body should complete promptly after shutdown cancellation"); + } + + #[tokio::test] + async fn attach_events_ends_when_shutdown_fires() { + let state = test_app_state(); + let app = build_test_router(state.clone()); + + let response = app + .oneshot(req_get("/api/v1/attach")) + .await + .expect("attach request should complete"); + assert_event_stream_response(&response); + + let mut body = response.into_body(); + assert_sse_body_is_live(&mut body).await; + + state.shutdown_token().cancel(); + + assert_sse_body_completes_after_shutdown(body).await; + } + + #[tokio::test] + async fn attach_run_events_ends_when_shutdown_fires() { + let state = test_app_state(); + let app = build_test_router(state.clone()); + let run_id = RunId::new(); + let run_store = state + .store_ref() + .create_run(&run_id) + .await + .expect("test run should be creatable"); + for event in [ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ] { + workflow_event::append_event(&run_store, &run_id, &event) + .await + .expect("run lifecycle event should append"); + } + + let response = app + .oneshot(req_get(&format!("/api/v1/runs/{run_id}/attach"))) + .await + .expect("run attach request should complete"); + assert_event_stream_response(&response); + + let mut body = response.into_body(); + assert_sse_body_is_live(&mut body).await; + + state.shutdown_token().cancel(); + + assert_sse_body_completes_after_shutdown(body).await; + } + async fn seed_run_with_mixed_events() -> (RunId, axum::Router) { let state = test_app_state(); let app = build_test_router(state.clone()); diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index cdad67c6f..de9c20afe 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -1755,6 +1755,7 @@ methods = ["dev-token"] env_lookup: default_env_lookup(), github_api_base_url: None, http_client: Some(fabro_http::test_http_client().expect("test HTTP client should build")), + shutdown: tokio_util::sync::CancellationToken::new(), }) else { panic!("build_app_state should require SESSION_SECRET") }; @@ -3247,6 +3248,7 @@ fn create_github_token_app_state_with_env_lookup( env_lookup: Arc::new(env_lookup), github_api_base_url, http_client: Some(fabro_http::test_http_client().expect("test HTTP client should build")), + shutdown: tokio_util::sync::CancellationToken::new(), }; let state = build_app_state(config).expect("test app state should build"); if let Some(token) = token { diff --git a/lib/crates/fabro-server/src/test_support.rs b/lib/crates/fabro-server/src/test_support.rs index d5eefb22a..a080f03f4 100644 --- a/lib/crates/fabro-server/src/test_support.rs +++ b/lib/crates/fabro-server/src/test_support.rs @@ -22,6 +22,7 @@ use fabro_types::{AuthMethod, IdpIdentity, ServerSettings}; use fabro_util::error::SharedError; use fabro_workflow::handler::HandlerRegistry; use object_store::memory::InMemory as MemoryObjectStore; +use tokio_util::sync::CancellationToken; use ulid::Ulid; use crate::auth; @@ -162,6 +163,7 @@ impl TestAppStateBuilder { http_client: Some( fabro_http::test_http_client().expect("test HTTP client should build"), ), + shutdown: CancellationToken::new(), }) .expect("test app state should build") }