diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index be9794922..22101219e 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -1127,6 +1127,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 tx = shutdown_tx.clone(); tokio::spawn(async move { shutdown_signal().await; @@ -1136,6 +1137,16 @@ 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(); + info!( + connections = closed, + "Sent restart close frame 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"); diff --git a/crates/buzz-relay/src/router.rs b/crates/buzz-relay/src/router.rs index 2375a4f9c..2af036079 100644 --- a/crates/buzz-relay/src/router.rs +++ b/crates/buzz-relay/src/router.rs @@ -299,9 +299,20 @@ async fn nip11_or_ws_handler( let max_frame_bytes = state.config.max_frame_bytes; match WebSocketUpgrade::from_request(req, &state).await { - Ok(ws) => limit_relay_websocket(ws, max_frame_bytes) - .on_upgrade(move |socket| handle_connection(socket, state, addr, tenant)) - .into_response(), + Ok(ws) => { + // Shutting down: refuse new sockets instead of accepting a + // connection onto a dying pod. Readiness already returns 503, but + // that only stops K8s routing — direct and in-flight upgrades + // still reach here during the pre-drain grace window. Clients + // treat the refusal as a normal dial failure and retry, landing + // on a healthy pod. + if state.shutting_down.load(Ordering::Relaxed) { + return (StatusCode::SERVICE_UNAVAILABLE, "relay restarting").into_response(); + } + limit_relay_websocket(ws, max_frame_bytes) + .on_upgrade(move |socket| handle_connection(socket, state, addr, tenant)) + .into_response() + } Err(_) => { // Browser requesting HTML and Git web GUI is enabled → serve SPA. if state.config.serve_git_web_gui { diff --git a/crates/buzz-relay/src/state.rs b/crates/buzz-relay/src/state.rs index 3a6ca4928..758c001b9 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -181,6 +181,10 @@ where /// Tracks active Nostr WebSocket connections and provides message routing by connection ID. pub struct ConnectionManager { connections: DashMap, + /// Sticky drain flag set by [`Self::drain_all`]. Registrations that land + /// after the drain snapshot self-signal, so no upgrade-vs-shutdown + /// interleaving can produce a connection that misses the restart close. + draining: AtomicBool, } impl ConnectionManager { @@ -188,6 +192,7 @@ impl ConnectionManager { pub fn new() -> Self { Self { connections: DashMap::new(), + draining: AtomicBool::new(false), } } @@ -208,6 +213,8 @@ impl ConnectionManager { subscriptions: ConnectionSubscriptions, grace_limit: u8, ) { + let drain_ctrl_tx = ctrl_tx.clone(); + let drain_cancel = cancel.clone(); self.connections.insert( conn_id, ConnEntry { @@ -221,6 +228,14 @@ impl ConnectionManager { grace_limit, }, ); + // 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. + if self.draining.load(Ordering::SeqCst) { + let _ = drain_ctrl_tx.try_send(Self::restart_close_frame()); + drain_cancel.cancel(); + } } /// Removes a connection from the registry. @@ -318,6 +333,45 @@ 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 + } + + /// The WS close frame announcing a graceful restart: 1012 Service Restart. + fn restart_close_frame() -> WsMessage { + WsMessage::Close(Some(axum::extract::ws::CloseFrame { + code: axum::extract::ws::close_code::RESTART, + reason: axum::extract::ws::Utf8Bytes::from_static("relay restarting"), + })) + } + /// Return the server-resolved community that the connection's host bound to. pub fn community_for_conn(&self, conn_id: Uuid) -> Option { self.connections @@ -1737,4 +1791,142 @@ mod tests { "community-B session stays live — ban does not cross the tenant fence" ); } + + #[tokio::test] + async fn drain_all_sends_restart_close_and_cancels_every_conn() { + // Graceful shutdown must tell every live client to reconnect — across + // all communities — with a 1012 restart close frame queued ahead of + // the cancel-driven socket close. + let mgr = ConnectionManager::new(); + + let register = |community| { + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, ctrl_rx) = mpsc::channel(8); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx, + cancel.clone(), + community, + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + (ctrl_rx, cancel) + }; + + let (mut ctrl_a, cancel_a) = register(buzz_core::tenant::CommunityId::from_uuid( + Uuid::from_u128(0xa), + )); + let (mut ctrl_b, cancel_b) = register(buzz_core::tenant::CommunityId::from_uuid( + Uuid::from_u128(0xb), + )); + + let closed = mgr.drain_all(); + + assert_eq!(closed, 2, "every connection is signalled, no tenant fence"); + assert!(cancel_a.is_cancelled(), "community-A session is cancelled"); + assert!(cancel_b.is_cancelled(), "community-B session is cancelled"); + + for ctrl_rx in [&mut ctrl_a, &mut ctrl_b] { + let frame = ctrl_rx.try_recv().expect("close frame delivered"); + match frame { + WsMessage::Close(Some(close)) => { + assert_eq!( + close.code, + axum::extract::ws::close_code::RESTART, + "close code is 1012 Service Restart" + ); + assert_eq!(close.reason.as_str(), "relay restarting"); + } + other => panic!("expected a restart close frame, got {other:?}"), + } + } + } + + #[tokio::test] + async fn drain_all_full_control_buffer_still_cancels() { + // Best-effort delivery: a wedged control channel must not block the + // drain — the cancel still closes the socket, just without the frame. + let mgr = ConnectionManager::new(); + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(8); + let (ctrl_tx, mut ctrl_rx) = mpsc::channel(1); + let cancel = CancellationToken::new(); + mgr.register( + conn_id, + tx, + ctrl_tx.clone(), + cancel.clone(), + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + Arc::new(AtomicU8::new(0)), + Arc::new(Mutex::new(HashMap::new())), + 3, + ); + // Wedge the 1-slot control channel. + ctrl_tx + .try_send(WsMessage::Text("wedge".into())) + .expect("fill control channel"); + + let closed = mgr.drain_all(); + + assert_eq!(closed, 1); + assert!( + cancel.is_cancelled(), + "cancel fires even when the close frame cannot be queued" + ); + // Only the wedge frame is present — the close was dropped, not queued. + assert!(matches!( + ctrl_rx.try_recv().expect("wedge frame"), + WsMessage::Text(_) + )); + assert!(ctrl_rx.try_recv().is_err(), "no second frame queued"); + } + + #[tokio::test] + async fn register_after_drain_self_signals_restart_close_and_cancel() { + // The shutdown-boundary race: an upgrade accepted before SIGTERM can + // finish its async admission check and register AFTER drain_all's + // one-shot snapshot. The sticky drain flag makes that interleaving + // deterministic — register itself queues the 1012 and cancels, so no + // late registration can ride out graceful shutdown unclosed. + let mgr = ConnectionManager::new(); + + // Drain with zero connections — sets the sticky flag. + assert_eq!(mgr.drain_all(), 0); + + // Late registration lands after the snapshot. + 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, + ); + + assert!( + cancel.is_cancelled(), + "late registration is cancelled by the sticky drain flag" + ); + match ctrl_rx.try_recv().expect("close frame delivered") { + WsMessage::Close(Some(close)) => { + assert_eq!( + close.code, + axum::extract::ws::close_code::RESTART, + "late registration still gets the 1012 restart close" + ); + assert_eq!(close.reason.as_str(), "relay restarting"); + } + other => panic!("expected a restart close frame, got {other:?}"), + } + } } diff --git a/desktop/playwright.config.ts b/desktop/playwright.config.ts index be9d73ed0..0469497d8 100644 --- a/desktop/playwright.config.ts +++ b/desktop/playwright.config.ts @@ -137,6 +137,7 @@ export default defineConfig({ "**/persona-sync.spec.ts", "**/team-snapshot.spec.ts", "**/agents-everywhere.live.spec.ts", + "**/relay-restart.live.spec.ts", "**/parity-ancestor-island.spec.ts", ], use: { diff --git a/desktop/src/shared/api/relayClientSession.ts b/desktop/src/shared/api/relayClientSession.ts index 6dcc85715..a5fd8a0eb 100644 --- a/desktop/src/shared/api/relayClientSession.ts +++ b/desktop/src/shared/api/relayClientSession.ts @@ -534,8 +534,6 @@ export class RelayClient { } private async connect() { - // Clear any pending stability timer from a previous connection — a new - // connect attempt resets the clock and must re-arm the timer on success. if (this.stabilityTimer !== null) { window.clearTimeout(this.stabilityTimer); this.stabilityTimer = null; @@ -545,51 +543,66 @@ export class RelayClient { this.hasConnectedOnce ? "reconnecting" : "connecting", ); - if (!this.relayUrl) { - this.relayUrl = await getRelayWsUrl(); - } - const generation = ++this.connectionGeneration; this.onMessageChannel = new Channel((message) => { - void this.handleWsMessage(message, generation); - }); - - this.wsId = await invoke("plugin:websocket|connect", { - url: this.relayUrl, - onMessage: this.onMessageChannel, - config: {}, - }); - - await new Promise((resolve, reject) => { - const timeout = window.setTimeout(() => { - this.authRequest = null; + void this.handleWsMessage(message, generation).catch((error) => { + if (generation !== this.connectionGeneration) return; this.resetConnection( - new Error("Timed out while waiting for relay authentication."), + this.normalizeRelayError(error, "Relay connection errored."), ); - reject(new Error("Timed out while waiting for relay authentication.")); - }, AUTH_TIMEOUT_MS); - - this.authRequest = { - pendingEventId: "", - resolve, - reject, - timeout, - }; + }); }); - // Start a stability timer instead of resetting backoff immediately. - // The backoff resets to its base value only after BACKOFF_RESET_STABLE_MS - // of uninterrupted uptime, preventing fast reconnect loops from erasing - // the exponential backoff that throttles them. - this.stabilityTimer = window.setTimeout(() => { - this.stabilityTimer = null; - this.reconnectDelayMs = RECONNECT_BASE_DELAY_MS; - }, BACKOFF_RESET_STABLE_MS); + try { + if (!this.relayUrl) { + this.relayUrl = await getRelayWsUrl(); + } + const wsId = await invoke("plugin:websocket|connect", { + url: this.relayUrl, + onMessage: this.onMessageChannel, + config: {}, + }); + if (generation !== this.connectionGeneration) { + void closeWebSocket(wsId, "stale connection attempt"); + throw new Error("Relay connection attempt was superseded."); + } + this.wsId = wsId; - await this.replayLiveSubscriptions(); - this.connectionStateEmitter.set("connected"); - this.stallWatchdog.start(); - this.emitReconnectIfNeeded(); + await new Promise((resolve, reject) => { + const timeout = window.setTimeout(() => { + const error = new Error("Relay authentication timed out."); + this.authRequest = null; + this.resetConnection(error); + reject(error); + }, AUTH_TIMEOUT_MS); + + this.authRequest = { + pendingEventId: "", + resolve, + reject, + timeout, + }; + }); + + this.stabilityTimer = window.setTimeout(() => { + this.stabilityTimer = null; + this.reconnectDelayMs = RECONNECT_BASE_DELAY_MS; + }, BACKOFF_RESET_STABLE_MS); + + await this.replayLiveSubscriptions(); + this.connectionStateEmitter.set("connected"); + this.stallWatchdog.start(); + this.emitReconnectIfNeeded(); + } catch (error) { + const connectionError = this.normalizeRelayError( + error, + "Failed to connect to relay.", + ); + if (generation === this.connectionGeneration) { + this.resetConnection(connectionError); + } + throw connectionError; + } } private async subscribe( diff --git a/desktop/src/testing/e2eBridge.ts b/desktop/src/testing/e2eBridge.ts index c89cdbfd5..ffcefee8a 100644 --- a/desktop/src/testing/e2eBridge.ts +++ b/desktop/src/testing/e2eBridge.ts @@ -250,6 +250,8 @@ type E2eConfig = { openerError?: string; /** Delay binding signatures so specs can exercise request supersession. */ nostrBindSignDelayMs?: number; + /** Reject successive mock WebSocket connect attempts, then resume. */ + websocketConnectErrors?: string[]; stallWebsocketSends?: boolean; userSearchDelayMs?: number; // NIP-IA gate inputs — see tests/helpers/bridge.ts:MockBridgeOptions for @@ -8496,6 +8498,11 @@ async function connectRealSocket(args: { url?: string; onMessage: unknown }) { } async function connectMockSocket(args: { onMessage: unknown }) { + const connectError = getConfig()?.mock?.websocketConnectErrors?.shift(); + if (connectError) { + throw new Error(connectError); + } + if (mockWebsocketSendMutexWedged) { return new Promise(() => {}); } diff --git a/desktop/tests/e2e/helpers/twoRelayHarness.ts b/desktop/tests/e2e/helpers/twoRelayHarness.ts index 2eded2d9b..98acd8639 100644 --- a/desktop/tests/e2e/helpers/twoRelayHarness.ts +++ b/desktop/tests/e2e/helpers/twoRelayHarness.ts @@ -17,7 +17,7 @@ type OwnedProcess = { name: string; child: ChildProcess; logPath: string }; export class TwoRelayHarness { readonly root: string; - readonly relays: readonly [RelaySpec, RelaySpec]; + readonly relays: readonly RelaySpec[]; private readonly processes: OwnedProcess[] = []; get ownedPids(): number[] { @@ -26,12 +26,12 @@ export class TwoRelayHarness { ); } - private constructor(root: string, relays: readonly [RelaySpec, RelaySpec]) { + private constructor(root: string, relays: readonly RelaySpec[]) { this.root = root; this.relays = relays; } - static async create(relays: readonly [RelaySpec, RelaySpec]) { + static async create(relays: readonly RelaySpec[]) { return new TwoRelayHarness( await mkdtemp(join(tmpdir(), "buzz-ae-e2e-")), relays, @@ -46,6 +46,45 @@ export class TwoRelayHarness { ); } + /** + * SIGTERM one relay by name and wait for it to exit on its own. Unlike + * `stop()`, this never escalates to SIGKILL: the relay's graceful drain + * (readiness 503 → 5s grace → 1012 close broadcast → listener drain) is + * exactly what restart tests exercise, so a forced kill would invalidate + * the run. Throws if the process does not exit within `timeoutMs`. + */ + async terminateRelayGracefully(name: string, timeoutMs = 45_000) { + const owned = [...this.processes] + .reverse() + .find( + (candidate) => + candidate.name === name && + candidate.child.exitCode === null && + candidate.child.signalCode === null, + ); + if (!owned) throw new Error(`no live relay process named ${name}`); + const exited = new Promise((resolveExit) => + owned.child.once("exit", () => resolveExit()), + ); + this.signal(owned.child, "SIGTERM"); + if ( + !(await Promise.race([exited.then(() => true), delay(timeoutMs, false)])) + ) { + throw new Error( + `${name} did not exit within ${timeoutMs}ms of SIGTERM — graceful drain is stuck`, + ); + } + } + + /** Start a fresh process for a relay spec on its original ports. */ + async restartRelay(name: string, binary = process.env.BUZZ_E2E_RELAY_BIN) { + if (!binary) + throw new Error("BUZZ_E2E_RELAY_BIN is required for the live gate"); + const relay = this.relays.find((candidate) => candidate.name === name); + if (!relay) throw new Error(`no relay spec named ${name}`); + await this.startRelay(binary, relay); + } + async startAcp( name: string, relayWsUrl: string, diff --git a/desktop/tests/e2e/relay-reconnect.spec.ts b/desktop/tests/e2e/relay-reconnect.spec.ts index 2df9075d6..4f47ceded 100644 --- a/desktop/tests/e2e/relay-reconnect.spec.ts +++ b/desktop/tests/e2e/relay-reconnect.spec.ts @@ -80,6 +80,33 @@ test.beforeEach(async ({ page }) => { await installMockBridge(page); }); +test("failed initial relay dial retries automatically", async ({ page }) => { + await installMockBridge(page, { + websocketConnectErrors: ["mock relay pod unavailable"], + }); + await page.goto("/"); + + // App-shell preconnect owns a keep-alive request. The first native dial is + // rejected before a socket ID exists; the session must still enter its + // backoff loop and recover without a click, query, or reload. + await expect + .poll( + () => + page.evaluate(() => { + const getState = ( + window as Window & { + __BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?: () => string; + } + ).__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__; + if (!getState) throw new Error("Relay state seam is not installed."); + return getState(); + }), + { timeout: 10_000 }, + ) + .toBe("connected"); + await expect(page.getByTestId("channel-general")).toBeVisible(); +}); + test("passive relay watchdog does not write while the websocket is half-open", async ({ page, }) => { diff --git a/desktop/tests/e2e/relay-restart.live.spec.ts b/desktop/tests/e2e/relay-restart.live.spec.ts new file mode 100644 index 000000000..057ef46f5 --- /dev/null +++ b/desktop/tests/e2e/relay-restart.live.spec.ts @@ -0,0 +1,96 @@ +import { expect, test, type Page } from "@playwright/test"; + +import { installBridge } from "../helpers/bridge"; +import { TwoRelayHarness, type RelaySpec } from "./helpers/twoRelayHarness"; + +// Live gate: boots a REAL buzz-relay process, points the app at it, SIGTERMs +// the relay mid-session, restarts it on the same port, and asserts the client +// converges back to "connected". This proves the full restart story end to +// end: the relay's graceful-drain 1012 close broadcast (server side) and the +// client's dial-failure retry + 1012 fast-reconnect (desktop side) — the two +// halves that synthetic mock-websocket specs cannot compose. +// +// Requires: BUZZ_E2E_RELAY_RESTART=1, BUZZ_E2E_RELAY_BIN, and +// BUZZ_E2E_DATABASE_URL (plus reachable Redis and media object store, same +// infra as the agents-everywhere live gate). +const enabled = process.env.BUZZ_E2E_RELAY_RESTART === "1"; + +function required(name: string, value: string | undefined): string { + if (!value) throw new Error(`${name} is required for the live gate`); + return value; +} + +async function connectionState(page: Page): Promise { + return page.evaluate(() => { + const win = window as Window & { + __BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?: () => string; + }; + return win.__BUZZ_E2E_GET_RELAY_CONNECTION_STATE__?.() ?? "uninstalled"; + }); +} + +test.describe("relay restart live gate", () => { + test.skip(!enabled, "set BUZZ_E2E_RELAY_RESTART=1 to run live gate"); + + test("client reconnects after the relay is SIGTERMed and restarted", async ({ + page, + }) => { + test.setTimeout(180_000); + const portBase = 26_000 + (process.pid % 3_000); + const spec: RelaySpec = { + name: "relay-restart", + ports: { + main: portBase, + health: portBase + 3_000, + metrics: portBase + 6_000, + }, + databaseUrl: required( + "BUZZ_E2E_DATABASE_URL", + process.env.BUZZ_E2E_DATABASE_URL, + ), + redisUrl: + process.env.BUZZ_E2E_REDIS_RESTART ?? "redis://127.0.0.1:6379/13", + }; + const harness = await TwoRelayHarness.create([spec]); + try { + await harness.startRelays(); + + const relayHttpUrl = `http://127.0.0.1:${spec.ports.main}`; + await installBridge(page, { + mode: "relay", + user: "tyler", + relayHttpUrl, + relayWsUrl: `ws://127.0.0.1:${spec.ports.main}`, + }); + await page.goto("/"); + + // Baseline: the app converges to a live authenticated session. + await expect + .poll(() => connectionState(page), { timeout: 60_000 }) + .toBe("connected"); + + // Roll the pod. Graceful drain: readiness 503 → 5s grace → 1012 close + // broadcast → process exit. The client must observe the close (not a + // silent stall) and start retrying. + await harness.terminateRelayGracefully(spec.name); + await expect + .poll(() => connectionState(page), { timeout: 30_000 }) + .not.toBe("connected"); + + // Bring the "new pod" up on the same address, exactly like a k8s + // restart behind a stable service endpoint. + await harness.restartRelay(spec.name); + + // The client's retry loop must find the fresh relay and converge back + // to connected without any user interaction. + await expect + .poll(() => connectionState(page), { timeout: 60_000 }) + .toBe("connected"); + } catch (error) { + console.error(await harness.logs()); + throw error; + } finally { + await harness.stop(); + } + }); +}); diff --git a/desktop/tests/helpers/bridge.ts b/desktop/tests/helpers/bridge.ts index 36b80aef2..21f69776d 100644 --- a/desktop/tests/helpers/bridge.ts +++ b/desktop/tests/helpers/bridge.ts @@ -261,6 +261,8 @@ type MockBridgeOptions = { openerError?: string; /** Delay binding signatures so specs can exercise request supersession. */ nostrBindSignDelayMs?: number; + /** Reject successive mock WebSocket connect attempts, then resume. */ + websocketConnectErrors?: string[]; stallWebsocketSends?: boolean; userSearchDelayMs?: number; /**