mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
relay: fuzz WebSocket 1012 restart-close timing on graceful drain (BUZZ_DRAIN_JITTER_MS) (#4542)
## Problem On SIGTERM the relay sends every live WebSocket a **1012 Service Restart** close frame via `ConnectionManager::drain_all()` — all in the same instant (`main.rs` shutdown task → `state.rs::drain_all`). On a pod holding thousands of sessions, that makes every client reconnect simultaneously: the thundering-herd reconnect behind the DB pool-timeout bursts observed on each rolling deploy. Client-side jitter can't fix this — the desktop client *resets* its backoff to base on a 1012 and reconnects with only ±25% jitter (`relayClientSession.ts`), so the spread has to come from the server. ## Change Add `BUZZ_DRAIN_JITTER_MS` (default `0` = unchanged behavior). The two paths are kept **deliberately separate** so the default is byte-for-byte the previously shipped shutdown: - **Jitter off (`0`/unset, the default):** the original synchronous, all-at-once `drain_all()` runs unchanged — queue the 1012 on each connection's control channel, cancel, return. No new machinery on the default path. - **Jitter on (`> 0`):** a separate async `drain_all_jittered(jitter_ms)` spreads each connection's restart close over an independent uniform delay in **`[1, jitter_ms]`**. Each delayed close travels a dedicated `RestartClose` channel; the writer flushes the 1012 frame and **acknowledges the flush over a oneshot**, so drain waits for confirmed delivery (up to `RESTART_CLOSE_ACK_TIMEOUT` = 5s) rather than assuming it, falling back to cancellation if the channel is full/closed or the ack times out. The drain future is **owned and awaited** by the shutdown task, and the 30s hard-drain backstop is aborted only after a clean drain — so a clean roll exits `0`. The two methods can be unified and the old one dropped later once the jittered path is proven for all cases. - **`config.rs`** — `drain_jitter_ms`: non-negative parse, clamped to `MAX_DRAIN_JITTER_MS` = **20s** (leaving 10s of the 30s budget for flush). Junk fails loudly at startup; **empty/whitespace-only is treated as unset (jitter off)** so a `BUZZ_DRAIN_JITTER_MS=""` kill switch does not crashloop the relay (matches the sibling env vars in this file). - **`state.rs`** — `drain_all()` (unchanged synchronous default) + `drain_all_jittered()` (jittered + flush-ack). Both set the sticky `draining` flag before the first await. A registration that lands mid-shutdown always self-signals via the **immediate** control-frame + cancel path — jitter smears already-established sockets, not late arrivals. - **`main.rs`** — shutdown task dispatches: `drain_jitter_ms == 0` → `drain_all()`, else `drain_all_jittered(...).await`. ## Safety - **Default off is the currently-committed path.** With jitter unset/0 the shutdown runs the original synchronous `drain_all()` — no restart channel, no ack wait. Safe to deploy dark and dial up. - **Shutdown-boundary race preserved.** Sticky flag set before any await; a late registration self-signals its close with no jitter. - **Owned + backstopped.** The jittered drain future is awaited; the 30s hard-drain `process::exit(1)` remains the ceiling. `MAX_DRAIN_JITTER_MS` (20s) + `RESTART_CLOSE_ACK_TIMEOUT` (5s) = 25s, inside the 30s budget; 5s pre-sleep + 25s = 30s against `terminationGracePeriodSeconds: 60`. ## Known behavior to note (not a blocker, flagged from review) On a **successful** flush the jittered path deliberately does not cancel the connection token — teardown then depends on the client echoing our Close, or on process exit. Compliant clients echo; a silent client rides to the 30s hard exit. The default (jitter-off) path cancels deterministically as before. ## Tests - `config::tests::drain_jitter_defaults_off_and_rejects_junk` — default off, `20000`, clamp `60000`→`20000`, explicit `0`, junk `"soon"` fails, **empty `""` and whitespace-only treated as off**. - `state::tests::drain_all_is_immediate` — default path queues frame + cancels synchronously. - `state::tests::drain_all_sends_restart_close_and_cancels_every_conn`, `drain_all_full_control_buffer_still_cancels`, `register_after_drain_self_signals_restart_close_and_cancel`. - `state::tests::drain_all_jittered_defers_close_until_within_jitter_window` (paused time). - `state::tests::drain_all_jittered_waits_for_writer_acknowledgement_without_cancelling`. - `state::tests::drain_all_jittered_cancels_when_restart_channel_is_full_or_closed`. - `state::tests::drain_all_jittered_cancels_when_flush_ack_times_out` (paused time — the 5s ack-timeout fallback). Validation at `46c690940`: `cargo fmt -p buzz-relay --check`, `cargo clippy -p buzz-relay --all-targets -- -D warnings`, and the drain/config unit suite all clean. Local live SIGTERM test with a real relay process + 200 NIP-42-authenticated sockets — see the PR comment for the before/after distribution and exit codes. ## Rollout Ship with default `0`, then set `BUZZ_DRAIN_JITTER_MS` (e.g. 10000–20000) on bb-block first, watch the roll-window pool-timeout metric, then bb-public. `""` is a safe kill switch. Complements the preStop `sleep` (stops routing before close). --------- Signed-off-by: npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w <80ffe7a6e22722984fb7fbf4c72b026030a745b08fae413253b1d993a5f19a30@buzz.block.builderlab.xyz> Signed-off-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz> Signed-off-by: Brad Seiler <seiler@squareup.com> Co-authored-by: npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w <80ffe7a6e22722984fb7fbf4c72b026030a745b08fae413253b1d993a5f19a30@buzz.block.builderlab.xyz> Co-authored-by: npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch <51cde9442e46eac81c83ec71552708c3b0cb96a88ff8e1581d8fb4de281afbf3@buzz.block.builderlab.xyz>
This commit is contained in:
Brad Seiler
co-authored by
npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w
npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch
npub1srl70fhzyu3fsnahl06vw2czvqc2w3ds37hyzvjnk8ve8f03ngcqg9le2w
npub128x7j3pwgm4vs8yra3c42fcgcwcvh94g3luwzkqa376du2q6l0esqcrwch
parent
005fe54d02
commit
e14fff74d0
@@ -278,6 +278,7 @@ out of the box with `just setup` or `just relay`. Common overrides:
|
||||
| `BUZZ_REQUIRE_AUTH_TOKEN` | `false` | When true, REST requires NIP-98 (no `X-Pubkey` fallback) |
|
||||
| `BUZZ_REQUIRE_RELAY_MEMBERSHIP` | `false` | When true, only pubkeys in `relay_members` can connect |
|
||||
| `BUZZ_REQUIRE_MEDIA_GET_AUTH` | `false` | When true, `GET`/`HEAD /media/*` require Blossom kind 24242 `t=get` auth plus relay membership. |
|
||||
| `BUZZ_DRAIN_JITTER_MS` | `0` (off) | Per-connection upper bound, in ms, for the random delay before each live WebSocket gets its `1012 Service Restart` close on graceful shutdown. `0` closes every socket at once (the previous behavior). A positive value spreads closes uniformly over `[1, value]` ms to avoid a reconnect thundering herd on rolling deploys. Values above `20000` are capped to `20000` (`MAX_DRAIN_JITTER_MS`) to leave close-frame delivery headroom under the relay's 30s hard-drain timeout. Empty or whitespace-only is treated as unset (off); a non-integer fails startup loudly. |
|
||||
| `BUZZ_AUDIT_ENABLED` | `true` | Tamper-evident event/media audit log. Set `false`/`0`/`off` to skip its DB pool and writes. Does not disable the separate moderation audit trail. |
|
||||
| `BUZZ_AUTO_MIGRATE` | `false` | Opt in with `true`/`1`/`yes`/`on` to run embedded SQLx migrations on relay startup |
|
||||
| `RELAY_OWNER_PUBKEY` | unset | Bootstrapped as `owner` in `relay_members` at first start |
|
||||
|
||||
@@ -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::<u64>()
|
||||
.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();
|
||||
|
||||
@@ -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<Mutex<HashMap<String, Vec<Filter>>>>;
|
||||
|
||||
/// Request for the writer to flush a restart close and report the result.
|
||||
pub(crate) struct RestartClose {
|
||||
pub(crate) flushed: tokio::sync::oneshot::Sender<bool>,
|
||||
}
|
||||
|
||||
/// 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::<WsMessage>(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::<RestartClose>(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<WebSocket, WsMessage>,
|
||||
data_rx: mpsc::Receiver<WsMessage>,
|
||||
ctrl_rx: mpsc::Receiver<WsMessage>,
|
||||
restart_rx: mpsc::Receiver<RestartClose>,
|
||||
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<S>(
|
||||
mut ws_send: S,
|
||||
mut data_rx: mpsc::Receiver<WsMessage>,
|
||||
mut ctrl_rx: mpsc::Receiver<WsMessage>,
|
||||
mut restart_rx: mpsc::Receiver<RestartClose>,
|
||||
cancel: CancellationToken,
|
||||
) where
|
||||
S: Sink<WsMessage> + Unpin,
|
||||
@@ -319,9 +332,21 @@ async fn send_loop_inner<S>(
|
||||
}
|
||||
|
||||
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!(
|
||||
|
||||
@@ -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)),
|
||||
|
||||
@@ -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,37 @@ async fn run_periodic_until_cancelled<Tick, TickFuture>(
|
||||
/// │ → graceful drain (30s) → exit │
|
||||
/// └─────────────────────────────────────────────────────────┘
|
||||
/// ```
|
||||
///
|
||||
/// ## Shutdown budget
|
||||
///
|
||||
/// The full teardown, measured from SIGTERM, is bounded as follows:
|
||||
///
|
||||
/// 1. `5s` grace. Readiness returns 503 immediately, then the process
|
||||
/// sleeps 5 seconds so Kubernetes stops routing new traffic before any
|
||||
/// listener closes.
|
||||
/// 2. `GRACEFUL_DRAIN_TIMEOUT` (`30s`) hard drain. Started at the end of the
|
||||
/// grace, this backstops the whole drain and force-exits the process if
|
||||
/// exceeded. It bounds everything after the grace, not the grace itself.
|
||||
///
|
||||
/// A single WebSocket can therefore stay open, from SIGTERM, for up to:
|
||||
///
|
||||
/// ```text
|
||||
/// 5s grace + up to 20s jitter + up to 5s close-frame ack = 30s
|
||||
/// (fixed) (MAX_DRAIN_JITTER_MS) (RESTART_CLOSE_ACK_TIMEOUT)
|
||||
/// ```
|
||||
///
|
||||
/// The 5s grace runs before the 30s hard-drain clock starts, so the jitter
|
||||
/// (capped at [`buzz_relay::config::MAX_DRAIN_JITTER_MS`] = 20s) plus the
|
||||
/// per-connection close-frame ack wait (`RESTART_CLOSE_ACK_TIMEOUT` = 5s in
|
||||
/// `state.rs`) sum to 25s and stay inside the 30s hard drain. Total worst
|
||||
/// case from SIGTERM to forced exit is 5s + 30s = 35s. Both fit inside the
|
||||
/// chart's `terminationGracePeriodSeconds: 60` (`deploy/charts/buzz/values.yaml`),
|
||||
/// which leaves headroom but assumes no `preStop` hook adds further delay.
|
||||
/// With jitter off (`BUZZ_DRAIN_JITTER_MS=0`, the default) sockets close
|
||||
/// all-at-once right after the grace, so the per-socket delay collapses to
|
||||
/// roughly the 5s grace plus the ack wait.
|
||||
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 +1238,36 @@ 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 {
|
||||
// TODO(coverage): `serve`'s shutdown wiring has no automated test. The
|
||||
// jittered drain helper (`ConnectionManager::drain_all_jittered`) is
|
||||
// covered in `state.rs`, but coverage of the helper is not coverage of
|
||||
// its use here: the three wiring facts below are currently unguarded, and
|
||||
// mutating any one of them leaves the suite green.
|
||||
// 1. Jitter dispatch: `drain_jitter_ms == 0` must pick `drain_all`, and
|
||||
// a non-zero value must pick `drain_all_jittered(drain_jitter_ms)`.
|
||||
// A mutant that inverts this condition ships jitter-off in prod.
|
||||
// 2. The shutdown handle must be awaited before the abort. Dropping the
|
||||
// `shutdown_handle.await` (both the UDS and TCP-only return paths) is
|
||||
// the exact shape of the previously shipped detached-timer bug,
|
||||
// relocated from the helper to the call site: the runtime can exit
|
||||
// before delayed closes flush, so no client sees a 1012.
|
||||
// 3. `shutdown_tx.send(true)` must reach every listener's
|
||||
// `with_graceful_shutdown` future, on both the UDS and TCP-only paths.
|
||||
//
|
||||
// A focused test would refactor the drain/dispatch decision and the
|
||||
// listener-shutdown fan-out into a small seam that does not need a bound
|
||||
// socket or a real SIGTERM. One shape: extract the body of this spawned
|
||||
// task into a `run_graceful_shutdown(state, shutdown_tx)` fn parameterised
|
||||
// over a signal future and a clock, inject a fake `ConnectionManager`
|
||||
// (or a trait over `drain_all` / `drain_all_jittered`) that records which
|
||||
// path ran, drive it with `tokio::time` paused, and assert: (a) the right
|
||||
// drain path ran for jitter 0 vs non-zero, (b) the drain future completed
|
||||
// before the abort fired, and (c) each subscribed `watch` receiver
|
||||
// observed `true`. This keeps the test off real ports and off wall-clock
|
||||
// sleeps. Not implemented here. This comment records the plan only.
|
||||
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 +1275,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 +1347,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 +1372,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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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<ScopedPubkeyKey, SlidingWindowCounter>;
|
||||
|
||||
@@ -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<WsMessage>,
|
||||
restart_tx: Option<mpsc::Sender<RestartClose>>,
|
||||
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<WsMessage>,
|
||||
ctrl_tx: mpsc::Sender<WsMessage>,
|
||||
restart_tx: Option<mpsc::Sender<RestartClose>>,
|
||||
cancel: CancellationToken,
|
||||
community_id: CommunityId,
|
||||
backpressure_count: Arc<AtomicU8>,
|
||||
@@ -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::<u64>() % 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:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -128,6 +128,7 @@ spec:
|
||||
- { name: BUZZ_MAX_CONNECTIONS, value: {{ .Values.relay.maxConnections | quote }} }
|
||||
- { name: BUZZ_MAX_CONCURRENT_HANDLERS, value: {{ .Values.relay.maxConcurrentHandlers | quote }} }
|
||||
- { name: BUZZ_SEND_BUFFER, value: {{ .Values.relay.sendBuffer | quote }} }
|
||||
- { name: BUZZ_DRAIN_JITTER_MS, value: {{ .Values.relay.drainJitterMs | quote }} }
|
||||
- { name: BUZZ_REQUIRE_AUTH_TOKEN, value: {{ .Values.relay.requireAuthToken | quote }} }
|
||||
- { name: BUZZ_REQUIRE_RELAY_MEMBERSHIP, value: {{ .Values.relay.requireRelayMembership | quote }} }
|
||||
- { name: BUZZ_REQUIRE_MEDIA_GET_AUTH, value: {{ .Values.relay.requireMediaGetAuth | quote }} }
|
||||
|
||||
@@ -59,6 +59,7 @@
|
||||
"maxConnections": { "type": "integer", "minimum": 1 },
|
||||
"maxConcurrentHandlers": { "type": "integer", "minimum": 1 },
|
||||
"sendBuffer": { "type": "integer", "minimum": 1 },
|
||||
"drainJitterMs": { "type": "integer", "minimum": 0 },
|
||||
"requireAuthToken": { "type": "boolean" },
|
||||
"requireRelayMembership": { "type": "boolean" },
|
||||
"requireMediaGetAuth": { "type": "boolean" },
|
||||
|
||||
@@ -105,6 +105,16 @@ relay:
|
||||
maxConnections: 10000
|
||||
maxConcurrentHandlers: 1024
|
||||
sendBuffer: 1000
|
||||
# Graceful-shutdown reconnect jitter. On SIGTERM the relay closes every live
|
||||
# WebSocket with a 1012 Service Restart frame; with a rolling deploy this can
|
||||
# release a whole pod's sockets at once and stampede reconnects into the DB
|
||||
# pool. A positive value (milliseconds) spreads each close over a per-socket
|
||||
# random delay in [1, drainJitterMs], smoothing the reconnect herd. 0 (the
|
||||
# default) closes all sockets at once, preserving the previous behavior.
|
||||
# Values above 20000 are capped to 20000, leaving close-frame delivery
|
||||
# headroom under the relay's 30s hard-drain timeout (itself inside the 60s
|
||||
# terminationGracePeriodSeconds below).
|
||||
drainJitterMs: 0
|
||||
requireAuthToken: true
|
||||
requireRelayMembership: true
|
||||
# Authenticated media reads: relay GET/HEAD /media/* requires Blossom
|
||||
|
||||
Reference in New Issue
Block a user