From abf0d6cf84d766be428ad8929f18a87cccde07de Mon Sep 17 00:00:00 2001 From: kipavy Date: Wed, 11 Nov 2026 14:10:51 +0000 Subject: [PATCH] feat(server): drain connections on SIGTERM before exiting The server installed no signal handler. As PID 1 in its container it ignored SIGTERM, so every stop waited out Docker's 10 s timeout and was then killed mid-request, leaving shared terminal sessions marked live in the database with their key bytes still stored. On SIGTERM or Ctrl-C the server now stops accepting connections, lets in-flight requests finish, ends open sync streams (clients already reconnect on their own), and tells every shared-terminal socket the session ended with reason "server_restart" before closing it with 1012. The host's normal disconnect cleanup then marks the session ended. Upgraded sockets are tracked separately because hyper stops tracking a connection once it upgrades. Whatever is still open after SHUTDOWN_GRACE_SECS (default 8) is dropped and the process exits. Refs VoltiusApp/voltius#800 --- .env.example | 7 ++ Cargo.lock | 2 + Cargo.toml | 1 + src/main.rs | 22 +++++-- src/routes/sync.rs | 3 + src/routes/terminal.rs | 78 ++++++++++++++++++---- src/shutdown.rs | 142 +++++++++++++++++++++++++++++++++++++++++ 7 files changed, 236 insertions(+), 19 deletions(-) create mode 100644 src/shutdown.rs diff --git a/.env.example b/.env.example index d36b68b..c494c14 100644 --- a/.env.example +++ b/.env.example @@ -122,3 +122,10 @@ JWT_SECRET= # generate with: openssl rand -hex 32 # terminal sessions, rate limits and presence are per-process, so two # instances disagree without erroring. See CLAUDE.md before setting this. # ALLOW_MULTIPLE_INSTANCES=true + + +# Shutdown ---------------------------------------------------------------- +# Seconds a stopping server waits for open requests, sync streams and +# terminal sockets to close. Keep it below the container stop timeout +# (Docker's default is 10 s), or the server is killed mid-drain. +# SHUTDOWN_GRACE_SECS=8 diff --git a/Cargo.lock b/Cargo.lock index 4def701..9e0b0c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2939,6 +2939,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -3228,6 +3229,7 @@ dependencies = [ "tokio", "tokio-stream", "tokio-tungstenite", + "tokio-util", "tower", "tower-http", "tracing", diff --git a/Cargo.toml b/Cargo.toml index e471cf0..cd4149e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ base64 = "0.22" dotenvy = "0.15" rand = "0.8.6" tokio-stream = { version = "0.1.18", features = ["sync"] } +tokio-util = { version = "0.7", features = ["rt"] } hmac = "0.12" sha2 = "0.10" hex = "0.4" diff --git a/src/main.rs b/src/main.rs index d5d33ec..18f2b3e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ mod permissions; mod rate_limit; mod routes; mod self_host; +mod shutdown; mod single_instance; mod session_grants; mod team_authority; @@ -41,6 +42,7 @@ use rate_limit::{ SyncRateLimiter, WaitlistRateLimiter, }; use routes::audit::AuditClientRateLimiter; +use std::future::IntoFuture; use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; @@ -236,6 +238,7 @@ async fn main() { } }); } + let drain = shutdown::Drain::new(); let presence_map: PresenceMap = Arc::new(DashMap::new()); let usage_map: UsageMap = Arc::new(DashMap::new()); @@ -801,6 +804,7 @@ async fn main() { .layer(Extension(terminal_manager.clone())) .layer(Extension(presence_map.clone())) .layer(Extension(usage_map.clone())) + .layer(Extension(drain.clone())) .layer(Extension(min_client_version)); // Admin routes — auth + admin check, no rate limit (internal tool) @@ -869,7 +873,8 @@ async fn main() { get(routes::terminal::ws_handler), ) .layer(Extension(terminal_manager)) - .layer(Extension(notifier)); + .layer(Extension(notifier)) + .layer(Extension(drain.clone())); let app = Router::new() .merge(public) @@ -932,14 +937,19 @@ async fn main() { } }; - if let Err(error) = axum::serve( + let grace = Duration::from_secs( + std::env::var("SHUTDOWN_GRACE_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(8), + ); + let server = axum::serve( listener, app.into_make_service_with_connect_info::(), ) - .await - { - tracing::error!(%error, "Server exited with an error"); - } + .with_graceful_shutdown(drain.stopping()) + .into_future(); + drain.serve(server, grace).await; } #[cfg(test)] diff --git a/src/routes/sync.rs b/src/routes/sync.rs index cbd4fce..4b28fb4 100644 --- a/src/routes/sync.rs +++ b/src/routes/sync.rs @@ -290,6 +290,7 @@ pub async fn sync_stream( axum::Extension(notifier): axum::Extension, axum::Extension(presence): axum::Extension, axum::Extension(usage): axum::Extension, + axum::Extension(drain): axum::Extension, ) -> Sse>> { let user_id = auth.0; @@ -352,6 +353,8 @@ pub async fn sync_stream( } }); + // Ending the stream is what lets a draining server finish this response. + let stream = futures_util::StreamExt::take_until(stream, drain.stopping()); Sse::new(stream).keep_alive( KeepAlive::new() .interval(Duration::from_secs(30)) diff --git a/src/routes/terminal.rs b/src/routes/terminal.rs index 4f25574..28ece53 100644 --- a/src/routes/terminal.rs +++ b/src/routes/terminal.rs @@ -1373,6 +1373,7 @@ pub async fn ws_handler( State(pool): State, Extension(manager): Extension, Extension(notifier): Extension, + Extension(drain): Extension, ) -> impl IntoResponse { let user_id = match validate_token(&query.token, "access") { Ok(claims) => claims.sub, @@ -1385,7 +1386,7 @@ pub async fn ws_handler( let handle = resolve_participant_handle(&pool, user_id).await; ws.on_upgrade(move |socket| { - handle_socket( + drain.track(handle_socket( socket, session_id, user_id, @@ -1394,7 +1395,8 @@ pub async fn ws_handler( pool, manager, notifier, - ) + drain.stopping(), + )) }) } @@ -1550,6 +1552,13 @@ async fn check_admission( } const ACCESS_REVOKED_FRAME: &str = r#"{"type":"session_ended","reason":"access_revoked"}"#; +const SERVER_RESTART_FRAME: &str = r#"{"type":"session_ended","reason":"server_restart"}"#; + +async fn end_socket + Unpin>(sender: &mut S, frame: &str, close_code: Option) { + let _ = sender.send(Message::Text(frame.to_string())).await; + let close = close_code.map(|code| axum::extract::ws::CloseFrame { code, reason: "".into() }); + let _ = sender.send(Message::Close(close)).await; +} #[allow(clippy::too_many_arguments)] async fn handle_socket( @@ -1561,6 +1570,7 @@ async fn handle_socket( pool: PgPool, manager: TerminalManager, notifier: crate::sync_notifier::SyncNotifier, + stopping: tokio_util::sync::WaitForCancellationFutureOwned, ) { // Subscribed before admission so a revoke racing this join still triggers a recheck. let access_rx = manager.sessions.lock().await.get(&session_id).map(|s| s.access_changed.subscribe()); @@ -1669,15 +1679,19 @@ async fn handle_socket( let pool = pool.clone(); let manager = manager.clone(); tokio::spawn(async move { + tokio::pin!(stopping); loop { tokio::select! { + () = &mut stopping => { + end_socket(&mut ws_sender, SERVER_RESTART_FRAME, Some(axum::extract::ws::close_code::RESTART)).await; + break; + } Ok(()) = access_rx.changed() => { let admission = check_admission(&pool, &manager, session_id, user_id, invite_token.as_deref()).await; if matches!(admission, Some(Admission { authorized: false, .. })) { info!(session_id = %session_id, user_id = %user_id, "WS participant lost access"); - let _ = ws_sender.send(Message::Text(ACCESS_REVOKED_FRAME.to_string())).await; - let _ = ws_sender.send(Message::Close(None)).await; + end_socket(&mut ws_sender, ACCESS_REVOKED_FRAME, None).await; break; } } @@ -3872,20 +3886,23 @@ mod live_revoke_tests { use crate::test_pool_or_skip; use crate::test_support::{add_member, member_with_role, seed_session, seed_team, seed_user, set_user_tier}; use axum::{routing::get, Router}; + use crate::shutdown::Drain; use tokio_tungstenite::tungstenite::Message as WsMessage; type Client = tokio_tungstenite::WebSocketStream>; - async fn serve(pool: PgPool, manager: TerminalManager) -> std::net::SocketAddr { + async fn serve(pool: PgPool, manager: TerminalManager) -> (std::net::SocketAddr, Drain) { + let drain = Drain::new(); let app = Router::new() .route("/v1/terminal-sessions/:id/ws", get(ws_handler)) .layer(Extension(manager)) .layer(Extension(SyncNotifier::new())) + .layer(Extension(drain.clone())) .with_state(pool); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - addr + (addr, drain) } async fn try_connect(addr: std::net::SocketAddr, session_id: Uuid, user: Uuid) -> Option { @@ -3954,7 +3971,7 @@ mod live_revoke_tests { teams: &[Uuid], host: Uuid, mate: Uuid, - ) -> (TerminalManager, Uuid, std::net::SocketAddr) { + ) -> (TerminalManager, Uuid, std::net::SocketAddr, Drain) { set_user_tier(pool, host, "teams").await; for team in teams { add_member(pool, *team, host).await; @@ -3966,8 +3983,8 @@ mod live_revoke_tests { grant_invitee(pool, &SyncNotifier::new(), &manager, &crate::test_support::default_knock_limiter(), session_id, host, mate, "wrapped") .await .unwrap(); - let addr = serve(pool.clone(), manager.clone()).await; - (manager, session_id, addr) + let (addr, drain) = serve(pool.clone(), manager.clone()).await; + (manager, session_id, addr, drain) } async fn direct_session_with_joined_mate( @@ -3976,7 +3993,7 @@ mod live_revoke_tests { host: Uuid, mate: Uuid, ) -> (TerminalManager, Uuid, Client, Client) { - let (manager, session_id, addr) = direct_session_inviting(pool, teams, host, mate).await; + let (manager, session_id, addr, _) = direct_session_inviting(pool, teams, host, mate).await; let host_ws = connect(addr, session_id, host).await; let mate_ws = connect(addr, session_id, mate).await; (manager, session_id, host_ws, mate_ws) @@ -4032,13 +4049,48 @@ mod live_revoke_tests { assert_evicted(&mut host_ws).await; } + #[tokio::test] + async fn draining_ends_the_session_for_everyone_attached() { + let pool = test_pool_or_skip!(); + let host = seed_user(&pool).await; + let mate = seed_user(&pool).await; + let team = seed_team(&pool, host).await; + let (manager, session_id, addr, drain) = direct_session_inviting(&pool, &[team], host, mate).await; + let mut host_ws = connect(addr, session_id, host).await; + let mut mate_ws = connect(addr, session_id, mate).await; + + drain.begin(); + + for ws in [&mut host_ws, &mut mate_ws] { + let ended = next_of_type(ws, "session_ended").await.expect("told the session ended"); + assert_eq!(ended["reason"], "server_restart"); + let close = tokio::time::timeout(std::time::Duration::from_secs(5), ws.next()).await.unwrap(); + let Some(Ok(WsMessage::Close(Some(frame)))) = close else { panic!("expected a close frame, got {close:?}") }; + assert_eq!(u16::from(frame.code), 1012); + } + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while manager.sessions.lock().await.contains_key(&session_id) { + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("the host's cleanup still runs"); + let ended_at: Option> = + sqlx::query_scalar("SELECT ended_at FROM terminal_sessions WHERE id = $1") + .bind(session_id) + .fetch_one(&pool) + .await + .unwrap(); + assert!(ended_at.is_some(), "a drained session must not stay live in the database"); + } + #[tokio::test] async fn a_banned_invitee_cannot_join_with_an_unexpired_token() { let pool = test_pool_or_skip!(); let host = seed_user(&pool).await; let mate = seed_user(&pool).await; let team = seed_team(&pool, host).await; - let (_manager, session_id, addr) = direct_session_inviting(&pool, &[team], host, mate).await; + let (_manager, session_id, addr, _) = direct_session_inviting(&pool, &[team], host, mate).await; let _host_ws = connect(addr, session_id, host).await; lose_account(&pool, mate, BAN).await; @@ -4144,7 +4196,7 @@ mod live_revoke_tests { state.visibility = "vault".to_string(); state.vault_ids = vec![team]; } - let addr = serve(pool.clone(), manager.clone()).await; + let (addr, _) = serve(pool.clone(), manager.clone()).await; let host_ws = connect(addr, session_id, host).await; let mate_ws = connect(addr, session_id, mate).await; VaultSession { manager, team, host, mate, session_id, host_ws, mate_ws } @@ -4296,7 +4348,7 @@ mod live_revoke_tests { let manager = TerminalManager::new(); manager.insert_test_session(session_id, host).await; manager.sessions.lock().await.get_mut(&session_id).unwrap().visibility = "invite_link".to_string(); - let addr = serve(pool.clone(), manager.clone()).await; + let (addr, _) = serve(pool.clone(), manager.clone()).await; let mut host_ws = connect(addr, session_id, host).await; let mut guest_ws = try_connect_with(addr, session_id, guest, &format!("&invite_token={token}")) diff --git a/src/shutdown.rs b/src/shutdown.rs new file mode 100644 index 0000000..c46e1b0 --- /dev/null +++ b/src/shutdown.rs @@ -0,0 +1,142 @@ +use std::future::Future; +use std::time::Duration; +use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned}; +use tokio_util::task::{task_tracker::TrackedFuture, TaskTracker}; +use tracing::{error, info, warn}; + +#[derive(Clone, Default)] +pub struct Drain { + token: CancellationToken, + sockets: TaskTracker, +} + +impl Drain { + pub fn new() -> Self { + Self::default() + } + + pub fn stopping(&self) -> WaitForCancellationFutureOwned { + self.token.clone().cancelled_owned() + } + + // Upgraded sockets leave hyper's connection tracking, so they are waited on here. + pub fn track(&self, socket: F) -> TrackedFuture { + self.sockets.track_future(socket) + } + + pub fn begin(&self) { + self.token.cancel(); + } + + pub async fn serve(self, server: impl Future>, grace: Duration) { + tokio::spawn({ + let drain = self.clone(); + async move { + stop_signal().await; + info!( + grace_secs = grace.as_secs(), + "Stop signal received, draining connections" + ); + drain.begin(); + } + }); + + let drained = async { + if let Err(error) = server.await { + error!(%error, "Server exited with an error"); + } + self.begin(); + self.sockets.close(); + self.sockets.wait().await; + }; + let deadline = async { + self.token.cancelled().await; + tokio::time::sleep(grace).await; + }; + tokio::select! { + () = drained => info!("Server drained"), + () = deadline => warn!(open_sockets = self.sockets.len(), "Drain deadline reached, exiting"), + } + } +} + +async fn stop_signal() { + let interrupt = async { + let _ = tokio::signal::ctrl_c().await; + }; + #[cfg(unix)] + let terminate = async { + match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) { + Ok(mut sigterm) => { + sigterm.recv().await; + } + Err(error) => { + error!(%error, "Could not listen for SIGTERM"); + std::future::pending::<()>().await; + } + } + }; + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + tokio::select! { + () = interrupt => {} + () = terminate => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn waits_for_tracked_sockets_once_the_server_returns() { + let drain = Drain::new(); + let (release, released) = tokio::sync::oneshot::channel::<()>(); + let socket = tokio::spawn(drain.track(async move { + let _ = released.await; + })); + let serving = tokio::spawn( + drain + .clone() + .serve(async { Ok(()) }, Duration::from_secs(30)), + ); + + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!serving.is_finished(), "must wait for the open socket"); + release.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(5), serving) + .await + .unwrap() + .unwrap(); + socket.await.unwrap(); + } + + #[tokio::test(start_paused = true)] + async fn gives_up_on_stuck_sockets_after_the_grace_period() { + let drain = Drain::new(); + let _stuck = tokio::spawn(drain.track(std::future::pending::<()>())); + let mut serving = tokio::spawn( + drain + .clone() + .serve(async { Ok(()) }, Duration::from_secs(8)), + ); + + assert!(tokio::time::timeout(Duration::from_secs(7), &mut serving) + .await + .is_err()); + tokio::time::timeout(Duration::from_secs(2), serving) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn stopping_resolves_when_the_drain_begins() { + let drain = Drain::new(); + let stopping = drain.stopping(); + drain.begin(); + tokio::time::timeout(Duration::from_secs(1), stopping) + .await + .unwrap(); + } +}