diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 85a0ca2efe..dd50973d03 100644 --- a/crates/buzz-relay/src/config.rs +++ b/crates/buzz-relay/src/config.rs @@ -46,6 +46,10 @@ pub struct JoinPolicyConfig { pub version: String, } +/// Maximum configured jitter, leaving ten seconds of the hard-drain budget for +/// WebSocket close-frame delivery after the final delayed cancellation. +pub const MAX_DRAIN_JITTER_MS: u64 = 20_000; + /// Relay runtime configuration, loaded from environment variables. #[derive(Debug, Clone)] pub struct Config { @@ -60,6 +64,20 @@ pub struct Config { /// `0` (the default) disables bounded-staleness replica routing; see /// [`buzz_db::DbConfig::replica_read_max_age_ms`]. pub replica_read_max_age_ms: u64, + + /// Upper bound, in milliseconds, of the per-connection random delay applied + /// when sending the `1012 Service Restart` close frame during graceful + /// shutdown (`BUZZ_DRAIN_JITTER_MS`). Each live connection is closed after + /// an independent delay drawn uniformly from `[1, drain_jitter_ms]` when + /// jitter is enabled, which + /// spreads client reconnects across the window instead of releasing the + /// whole pod's sockets in one instant (the reconnect thundering herd that + /// drives DB pool-timeout bursts on rolling deploys). + /// + /// Default `0` reproduces the previous all-at-once close. Values above + /// [`MAX_DRAIN_JITTER_MS`] are capped, leaving headroom under the relay's + /// 30-second hard-drain timeout for close-frame delivery. + pub drain_jitter_ms: u64, /// Redis connection URL used by the pub/sub manager. pub redis_url: String, /// Maximum connections in the shared Redis pool. Defaults to 16. @@ -453,6 +471,25 @@ impl Config { Err(_) => 0, }; + // Drain jitter: 0 = off (default). Clamp oversized values so every + // delayed close is initiated with ten seconds left in the relay's + // hard-drain budget. An empty/whitespace-only value is treated as unset + // (jitter off), matching the sibling vars in this file — so setting the + // var to "" is a valid kill switch, not a crashloop. + let drain_jitter_ms = match std::env::var("BUZZ_DRAIN_JITTER_MS") { + Ok(raw) if raw.trim().is_empty() => 0, + Ok(raw) => raw + .trim() + .parse::() + .map_err(|_| { + ConfigError::InvalidValue( + "BUZZ_DRAIN_JITTER_MS must be a non-negative integer".to_string(), + ) + })? + .min(MAX_DRAIN_JITTER_MS), + Err(_) => 0, + }; + let redis_url = std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://localhost:6379".to_string()); @@ -934,6 +971,7 @@ impl Config { database_url, read_database_url, replica_read_max_age_ms, + drain_jitter_ms, redis_url, redis_pool_size, db_pool_size, @@ -1267,6 +1305,60 @@ mod tests { } } + #[test] + fn drain_jitter_defaults_off_and_rejects_junk() { + let _guard = ENV_MUTEX.lock().unwrap(); + let previous = std::env::var_os("BUZZ_DRAIN_JITTER_MS"); + + std::env::remove_var("BUZZ_DRAIN_JITTER_MS"); + let unset = Config::from_env().expect("config").drain_jitter_ms; + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", "20000"); + let set = Config::from_env().expect("config").drain_jitter_ms; + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", "60000"); + let capped = Config::from_env().expect("config").drain_jitter_ms; + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", "0"); + let zero = Config::from_env().expect("config").drain_jitter_ms; + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", "soon"); + let junk = Config::from_env(); + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", ""); + let empty = Config::from_env() + .expect("empty is a valid kill switch") + .drain_jitter_ms; + + std::env::set_var("BUZZ_DRAIN_JITTER_MS", " "); + let blank = Config::from_env() + .expect("whitespace-only is a valid kill switch") + .drain_jitter_ms; + + if let Some(value) = previous { + std::env::set_var("BUZZ_DRAIN_JITTER_MS", value); + } else { + std::env::remove_var("BUZZ_DRAIN_JITTER_MS"); + } + + assert_eq!(unset, 0, "drain jitter must default off"); + assert_eq!(set, MAX_DRAIN_JITTER_MS); + assert_eq!( + capped, MAX_DRAIN_JITTER_MS, + "oversized jitter leaves close-frame flush headroom" + ); + assert_eq!(zero, 0, "explicit 0 is off"); + assert!( + junk.is_err(), + "an unparsable jitter must fail loudly, not silently disable" + ); + assert_eq!( + empty, 0, + "an empty value is treated as unset — a kill switch, not a crashloop" + ); + assert_eq!(blank, 0, "a whitespace-only value is treated as unset"); + } + #[test] fn audit_logging_defaults_on_and_accepts_explicit_off() { let _guard = ENV_MUTEX.lock().unwrap(); diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index 96e266779f..72a7eb9126 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -29,6 +29,11 @@ const AUTH_TIMEOUT: Duration = Duration::from_secs(5); /// Shared mutable subscription map for a single WebSocket connection. pub(crate) type ConnectionSubscriptions = Arc>>>; +/// Request for the writer to flush a restart close and report the result. +pub(crate) struct RestartClose { + pub(crate) flushed: tokio::sync::oneshot::Sender, +} + /// Maximum outbound data frames buffered into the websocket sink before one flush. const MAX_WS_SEND_BATCH: usize = 64; @@ -161,6 +166,11 @@ async fn handle_active_connection( // even when the data buffer is full. let (ctrl_tx, ctrl_rx) = mpsc::channel::(8); + // Dedicated restart-close channel carries a flush acknowledgement. Keeping + // ordinary control frames unchanged avoids coupling heartbeat/ban traffic + // to graceful-shutdown delivery tracking. + let (restart_tx, restart_rx) = mpsc::channel::(1); + let backpressure_count = Arc::new(AtomicU8::new(0)); let subscriptions = Arc::new(Mutex::new(HashMap::new())); @@ -205,6 +215,7 @@ async fn handle_active_connection( conn_id, tx.clone(), ctrl_tx.clone(), + Some(restart_tx), cancel.clone(), conn.tenant.community(), Arc::clone(&backpressure_count), @@ -215,7 +226,7 @@ async fn handle_active_connection( let (ws_send, ws_recv) = socket.split(); let send_cancel = cancel.child_token(); - let send_task = tokio::spawn(send_loop(ws_send, rx, ctrl_rx, send_cancel)); + let send_task = tokio::spawn(send_loop(ws_send, rx, ctrl_rx, restart_rx, send_cancel)); let missed_pongs = Arc::new(AtomicU8::new(0)); let heartbeat_cancel = cancel.clone(); @@ -297,15 +308,17 @@ async fn send_loop( ws_send: futures_util::stream::SplitSink, data_rx: mpsc::Receiver, ctrl_rx: mpsc::Receiver, + restart_rx: mpsc::Receiver, cancel: CancellationToken, ) { - send_loop_inner(ws_send, data_rx, ctrl_rx, cancel).await; + send_loop_inner(ws_send, data_rx, ctrl_rx, restart_rx, cancel).await; } async fn send_loop_inner( mut ws_send: S, mut data_rx: mpsc::Receiver, mut ctrl_rx: mpsc::Receiver, + mut restart_rx: mpsc::Receiver, cancel: CancellationToken, ) where S: Sink + Unpin, @@ -319,9 +332,21 @@ async fn send_loop_inner( } tokio::select! { - // Biased: cancel > control > data. Cancel must win immediately - // so backpressure-triggered shutdown isn't starved by queued data. + // Biased: restart > cancel > ordinary control > data. A restart + // command owns shutdown delivery and must flush its 1012 before + // cancellation can fall back to an unacknowledged close. biased; + Some(restart) = restart_rx.recv() => { + let sent = ws_send + .send(WsMessage::Close(Some(axum::extract::ws::CloseFrame { + code: axum::extract::ws::close_code::RESTART, + reason: axum::extract::ws::Utf8Bytes::from_static("relay restarting"), + }))) + .await + .is_ok(); + let _ = restart.flushed.send(sent); + break; + } _ = cancel.cancelled() => { // Drain any queued control frames before closing. A ban // disconnect queues its `OK false "blocked: …"` reason frame on @@ -797,7 +822,8 @@ mod tests { } let (sink, state) = MockSink::new(Some(1)); - send_loop_inner(sink, data_rx, ctrl_rx, CancellationToken::new()).await; + let (_restart_tx, restart_rx) = mpsc::channel(1); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, CancellationToken::new()).await; let state = state.lock().expect("mock sink poisoned"); assert_eq!(state.flush_count, 1); @@ -817,7 +843,8 @@ mod tests { .expect("queue data frame"); let (sink, state) = MockSink::new(Some(1)); - send_loop_inner(sink, data_rx, ctrl_rx, CancellationToken::new()).await; + let (_restart_tx, restart_rx) = mpsc::channel(1); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, CancellationToken::new()).await; let state = state.lock().expect("mock sink poisoned"); assert_eq!(state.flush_count, 1); @@ -842,7 +869,8 @@ mod tests { .expect("queue control frame"); let (sink, state) = MockSink::new(Some(2)); - send_loop_inner(sink, data_rx, ctrl_rx, CancellationToken::new()).await; + let (_restart_tx, restart_rx) = mpsc::channel(1); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, CancellationToken::new()).await; let state = state.lock().expect("mock sink poisoned"); assert_eq!(state.flush_count, 2); @@ -852,6 +880,57 @@ mod tests { ); } + #[tokio::test] + async fn send_loop_acknowledges_restart_after_flushing_exactly_one_1012() { + let (_data_tx, data_rx) = mpsc::channel(1); + let (_ctrl_tx, ctrl_rx) = mpsc::channel(1); + let (restart_tx, restart_rx) = mpsc::channel(1); + let (flushed_tx, flushed_rx) = tokio::sync::oneshot::channel(); + restart_tx + .send(RestartClose { + flushed: flushed_tx, + }) + .await + .expect("queue restart close"); + + let (sink, state) = MockSink::new(None); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, CancellationToken::new()).await; + + assert_eq!(flushed_rx.await, Ok(true)); + let state = state.lock().expect("mock sink poisoned"); + assert_eq!(state.flush_count, 1, "ack follows the close flush"); + assert_eq!(state.messages.len(), 1, "writer exits after restart close"); + match &state.messages[0] { + WsMessage::Close(Some(close)) => { + assert_eq!(close.code, axum::extract::ws::close_code::RESTART); + assert_eq!(close.reason.as_str(), "relay restarting"); + } + other => panic!("expected one 1012 restart close, got {other:?}"), + } + } + + #[tokio::test] + async fn send_loop_reports_restart_flush_failure() { + let (_data_tx, data_rx) = mpsc::channel(1); + let (_ctrl_tx, ctrl_rx) = mpsc::channel(1); + let (restart_tx, restart_rx) = mpsc::channel(1); + let (flushed_tx, flushed_rx) = tokio::sync::oneshot::channel(); + restart_tx + .send(RestartClose { + flushed: flushed_tx, + }) + .await + .expect("queue restart close"); + + let (sink, state) = MockSink::new(Some(1)); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, CancellationToken::new()).await; + + assert_eq!(flushed_rx.await, Ok(false)); + let state = state.lock().expect("mock sink poisoned"); + assert_eq!(state.flush_count, 1); + assert_eq!(state.messages.len(), 1, "no fallback close is appended"); + } + #[tokio::test] async fn send_loop_flushes_queued_control_before_close_on_cancel() { // A ban disconnect queues its `OK false "blocked: …"` reason frame on @@ -871,7 +950,8 @@ mod tests { cancel.cancel(); let (sink, state) = MockSink::new(None); - send_loop_inner(sink, data_rx, ctrl_rx, cancel).await; + let (_restart_tx, restart_rx) = mpsc::channel(1); + send_loop_inner(sink, data_rx, ctrl_rx, restart_rx, cancel).await; let state = state.lock().expect("mock sink poisoned"); assert_eq!( diff --git a/crates/buzz-relay/src/handlers/event.rs b/crates/buzz-relay/src/handlers/event.rs index a9cdffcdec..a67797385b 100644 --- a/crates/buzz-relay/src/handlers/event.rs +++ b/crates/buzz-relay/src/handlers/event.rs @@ -1459,6 +1459,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, CancellationToken::new(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -2098,6 +2099,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, CancellationToken::new(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -2423,6 +2425,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, CancellationToken::new(), community_id, Arc::new(AtomicU8::new(0)), diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index 799cf9cf60..023a1970b7 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -17,7 +17,7 @@ use buzz_db::{Db, DbConfig}; use buzz_pubsub::PubSubManager; use buzz_search::SearchService; -use buzz_relay::config::Config; +use buzz_relay::config::{Config, MAX_DRAIN_JITTER_MS}; use buzz_relay::metrics as relay_metrics; use buzz_relay::router::{build_health_router, build_router}; use buzz_relay::state::AppState; @@ -1189,6 +1189,8 @@ async fn run_periodic_until_cancelled( /// │ → graceful drain (30s) → exit │ /// └─────────────────────────────────────────────────────────┘ /// ``` +const GRACEFUL_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + async fn serve( router: axum::Router, health_router: axum::Router, @@ -1207,8 +1209,9 @@ async fn serve( let (shutdown_tx, _) = tokio::sync::watch::channel(false); let shutdown_flag = Arc::clone(&state.shutting_down); let drain_conn_manager = Arc::clone(&state.conn_manager); + let drain_jitter_ms = state.config.drain_jitter_ms; let tx = shutdown_tx.clone(); - tokio::spawn(async move { + let shutdown_handle = tokio::spawn(async move { shutdown_signal().await; shutdown_flag.store(true, Ordering::Relaxed); info!("Shutdown signal received — readiness now returns 503"); @@ -1216,20 +1219,31 @@ async fn serve( tokio::time::sleep(std::time::Duration::from_secs(5)).await; info!("Starting graceful drain (30s timeout)"); let _ = tx.send(true); - // Tell every connected client to reconnect NOW. Without this, upgraded - // WebSocket connections outlive the listener drain: clients ride the - // dying pod until the forced exit below and only learn about the - // restart from a TCP reset. The 1012 close frame turns a 35s silent - // death into an immediate, well-attributed reconnect. - let closed = drain_conn_manager.drain_all(); + // Keep the original process-level backstop alive while listener and + // upgraded-socket shutdown proceeds. The caller aborts it only after + // Axum and the owned jitter drain have both completed. + let hard_shutdown = tokio::spawn(async { + tokio::time::sleep(GRACEFUL_DRAIN_TIMEOUT).await; + tracing::error!("Drain timeout exceeded — forcing exit"); + std::process::exit(1); + }); + let hard_shutdown_abort = hard_shutdown.abort_handle(); + // Stop accepting first, then close every live socket. Jitter off (the + // default) uses the original synchronous all-at-once drain; jitter on + // retains ownership of every delayed close until its 1012 frame has + // been flushed and acknowledged (or its send loop cancelled). + let closed = if drain_jitter_ms == 0 { + drain_conn_manager.drain_all() + } else { + drain_conn_manager.drain_all_jittered(drain_jitter_ms).await + }; info!( connections = closed, - "Sent restart close frame to all live WebSocket connections" + jitter_ms = drain_jitter_ms, + max_jitter_ms = MAX_DRAIN_JITTER_MS, + "Signalled restart close to all live WebSocket connections" ); - // Hard timeout: force exit if connections don't drain within 30s. - tokio::time::sleep(std::time::Duration::from_secs(30)).await; - tracing::error!("Drain timeout exceeded — forcing exit"); - std::process::exit(1); + hard_shutdown_abort }); let tcp_listener = tokio::net::TcpListener::bind(&config.bind_addr) @@ -1277,7 +1291,11 @@ async fn serve( .await .map_err(|e| anyhow::anyhow!("TCP server error: {e}"))?; + let hard_shutdown = shutdown_handle + .await + .map_err(|e| anyhow::anyhow!("Shutdown task failed: {e}"))?; uds_handle.abort(); + hard_shutdown.abort(); return Ok(()); } @@ -1298,6 +1316,10 @@ async fn serve( .await .map_err(|e| anyhow::anyhow!("Server error: {e}"))?; + let hard_shutdown = shutdown_handle + .await + .map_err(|e| anyhow::anyhow!("Shutdown task failed: {e}"))?; + hard_shutdown.abort(); Ok(()) } diff --git a/crates/buzz-relay/src/state.rs b/crates/buzz-relay/src/state.rs index 58a869a995..14a50df7b7 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -9,6 +9,7 @@ use std::time::Instant; use axum::body::Bytes; use axum::extract::ws::{Message as WsMessage, Utf8Bytes as WsUtf8Bytes}; use dashmap::DashMap; +use futures_util::future::join_all; use tokio::sync::mpsc; use tokio::sync::Semaphore; use tokio::task::JoinHandle; @@ -31,10 +32,13 @@ use deadpool_redis; use crate::audio::AudioRoomManager; use crate::config::Config; -use crate::connection::ConnectionSubscriptions; +use crate::connection::{ConnectionSubscriptions, RestartClose}; use crate::subscription::SubscriptionRegistry; pub(crate) type ScopedPubkeyKey = (CommunityId, [u8; 32]); + +/// Leaves headroom under the process-wide drain deadline for a stalled writer. +const RESTART_CLOSE_ACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); type SlidingWindowCounter = (u32, Instant); type ScopedRateLimiter = DashMap; @@ -45,6 +49,7 @@ struct ConnEntry { /// the send loop. Used to deliver a ban-disconnect frame that must reach /// the client before the socket is closed (see [`ConnectionManager::disconnect_pubkey`]). ctrl_tx: mpsc::Sender, + restart_tx: Option>, cancel: CancellationToken, /// Community resolved from the connection host at handshake. This is the /// receiver-side tenant label fan-out must compare against the event label. @@ -202,11 +207,12 @@ impl ConnectionManager { // Each argument is a distinct per-connection attribute stored verbatim in // `ConnEntry`; a params struct would only relocate the same fields. #[allow(clippy::too_many_arguments)] - pub fn register( + pub(crate) fn register( &self, conn_id: Uuid, tx: mpsc::Sender, ctrl_tx: mpsc::Sender, + restart_tx: Option>, cancel: CancellationToken, community_id: CommunityId, backpressure_count: Arc, @@ -220,6 +226,7 @@ impl ConnectionManager { ConnEntry { tx, ctrl_tx, + restart_tx, cancel, community_id, backpressure_count, @@ -231,7 +238,11 @@ impl ConnectionManager { // Insert-then-check pairs with drain_all's store-then-iterate: either // the drain iteration sees this entry, or this check sees the flag. // A registration that raced past the snapshot self-signals here, so - // no connection can outlive graceful shutdown unclosed. + // no connection can outlive graceful shutdown unclosed. A client that + // arrives mid-shutdown should be closed at once, so the self-signal + // always uses the immediate control-frame + cancel path regardless of + // whether jittered drain is enabled — jitter smears the sockets that + // were already established, not late arrivals. if self.draining.load(Ordering::SeqCst) { let _ = drain_ctrl_tx.try_send(Self::restart_close_frame()); drain_cancel.cancel(); @@ -335,6 +346,11 @@ impl ConnectionManager { /// Closes every live connection with a `1012 Service Restart` close frame. /// + /// This is the original, all-at-once drain, retained as the default path + /// (`BUZZ_DRAIN_JITTER_MS` unset or `0`). It is synchronous and returns as + /// soon as every close is queued and every connection cancelled, so the + /// caller's hard-drain timeout backstops delivery unchanged. + /// /// Called when graceful shutdown starts draining. Without this, upgraded /// WebSocket connections outlive the axum listener drain: clients ride the /// dying pod until the forced exit and then learn about the restart from a @@ -364,6 +380,82 @@ impl ConnectionManager { closed } + /// Closes every live connection with a `1012 Service Restart` frame, + /// spreading closes across `[1, jitter_ms]`. + /// + /// This is the jittered drain, used only when `BUZZ_DRAIN_JITTER_MS > 0`. + /// It is kept deliberately separate from [`Self::drain_all`] so that the + /// default (jitter-off) shutdown path is byte-for-byte the previously + /// shipped behavior; the new close-acknowledgement machinery only runs when + /// jitter is explicitly enabled. Once the jittered path is proven in + /// production for all cases, the two can be unified and the old one dropped. + /// + /// A pod under a rolling deploy can hold thousands of WebSocket sessions. + /// Closing them simultaneously ([`Self::drain_all`]) makes every client + /// reconnect at the same moment — a thundering herd that drives the DB + /// pool-timeout bursts observed on each roll. Delaying each connection's + /// close by an independent uniform random offset in `[1, jitter_ms]` + /// smears the reconnects across the window while keeping the well-attributed + /// 1012 close. + /// + /// Each delayed close is delivered over the connection's dedicated + /// [`RestartClose`] channel: the writer flushes the 1012 frame and + /// acknowledges the flush, so drain waits for confirmed delivery (up to + /// [`RESTART_CLOSE_ACK_TIMEOUT`]) rather than assuming it. If the channel is + /// full/closed or the ack times out, drain falls back to cancellation. + /// + /// The sticky drain flag is set before the first await, preserving + /// [`Self::drain_all`]'s shutdown-boundary race guarantee: a registration + /// that lands after the snapshot self-signals immediately (no jitter — a + /// client arriving mid-shutdown should be closed at once). The returned + /// future owns every delayed close, so the caller must await it before the + /// relay runtime is allowed to stop. + /// + /// Returns the number of connections signalled. + pub async fn drain_all_jittered(&self, jitter_ms: u64) -> usize { + // Store-then-snapshot pairs with register's insert-then-check: either + // the snapshot captures a registration, or it observes the sticky flag + // and self-signals immediately. + self.draining.store(true, Ordering::SeqCst); + let jitter_ms = jitter_ms.max(1); + let pending: Vec<_> = self + .connections + .iter() + .map(|entry| { + let ctrl_tx = entry.ctrl_tx.clone(); + let restart_tx = entry.restart_tx.clone(); + let cancel = entry.cancel.clone(); + let delay_ms = 1 + rand::random::() % jitter_ms; + async move { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + let Some(restart_tx) = restart_tx else { + // Unit-only registrations do not own a writer task. + let _ = ctrl_tx.try_send(Self::restart_close_frame()); + cancel.cancel(); + return; + }; + let (flushed_tx, flushed_rx) = tokio::sync::oneshot::channel(); + if restart_tx + .try_send(RestartClose { + flushed: flushed_tx, + }) + .is_err() + { + cancel.cancel(); + return; + } + let flushed = tokio::time::timeout(RESTART_CLOSE_ACK_TIMEOUT, flushed_rx).await; + if !matches!(flushed, Ok(Ok(true))) { + cancel.cancel(); + } + } + }) + .collect(); + let count = pending.len(); + join_all(pending).await; + count + } + /// The WS close frame announcing a graceful restart: 1012 Service Restart. fn restart_close_frame() -> WsMessage { WsMessage::Close(Some(axum::extract::ws::CloseFrame { @@ -1246,6 +1338,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::clone(&bp), @@ -1371,6 +1464,7 @@ mod tests { conn_id, tx, conn.ctrl_tx.clone(), + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::clone(&bp), @@ -1417,6 +1511,7 @@ mod tests { conn_a, tx_a, ctrl_tx_a, + None, CancellationToken::new(), community_a, Arc::new(AtomicU8::new(0)), @@ -1427,6 +1522,7 @@ mod tests { conn_b, tx_b, ctrl_tx_b, + None, CancellationToken::new(), community_b, Arc::new(AtomicU8::new(0)), @@ -1463,6 +1559,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel, buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), bp, @@ -1765,6 +1862,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1793,6 +1891,117 @@ mod tests { ); } + #[tokio::test] + async fn drain_all_jittered_waits_for_writer_acknowledgement_without_cancelling() { + let mgr = Arc::new(ConnectionManager::new()); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, _ctrl_rx) = mpsc::channel(8); + let (restart_tx, mut restart_rx) = mpsc::channel(1); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + Some(restart_tx), + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + + let drain_mgr = Arc::clone(&mgr); + let drain = tokio::spawn(async move { drain_mgr.drain_all_jittered(1).await }); + let restart = restart_rx.recv().await.expect("restart command delivered"); + assert!(!drain.is_finished(), "drain waits for the writer flush"); + restart.flushed.send(true).expect("acknowledge flush"); + + assert_eq!(drain.await.expect("drain task"), 1); + assert!( + !cancel.is_cancelled(), + "successful flush does not use cancellation fallback" + ); + } + + #[tokio::test] + async fn drain_all_jittered_cancels_when_restart_channel_is_full_or_closed() { + for keep_receiver in [true, false] { + let mgr = ConnectionManager::new(); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, _ctrl_rx) = mpsc::channel(8); + let (restart_tx, restart_rx) = mpsc::channel(1); + let (pending_tx, _pending_rx) = tokio::sync::oneshot::channel(); + if keep_receiver { + restart_tx + .try_send(RestartClose { + flushed: pending_tx, + }) + .expect("fill restart channel"); + } else { + drop(restart_rx); + } + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + Some(restart_tx), + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + + assert_eq!(mgr.drain_all_jittered(1).await, 1); + assert!( + cancel.is_cancelled(), + "unavailable writer cancels as fallback" + ); + } + } + + #[tokio::test(start_paused = true)] + async fn drain_all_jittered_cancels_when_flush_ack_times_out() { + // A writer that accepts the restart command but never acknowledges the + // flush (e.g. wedged mid-send) must not stall the drain: after + // RESTART_CLOSE_ACK_TIMEOUT the connection falls back to cancellation. + let mgr = Arc::new(ConnectionManager::new()); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, _ctrl_rx) = mpsc::channel(8); + let (restart_tx, mut restart_rx) = mpsc::channel(1); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + Some(restart_tx), + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + + let drain_mgr = Arc::clone(&mgr); + let drain = tokio::spawn(async move { drain_mgr.drain_all_jittered(1).await }); + // Take the restart command but hold the ack sender forever. + let restart = restart_rx.recv().await.expect("restart command delivered"); + assert!(!drain.is_finished(), "drain waits on the ack timeout"); + // Advance past the 5s ack timeout under paused time. + tokio::time::sleep(RESTART_CLOSE_ACK_TIMEOUT + std::time::Duration::from_millis(1)).await; + + assert_eq!(drain.await.expect("drain task"), 1); + assert!( + cancel.is_cancelled(), + "an un-acknowledged flush falls back to cancellation" + ); + drop(restart); + } + #[tokio::test] async fn drain_all_sends_restart_close_and_cancels_every_conn() { // Graceful shutdown must tell every live client to reconnect — across @@ -1809,6 +2018,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1860,6 +2070,7 @@ mod tests { conn_id, tx, ctrl_tx.clone(), + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -1907,6 +2118,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -1930,4 +2142,133 @@ mod tests { other => panic!("expected a restart close frame, got {other:?}"), } } + + #[tokio::test] + async fn drain_all_is_immediate() { + // The default (jitter-off) drain queues the frame and cancels + // synchronously — the frame is present the moment drain_all() returns. + let mgr = Arc::new(ConnectionManager::new()); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, mut ctrl_rx) = mpsc::channel(8); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + None, + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + + let closed = mgr.drain_all(); + + assert_eq!(closed, 1); + assert!(cancel.is_cancelled(), "default drain cancels synchronously"); + assert!( + matches!( + ctrl_rx + .try_recv() + .expect("close frame delivered synchronously"), + WsMessage::Close(Some(_)) + ), + "the restart close is queued before drain_all() returns" + ); + } + + #[tokio::test(start_paused = true)] + async fn drain_all_jittered_defers_close_until_within_jitter_window() { + // With jitter, the close is deferred within the owned drain future. + // The sticky drain flag is still set immediately, so a late + // registration self-signals with no delay. + let mgr = Arc::new(ConnectionManager::new()); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, mut ctrl_rx) = mpsc::channel(8); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + None, + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + + let jitter_ms = 20_000u64; + // Poll the owned drain through its first await. Dropping this future + // would drop the timers too; the shutdown path must retain and await it. + let drain = mgr.drain_all_jittered(jitter_ms); + tokio::pin!(drain); + assert!( + futures_util::poll!(&mut drain).is_pending(), + "jittered drain remains pending while its timers are owned" + ); + + // Not closed yet — the delayed drain is parked on its timer. + assert!( + !cancel.is_cancelled(), + "jittered close is deferred, not synchronous" + ); + assert!( + ctrl_rx.try_recv().is_err(), + "no close frame queued before the delay elapses" + ); + + // A registration racing past the snapshot still self-signals at once, + // regardless of jitter — clients arriving mid-shutdown are closed now. + let late_id = Uuid::new_v4(); + let (late_tx, _late_rx) = mpsc::channel(8); + let (late_ctrl_tx, mut late_ctrl_rx) = mpsc::channel(8); + let late_cancel = CancellationToken::new(); + mgr.register( + late_id, + late_tx, + late_ctrl_tx, + None, + late_cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + assert!( + late_cancel.is_cancelled(), + "late registration self-signals immediately, unaffected by jitter" + ); + assert!( + matches!( + late_ctrl_rx.try_recv().expect("late close frame"), + WsMessage::Close(Some(_)) + ), + "late registration gets the restart close with no delay" + ); + + // Advance past the whole jitter window; awaiting the owned drain must + // complete only after the deferred close has fired. + tokio::time::advance(std::time::Duration::from_millis(jitter_ms + 1)).await; + assert_eq!(drain.await, 1, "one captured connection drained"); + + assert!( + cancel.is_cancelled(), + "the jittered connection is closed within the jitter window" + ); + match ctrl_rx.try_recv().expect("deferred close frame delivered") { + WsMessage::Close(Some(close)) => { + assert_eq!( + close.code, + axum::extract::ws::close_code::RESTART, + "jittered close is still 1012 Service Restart" + ); + assert_eq!(close.reason.as_str(), "relay restarting"); + } + other => panic!("expected a restart close frame, got {other:?}"), + } + } }