From 344bfdc486537f5d0036263a5d5ff4cfa40acda3 Mon Sep 17 00:00:00 2001 From: npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w <80ffe7a6e22722984fb7fbf4c72b026030a745b08fae413253b1d993a5f19a30@buzz.block.builderlab.xyz> Date: Mon, 3 Aug 2026 11:02:31 -0400 Subject: [PATCH 1/4] relay: fuzz WebSocket 1012 restart-close timing on graceful drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On SIGTERM the relay sends every live WebSocket a 1012 Service Restart close frame via ConnectionManager::drain_all(), all in one instant. On a pod holding thousands of sessions that makes every client reconnect simultaneously — the thundering herd behind the DB pool-timeout bursts seen on each rolling deploy. Add BUZZ_DRAIN_JITTER_MS (default 0 = unchanged all-at-once behavior). When set, drain_all_jittered spreads each connection's restart close over an independent uniform delay in [0, jitter_ms], smearing reconnects across the window while keeping the well-attributed 1012 close and the 30s hard-drain backstop. The sticky drain flag is still set synchronously before returning, so the shutdown-boundary race guarantee holds: a registration that lands after the snapshot self-signals immediately with no jitter. jitter_ms == 0 delegates to drain_all for identical synchronous behavior and no task spawns. Tests: config default-off/reject-junk parse; jittered drain defers the close until within the window (paused-time), zero-jitter stays synchronous, and a late registration self-signals unaffected by jitter. Co-authored-by: npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w <80ffe7a6e22722984fb7fbf4c72b026030a745b08fae413253b1d993a5f19a30@buzz.block.builderlab.xyz> Signed-off-by: npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w <80ffe7a6e22722984fb7fbf4c72b026030a745b08fae413253b1d993a5f19a30@buzz.block.builderlab.xyz> --- crates/buzz-relay/src/config.rs | 58 +++++++++++ crates/buzz-relay/src/main.rs | 11 +- crates/buzz-relay/src/state.rs | 172 ++++++++++++++++++++++++++++++++ 3 files changed, 239 insertions(+), 2 deletions(-) diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 85a0ca2efe..244f20ce89 100644 --- a/crates/buzz-relay/src/config.rs +++ b/crates/buzz-relay/src/config.rs @@ -60,6 +60,19 @@ 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 `[0, drain_jitter_ms]`, 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 are capped + /// at shutdown to leave headroom under the 30s hard-drain timeout; keep it + /// well below `terminationGracePeriodSeconds`. + 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 +466,18 @@ impl Config { Err(_) => 0, }; + // Drain jitter: 0 = off (default). Non-negative parse, same shape as + // the replica-read budget above. Bound is enforced (capped) at + // shutdown, not here, so config never fails on a large value. + 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(), + ) + })?, + Err(_) => 0, + }; + let redis_url = std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://localhost:6379".to_string()); @@ -934,6 +959,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 +1293,38 @@ 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", "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, 20_000); + 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/main.rs b/crates/buzz-relay/src/main.rs index 799cf9cf60..9d334d3750 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -1207,6 +1207,7 @@ 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 { shutdown_signal().await; @@ -1221,10 +1222,16 @@ async fn serve( // 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(); + // + // With BUZZ_DRAIN_JITTER_MS > 0, the closes are spread across the + // jitter window instead of firing all at once, so a pod's clients do + // not reconnect in a single thundering herd. The 30s hard timeout + // below backstops the window; keep the jitter well under it. + let closed = drain_conn_manager.drain_all_jittered(drain_jitter_ms); info!( connections = closed, - "Sent restart close frame to all live WebSocket connections" + jitter_ms = 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; diff --git a/crates/buzz-relay/src/state.rs b/crates/buzz-relay/src/state.rs index 58a869a995..1458f8b8a1 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -364,6 +364,54 @@ impl ConnectionManager { closed } + /// Staggered variant of [`Self::drain_all`]: closes every live connection + /// with the `1012 Service Restart` frame, but spreads the closes across + /// `[0, jitter_ms]` instead of firing them all in one instant. + /// + /// 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 `[0, jitter_ms]` + /// smears the reconnects across the window while keeping the well-attributed + /// 1012 close. + /// + /// The sticky drain flag is set **synchronously before returning**, so this + /// preserves [`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 + /// per-connection closes run in detached tasks; the caller's hard-drain + /// timeout is the backstop, so `jitter_ms` must stay well under it. + /// + /// `jitter_ms == 0` delegates to [`Self::drain_all`] for identical + /// synchronous behavior and no task spawns. + /// + /// Returns the number of connections scheduled to close. + pub fn drain_all_jittered(self: &Arc, jitter_ms: u64) -> usize { + if jitter_ms == 0 { + return self.drain_all(); + } + // Set the sticky flag first (store-then-iterate), so any registration + // racing past the snapshot below observes it and self-signals with no + // delay. Only the connections captured in this snapshot are jittered. + self.draining.store(true, Ordering::SeqCst); + let mut scheduled = 0usize; + for entry in self.connections.iter() { + let ctrl_tx = entry.ctrl_tx.clone(); + let cancel = entry.cancel.clone(); + // `rand::random % n` matches the jitter idiom used elsewhere in the + // relay (see main.rs cron jitter). Uniform over [0, jitter_ms). + let delay_ms = rand::random::() % jitter_ms; + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + let _ = ctrl_tx.try_send(Self::restart_close_frame()); + cancel.cancel(); + }); + scheduled += 1; + } + scheduled + } + /// The WS close frame announcing a graceful restart: 1012 Service Restart. fn restart_close_frame() -> WsMessage { WsMessage::Close(Some(axum::extract::ws::CloseFrame { @@ -1930,4 +1978,128 @@ mod tests { other => panic!("expected a restart close frame, got {other:?}"), } } + + #[tokio::test] + async fn drain_all_jittered_zero_is_synchronous_drain_all() { + // jitter_ms == 0 must reproduce drain_all exactly: synchronous close, + // no task spawns, frame already queued when the call 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, + 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_jittered(0); + + 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_jittered(0) returns" + ); + } + + #[tokio::test(start_paused = true)] + async fn drain_all_jittered_defers_close_until_within_jitter_window() { + // With jitter, the close must NOT be queued synchronously: it fires + // from a spawned task after a delay bounded by the jitter window. But + // 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, + 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; + let scheduled = mgr.drain_all_jittered(jitter_ms); + assert_eq!(scheduled, 1, "connection is scheduled to close"); + + // Let the spawned close task run up to its `sleep` await point without + // advancing the (paused) clock, so its timer is armed before we jump. + tokio::task::yield_now().await; + + // Not closed synchronously — the delayed task 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, + 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; the deferred close must fire. + tokio::time::advance(std::time::Duration::from_millis(jitter_ms + 1)).await; + tokio::task::yield_now().await; + + 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:?}"), + } + } } From 3930255ddb20e6bb0723abc39cd3be543b16ba5a Mon Sep 17 00:00:00 2001 From: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Date: Mon, 3 Aug 2026 18:27:26 -0400 Subject: [PATCH 2/4] fix(relay): await jittered websocket drain Co-authored-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Signed-off-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> --- crates/buzz-relay/src/config.rs | 42 +++++++++----- crates/buzz-relay/src/main.rs | 43 +++++++++------ crates/buzz-relay/src/state.rs | 98 ++++++++++++++++++--------------- 3 files changed, 110 insertions(+), 73 deletions(-) diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 244f20ce89..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 { @@ -64,14 +68,15 @@ pub struct Config { /// 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 `[0, drain_jitter_ms]`, which + /// 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 are capped - /// at shutdown to leave headroom under the 30s hard-drain timeout; keep it - /// well below `terminationGracePeriodSeconds`. + /// 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, @@ -466,15 +471,19 @@ impl Config { Err(_) => 0, }; - // Drain jitter: 0 = off (default). Non-negative parse, same shape as - // the replica-read budget above. Bound is enforced (capped) at - // shutdown, not here, so config never fails on a large value. + // 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(), - ) - })?, + 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, }; @@ -1304,6 +1313,9 @@ mod tests { 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; @@ -1317,7 +1329,11 @@ mod tests { } assert_eq!(unset, 0, "drain jitter must default off"); - assert_eq!(set, 20_000); + 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(), diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index 9d334d3750..32d26dbb9d 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, @@ -1209,7 +1211,7 @@ async fn serve( 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"); @@ -1217,26 +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. - // - // With BUZZ_DRAIN_JITTER_MS > 0, the closes are spread across the - // jitter window instead of firing all at once, so a pod's clients do - // not reconnect in a single thundering herd. The 30s hard timeout - // below backstops the window; keep the jitter well under it. - let closed = drain_conn_manager.drain_all_jittered(drain_jitter_ms); + // 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_jittered(drain_jitter_ms).await; info!( connections = closed, 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) @@ -1284,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(()); } @@ -1305,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 1458f8b8a1..5b2a308519 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; @@ -366,50 +367,58 @@ impl ConnectionManager { /// Staggered variant of [`Self::drain_all`]: closes every live connection /// with the `1012 Service Restart` frame, but spreads the closes across - /// `[0, jitter_ms]` instead of firing them all in one instant. + /// `[1, jitter_ms]` when jitter is enabled instead of firing them all in + /// one instant. /// /// 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 `[0, jitter_ms]` + /// close by an independent uniform random offset in `[1, jitter_ms]` /// smears the reconnects across the window while keeping the well-attributed /// 1012 close. /// - /// The sticky drain flag is set **synchronously before returning**, so this - /// preserves [`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 - /// per-connection closes run in detached tasks; the caller's hard-drain - /// timeout is the backstop, so `jitter_ms` must stay well under it. + /// 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. /// - /// `jitter_ms == 0` delegates to [`Self::drain_all`] for identical - /// synchronous behavior and no task spawns. + /// `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 scheduled to close. - pub fn drain_all_jittered(self: &Arc, jitter_ms: u64) -> usize { - if jitter_ms == 0 { - return self.drain_all(); - } - // Set the sticky flag first (store-then-iterate), so any registration - // racing past the snapshot below observes it and self-signals with no - // delay. Only the connections captured in this snapshot are jittered. + /// 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 mut scheduled = 0usize; - for entry in self.connections.iter() { - let ctrl_tx = entry.ctrl_tx.clone(); - let cancel = entry.cancel.clone(); - // `rand::random % n` matches the jitter idiom used elsewhere in the - // relay (see main.rs cron jitter). Uniform over [0, jitter_ms). - let delay_ms = rand::random::() % jitter_ms; - tokio::spawn(async move { - tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; - let _ = ctrl_tx.try_send(Self::restart_close_frame()); - cancel.cancel(); - }); - scheduled += 1; - } - scheduled + let frame = Self::restart_close_frame(); + let pending: Vec<_> = self + .connections + .iter() + .map(|entry| { + let ctrl_tx = entry.ctrl_tx.clone(); + let cancel = entry.cancel.clone(); + let frame = frame.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 _ = ctrl_tx.try_send(frame); + cancel.cancel(); + } + }) + .collect(); + let count = pending.len(); + join_all(pending).await; + count } /// The WS close frame announcing a graceful restart: 1012 Service Restart. @@ -1999,7 +2008,7 @@ mod tests { 3, ); - let closed = mgr.drain_all_jittered(0); + let closed = mgr.drain_all_jittered(0).await; assert_eq!(closed, 1); assert!(cancel.is_cancelled(), "zero jitter cancels synchronously"); @@ -2037,14 +2046,16 @@ mod tests { ); let jitter_ms = 20_000u64; - let scheduled = mgr.drain_all_jittered(jitter_ms); - assert_eq!(scheduled, 1, "connection is scheduled to close"); - - // Let the spawned close task run up to its `sleep` await point without - // advancing the (paused) clock, so its timer is armed before we jump. - tokio::task::yield_now().await; + // 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 synchronously — the delayed task is parked on its timer. + // Not closed yet — the delayed drain is parked on its timer. assert!( !cancel.is_cancelled(), "jittered close is deferred, not synchronous" @@ -2082,9 +2093,10 @@ mod tests { "late registration gets the restart close with no delay" ); - // Advance past the whole jitter window; the deferred close must fire. + // 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; - tokio::task::yield_now().await; + assert_eq!(drain.await, 1, "one captured connection drained"); assert!( cancel.is_cancelled(), From c47c1e5f7497f71599a9aaa8bbbdd89f4f799e72 Mon Sep 17 00:00:00 2001 From: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Date: Mon, 3 Aug 2026 19:15:19 -0400 Subject: [PATCH 3/4] refactor(relay): unify websocket drain API Co-authored-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Signed-off-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> --- crates/buzz-relay/src/main.rs | 2 +- crates/buzz-relay/src/state.rs | 64 ++++++++-------------------------- 2 files changed, 16 insertions(+), 50 deletions(-) diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index 32d26dbb9d..e26f5b3c6a 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -1230,7 +1230,7 @@ async fn serve( 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_jittered(drain_jitter_ms).await; + let closed = drain_conn_manager.drain_all(drain_jitter_ms).await; info!( connections = closed, jitter_ms = drain_jitter_ms, diff --git a/crates/buzz-relay/src/state.rs b/crates/buzz-relay/src/state.rs index 5b2a308519..13a45ca418 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -334,41 +334,8 @@ impl ConnectionManager { closed } - /// Closes every live connection with a `1012 Service Restart` close frame. - /// - /// 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. - /// - /// 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. - /// - /// 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. - 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 - } - - /// Staggered variant of [`Self::drain_all`]: closes every live connection - /// with the `1012 Service Restart` frame, but spreads the closes across - /// `[1, jitter_ms]` when jitter is enabled instead of firing them all in - /// one instant. + /// 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 @@ -389,7 +356,7 @@ impl ConnectionManager { /// first await, preserving the previous all-at-once behavior. /// /// Returns the number of connections signalled. - pub async fn drain_all_jittered(&self, jitter_ms: u64) -> usize { + 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. @@ -1882,7 +1849,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"); @@ -1928,7 +1895,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!( @@ -1953,7 +1920,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(); @@ -1989,9 +1956,9 @@ mod tests { } #[tokio::test] - async fn drain_all_jittered_zero_is_synchronous_drain_all() { - // jitter_ms == 0 must reproduce drain_all exactly: synchronous close, - // no task spawns, frame already queued when the call returns. + 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); @@ -2008,7 +1975,7 @@ mod tests { 3, ); - let closed = mgr.drain_all_jittered(0).await; + let closed = mgr.drain_all(0).await; assert_eq!(closed, 1); assert!(cancel.is_cancelled(), "zero jitter cancels synchronously"); @@ -2019,15 +1986,14 @@ mod tests { .expect("close frame delivered synchronously"), WsMessage::Close(Some(_)) ), - "the restart close is queued before drain_all_jittered(0) returns" + "the restart close is queued before drain_all(0) returns" ); } #[tokio::test(start_paused = true)] - async fn drain_all_jittered_defers_close_until_within_jitter_window() { - // With jitter, the close must NOT be queued synchronously: it fires - // from a spawned task after a delay bounded by the jitter window. But - // the sticky drain flag is still set immediately, so a late + 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(); @@ -2048,7 +2014,7 @@ mod tests { 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); + let drain = mgr.drain_all(jitter_ms); tokio::pin!(drain); assert!( futures_util::poll!(&mut drain).is_pending(), From 21990063e04429aa6fba2e8576b912b4799dabf3 Mon Sep 17 00:00:00 2001 From: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Date: Mon, 3 Aug 2026 19:55:55 -0400 Subject: [PATCH 4/4] fix(relay): acknowledge restart close flush Co-authored-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Signed-off-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> --- crates/buzz-relay/src/connection.rs | 96 +++++++++++++++-- crates/buzz-relay/src/handlers/event.rs | 3 + crates/buzz-relay/src/state.rs | 131 ++++++++++++++++++++++-- 3 files changed, 214 insertions(+), 16 deletions(-) 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/state.rs b/crates/buzz-relay/src/state.rs index 13a45ca418..b581ec4766 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -32,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; @@ -46,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. @@ -203,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( @@ -221,6 +227,7 @@ impl ConnectionManager { ConnEntry { tx, ctrl_tx, + restart_tx, cancel, community_id, backpressure_count, @@ -234,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(); + } } } @@ -361,14 +375,13 @@ impl ConnectionManager { // 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 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 frame = frame.clone(); let delay_ms = if jitter_ms == 0 { 0 } else { @@ -378,8 +391,26 @@ impl ConnectionManager { if delay_ms > 0 { tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; } - let _ = ctrl_tx.try_send(frame); - cancel.cancel(); + 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(); @@ -1270,6 +1301,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::clone(&bp), @@ -1395,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), @@ -1441,6 +1474,7 @@ mod tests { conn_a, tx_a, ctrl_tx_a, + None, CancellationToken::new(), community_a, Arc::new(AtomicU8::new(0)), @@ -1451,6 +1485,7 @@ mod tests { conn_b, tx_b, ctrl_tx_b, + None, CancellationToken::new(), community_b, Arc::new(AtomicU8::new(0)), @@ -1487,6 +1522,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel, buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), bp, @@ -1789,6 +1825,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1817,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 @@ -1833,6 +1942,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), community, Arc::new(AtomicU8::new(0)), @@ -1884,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)), @@ -1931,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)), @@ -1968,6 +2080,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -2004,6 +2117,7 @@ mod tests { conn_id, tx, ctrl_tx, + None, cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)), @@ -2041,6 +2155,7 @@ mod tests { late_id, late_tx, late_ctrl_tx, + None, late_cancel.clone(), buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), Arc::new(AtomicU8::new(0)),