diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 85a0ca2efe..87bd8f424e 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,22 @@ 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. + let drain_jitter_ms = match std::env::var("BUZZ_DRAIN_JITTER_MS") { + 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 +968,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 +1302,45 @@ 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(); + + 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" + ); + } + #[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..e26f5b3c6a 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,25 @@ 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 retain ownership of every delayed close + // until its 1012 frame has been queued and its send loop cancelled. + let closed = drain_conn_manager.drain_all(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 +1285,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 +1310,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..b581ec4766 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,17 +207,19 @@ 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, subscriptions: ConnectionSubscriptions, grace_limit: u8, ) { + let drain_restart_tx = restart_tx.clone(); let drain_ctrl_tx = ctrl_tx.clone(); let drain_cancel = cancel.clone(); self.connections.insert( @@ -220,6 +227,7 @@ impl ConnectionManager { ConnEntry { tx, ctrl_tx, + restart_tx, cancel, community_id, backpressure_count, @@ -233,8 +241,15 @@ impl ConnectionManager { // A registration that raced past the snapshot self-signals here, so // no connection can outlive graceful shutdown unclosed. if self.draining.load(Ordering::SeqCst) { - let _ = drain_ctrl_tx.try_send(Self::restart_close_frame()); - drain_cancel.cancel(); + if let Some(restart_tx) = drain_restart_tx { + let (flushed, _acknowledgement) = tokio::sync::oneshot::channel(); + if restart_tx.try_send(RestartClose { flushed }).is_err() { + drain_cancel.cancel(); + } + } else { + let _ = drain_ctrl_tx.try_send(Self::restart_close_frame()); + drain_cancel.cancel(); + } } } @@ -333,35 +348,75 @@ impl ConnectionManager { closed } - /// Closes every live connection with a `1012 Service Restart` close frame. + /// Closes every live connection with a `1012 Service Restart` frame, + /// spreading closes across `[1, jitter_ms]` when jitter is enabled. + /// + /// 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. /// - /// 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 - /// TCP reset (or, on an abrupt kill, from up to 60s of stall-watchdog - /// silence). The explicit close frame tells them to reconnect immediately - /// — and that the disconnect is a restart, not a policy action. + /// 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. /// - /// Uses the "queue frame on ctrl, then cancel" idiom (see - /// [`ConnectionManager::disconnect_pubkey`]): the send loop drains queued - /// control frames — including this close — before its cancel branch closes - /// the socket. Best-effort: a full control buffer still gets the close via - /// cancel, just without the restart code. + /// `jitter_ms == 0` queues and cancels every captured connection before the + /// first await, preserving the previous all-at-once behavior. /// /// Returns the number of connections signalled. - pub fn drain_all(&self) -> usize { - // Store-then-iterate pairs with register's insert-then-check: a - // registration that misses this iteration observes the flag and - // self-signals instead. The flag is sticky — drain is one-way. + pub async fn drain_all(&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 frame = Self::restart_close_frame(); - let mut closed = 0usize; - for entry in self.connections.iter() { - let _ = entry.ctrl_tx.try_send(frame.clone()); - entry.cancel.cancel(); - closed += 1; - } - closed + 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 = if jitter_ms == 0 { + 0 + } else { + 1 + rand::random::() % jitter_ms + }; + async move { + if delay_ms > 0 { + 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. @@ -1246,6 +1301,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::clone(&bp), @@ -1371,6 +1427,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 +1474,7 @@ mod tests { conn_a, tx_a, ctrl_tx_a, + None, CancellationToken::new(), community_a, Arc::new(AtomicU8::new(0)), @@ -1427,6 +1485,7 @@ mod tests { conn_b, tx_b, ctrl_tx_b, + None, CancellationToken::new(), community_b, Arc::new(AtomicU8::new(0)), @@ -1463,6 +1522,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel, buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), bp, @@ -1765,6 +1825,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1793,6 +1854,78 @@ mod tests { ); } + #[tokio::test] + async fn drain_all_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(0).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_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(0).await, 1); + assert!( + cancel.is_cancelled(), + "unavailable writer cancels as fallback" + ); + } + } + #[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 +1942,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1825,7 +1959,7 @@ mod tests { Uuid::from_u128(0xb), )); - let closed = mgr.drain_all(); + let closed = mgr.drain_all(0).await; assert_eq!(closed, 2, "every connection is signalled, no tenant fence"); assert!(cancel_a.is_cancelled(), "community-A session is cancelled"); @@ -1860,6 +1994,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)), @@ -1871,7 +2006,7 @@ mod tests { .try_send(WsMessage::Text("wedge".into())) .expect("fill control channel"); - let closed = mgr.drain_all(); + let closed = mgr.drain_all(0).await; assert_eq!(closed, 1); assert!( @@ -1896,7 +2031,7 @@ mod tests { let mgr = ConnectionManager::new(); // Drain with zero connections — sets the sticky flag. - assert_eq!(mgr.drain_all(), 0); + assert_eq!(mgr.drain_all(0).await, 0); // Late registration lands after the snapshot. let conn_id = Uuid::new_v4(); @@ -1907,6 +2042,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 +2066,133 @@ mod tests { other => panic!("expected a restart close frame, got {other:?}"), } } + + #[tokio::test] + async fn drain_all_zero_is_immediate() { + // jitter_ms == 0 preserves the original all-at-once behavior: the + // frame is queued and cancellation fires when the future resolves. + 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(0).await; + + assert_eq!(closed, 1); + assert!(cancel.is_cancelled(), "zero jitter cancels synchronously"); + assert!( + matches!( + ctrl_rx + .try_recv() + .expect("close frame delivered synchronously"), + WsMessage::Close(Some(_)) + ), + "the restart close is queued before drain_all(0) returns" + ); + } + + #[tokio::test(start_paused = true)] + async fn drain_all_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(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:?}"), + } + } }