From 6f67f36578e3b52f24aa99ff67f50257a8597406 Mon Sep 17 00:00:00 2001 From: Keith Smith Date: Wed, 23 Sep 2026 12:37:58 -0600 Subject: [PATCH] Add WS heartbeat and harden broadcast against dead sockets (#76) Messages could silently stop arriving until a room switch happened to paper over it with a REST refetch. Root cause: a connection that dies without a clean close (laptop sleep, an idle connection recycled by a proxy/NAT) looks identical to a healthy-but-quiet one -- readyState stays OPEN, onclose never fires, so the already-correct reconnect and rejoin logic never runs. The backend now pings every 30s and the frontend force-closes the socket if nothing (any frame, not just a pong) arrives within 65s, routing straight through the existing reconnect/rejoin path instead of needing separate recovery logic. Also found and fixed a more severe adjacent bug while tracing this: ConnectionManager.broadcast()/send_to_user() had no per-socket error handling, so one dead or racing socket could raise inside the app's single, never-restarted Redis pubsub listener task and silently kill live delivery for every room and every user on that process until restart. Now caught and logged per-socket instead of propagating. Co-Authored-By: Claude Sonnet 5 --- backend/app/ws/broadcaster.py | 31 +++++++++++++----- backend/app/ws/chat.py | 34 ++++++++++++++++++- backend/app/ws/connection_manager.py | 21 ++++++++++-- backend/tests/test_ws_chat.py | 18 ++++++++++ frontend/src/ws/useChatSocket.ts | 49 +++++++++++++++++++++++++--- 5 files changed, 137 insertions(+), 16 deletions(-) diff --git a/backend/app/ws/broadcaster.py b/backend/app/ws/broadcaster.py index 32f9543..8d01029 100644 --- a/backend/app/ws/broadcaster.py +++ b/backend/app/ws/broadcaster.py @@ -1,10 +1,13 @@ import json +import logging import uuid from redis.asyncio import Redis from app.ws.connection_manager import ConnectionManager +logger = logging.getLogger(__name__) + ROOM_CHANNEL_PREFIX = "room:" USER_CHANNEL_PREFIX = "user:" @@ -43,14 +46,26 @@ class Broadcaster: async for message in pubsub.listen(): if message["type"] != "pmessage": continue - channel = message["channel"] - payload = json.loads(message["data"]) - if channel.startswith(ROOM_CHANNEL_PREFIX): - room_id = uuid.UUID(channel.removeprefix(ROOM_CHANNEL_PREFIX)) - await self._manager.broadcast(room_id, payload) - elif channel.startswith(USER_CHANNEL_PREFIX): - user_id = uuid.UUID(channel.removeprefix(USER_CHANNEL_PREFIX)) - await self._manager.send_to_user(user_id, payload) + # #76: this is the one pubsub listener for the whole process + # (see main.py's lifespan) -- ConnectionManager.broadcast/ + # send_to_user already guard against one dead socket taking + # the rest of a single fan-out down, but this belt-and- + # suspenders catch is for anything else unexpected (a + # malformed payload, e.g.) doing the same. An uncaught + # exception here previously meant *this whole async for* + # loop died silently -- no more live delivery to any room on + # this instance, for anyone, until the process restarted. + try: + channel = message["channel"] + payload = json.loads(message["data"]) + if channel.startswith(ROOM_CHANNEL_PREFIX): + room_id = uuid.UUID(channel.removeprefix(ROOM_CHANNEL_PREFIX)) + await self._manager.broadcast(room_id, payload) + elif channel.startswith(USER_CHANNEL_PREFIX): + user_id = uuid.UUID(channel.removeprefix(USER_CHANNEL_PREFIX)) + await self._manager.send_to_user(user_id, payload) + except Exception: + logger.error("Error handling pubsub message on channel %s", message.get("channel"), exc_info=True) finally: await pubsub.punsubscribe(f"{ROOM_CHANNEL_PREFIX}*", f"{USER_CHANNEL_PREFIX}*") await pubsub.aclose() diff --git a/backend/app/ws/chat.py b/backend/app/ws/chat.py index d015b03..630dd16 100644 --- a/backend/app/ws/chat.py +++ b/backend/app/ws/chat.py @@ -1,3 +1,4 @@ +import asyncio import uuid from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect @@ -31,6 +32,18 @@ router = APIRouter(tags=["ws"]) WS_UNAUTHENTICATED = 4401 +# #76: comfortably under the ~60s idle-connection timeout common to reverse +# proxies/load balancers sitting in front of this endpoint, and short enough +# that the frontend's own zombie-detection watchdog (useChatSocket.ts) can +# notice a dead connection and reconnect well within a session -- a socket +# can go silently dead (no close frame at all, e.g. a laptop sleeping or a +# NAT dropping an idle mapping) without this, and no live message reaches it +# again until *something* else finally triggers a close, which used to be +# "switch rooms and back" purely by accident (that remounts the message +# list, which refetches history over plain REST -- unrelated to the socket +# itself actually recovering). +WS_PING_INTERVAL_SECONDS = 30 + class ClientEnvelope(BaseModel): type: str @@ -124,7 +137,18 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db) try: while True: - raw = await websocket.receive_json() + try: + raw = await asyncio.wait_for( + websocket.receive_json(), timeout=WS_PING_INTERVAL_SECONDS + ) + except asyncio.TimeoutError: + # Nothing received in a while -- ping rather than assume the + # connection's still good. A genuinely dead socket fails this + # send (raising, caught below same as any other disconnect); + # a merely quiet one just gets proded and the loop continues + # waiting. + await websocket.send_json({"type": "ping"}) + continue try: try: envelope = ClientEnvelope.model_validate(raw) @@ -146,6 +170,14 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db) joined_rooms.add(envelope.room_id) await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)}) + elif envelope.type == "pong": + # #76: the client's reply to our "ping" above -- nothing + # to do with it. Its only purpose is resetting the + # frontend's own liveness watchdog; the timeout on the + # receive_json() above already gets reset by receiving + # any frame at all, this one included. + pass + elif envelope.type == "leave": if envelope.room_id is None: await websocket.send_json({"type": "error", "detail": "room_id required"}) diff --git a/backend/app/ws/connection_manager.py b/backend/app/ws/connection_manager.py index 08fad8c..20b14a1 100644 --- a/backend/app/ws/connection_manager.py +++ b/backend/app/ws/connection_manager.py @@ -1,8 +1,11 @@ +import logging import uuid from collections import defaultdict from fastapi import WebSocket +logger = logging.getLogger(__name__) + class ConnectionManager: """Local, single-process WebSocket socket registry. @@ -43,8 +46,22 @@ class ConnectionManager: async def broadcast(self, room_id: uuid.UUID, payload: dict) -> None: for websocket in list(self._rooms.get(room_id, ())): - await websocket.send_json(payload) + # #76: one dead/racing socket (e.g. its own disconnect cleanup + # hasn't finished removing it from _rooms yet) must not take the + # rest of this broadcast down with it -- this runs inside + # Broadcaster.listen()'s single, never-restarted per-process + # pubsub loop, so an unhandled exception here previously meant + # one bad socket silently killed *live delivery for every room + # on this instance*, not just the one socket, until the process + # was restarted. + try: + await websocket.send_json(payload) + except Exception: + logger.warning("Dropping dead socket during broadcast to room %s", room_id, exc_info=True) async def send_to_user(self, user_id: uuid.UUID, payload: dict) -> None: for websocket in list(self._users.get(user_id, ())): - await websocket.send_json(payload) + try: + await websocket.send_json(payload) + except Exception: + logger.warning("Dropping dead socket sending to user %s", user_id, exc_info=True) diff --git a/backend/tests/test_ws_chat.py b/backend/tests/test_ws_chat.py index f808afe..e85e8da 100644 --- a/backend/tests/test_ws_chat.py +++ b/backend/tests/test_ws_chat.py @@ -109,6 +109,24 @@ def test_history_tiebreaks_identical_timestamps_by_id(ws_client): assert [m["content"] for m in history] == [m["content"] for m in expected_order] +def test_ws_sends_heartbeat_ping(ws_client, monkeypatch): + # #76: with nothing else sent on the connection, a "ping" must arrive + # proactively once the idle timeout elapses -- shrunk here so the test + # doesn't actually wait the real 30s. Replying "pong" (the frontend's + # own behavior) must be accepted silently and leave the connection + # fully usable afterward, not treated as an unknown message type. + monkeypatch.setattr("app.ws.chat.WS_PING_INTERVAL_SECONDS", 0.05) + username = _unique("alice") + _register(ws_client, username=username) + room = ws_client.post("/api/rooms", json={"name": _unique("general")}).json() + + with ws_client.websocket_connect("/ws/chat") as ws: + assert ws.receive_json() == {"type": "ping"} + ws.send_json({"type": "pong"}) + ws.send_json({"type": "join", "room_id": room["id"]}) + assert ws.receive_json() == {"type": "joined", "room_id": room["id"]} + + def test_ws_message_without_join_errors(ws_client): _register(ws_client, username=_unique("alice")) room = ws_client.post("/api/rooms", json={"name": _unique("general")}).json() diff --git a/frontend/src/ws/useChatSocket.ts b/frontend/src/ws/useChatSocket.ts index 1d36120..1b7ba4a 100644 --- a/frontend/src/ws/useChatSocket.ts +++ b/frontend/src/ws/useChatSocket.ts @@ -10,6 +10,13 @@ interface UseChatSocketOptions { const RECONNECT_BASE_DELAY_MS = 1000 const RECONNECT_MAX_DELAY_MS = 30000 +// #76: the backend pings every 30s (WS_PING_INTERVAL_SECONDS in chat.py) -- +// this needs to comfortably tolerate one missed ping plus ordinary network +// jitter before declaring the connection dead, not fire on the very first +// late beat. +const ZOMBIE_TIMEOUT_MS = 65000 +const ZOMBIE_CHECK_INTERVAL_MS = 10000 + const desktopMode = isDesktopNotificationsSupported() // Room join/leave (live message delivery) depends only on visibility -- @@ -51,6 +58,16 @@ export function useChatSocket({ onUnauthenticated }: UseChatSocketOptions) { // live WebSocket delivery the user can't actually see will do the job. const desiredRoomsRef = useRef(new Set()) const isVisibleRef = useRef(isVisible()) + // #76: last time *anything* arrived on the socket (a real message, a + // "ping", doesn't matter which) -- the only signal that a connection + // sitting at readyState OPEN is actually still alive rather than a + // zombie. Without this, a connection that dies without ever sending a + // close frame (laptop sleep, a NAT silently dropping an idle mapping) + // looks identical to a healthy-but-quiet one: onclose never fires, so the + // already-correct reconnect+rejoin logic below never runs, and new + // messages just stop arriving until something unrelated (e.g. switching + // rooms, which refetches history over plain REST) papers over it. + const lastActivityRef = useRef(Date.now()) const sendRoomFrame = useCallback((type: 'join' | 'leave', roomId: string) => { const ws = socketRef.current @@ -99,6 +116,7 @@ export function useChatSocket({ onUnauthenticated }: UseChatSocketOptions) { if (hasConnectedBefore) checkForUpdate() hasConnectedBefore = true reconnectDelay = RECONNECT_BASE_DELAY_MS + lastActivityRef.current = Date.now() setConnected(true) // Re-join whatever rooms were joined before a reconnect -- the // server has no memory of a dropped connection's prior state. Only @@ -122,7 +140,17 @@ export function useChatSocket({ onUnauthenticated }: UseChatSocketOptions) { ws.onmessage = (event) => { if (socketRef.current !== ws) return - const envelope = JSON.parse(event.data) as ServerEnvelope + lastActivityRef.current = Date.now() + const parsed = JSON.parse(event.data) + // #76: the server's heartbeat -- purely transport-level, not a + // domain event, so it's answered and swallowed right here instead + // of being forwarded to subscribers (ServerEnvelope's own type + // doesn't include it for exactly that reason). + if (parsed.type === 'ping') { + ws.send(JSON.stringify({ type: 'pong' })) + return + } + const envelope = parsed as ServerEnvelope for (const handler of subscribersRef.current) handler(envelope) } @@ -143,10 +171,9 @@ export function useChatSocket({ onUnauthenticated }: UseChatSocketOptions) { if (stopped) return // Unexpected close -- a deploy restarting the app server, a brief - // network blip, or (absent any app-level ping/pong) an idle - // connection getting recycled by a reverse proxy in front of it. - // Retry with exponential backoff instead of leaving the chat - // silently dead until the user manually reloads. + // network blip, or an idle connection getting recycled by a reverse + // proxy in front of it. Retry with exponential backoff instead of + // leaving the chat silently dead until the user manually reloads. reconnectTimer = setTimeout(connect, reconnectDelay) reconnectDelay = Math.min(reconnectDelay * 2, RECONNECT_MAX_DELAY_MS) } @@ -154,9 +181,21 @@ export function useChatSocket({ onUnauthenticated }: UseChatSocketOptions) { connect() + // #76: catches the case onclose can't -- a connection that dies without + // ever sending a close frame at all (readyState stays OPEN forever) is + // otherwise invisible to this hook. Forcing a close here just routes it + // through the exact same onclose/reconnect/rejoin path as a normal + // disconnect, rather than needing separate recovery logic of its own. + const zombieCheckTimer = setInterval(() => { + if (Date.now() - lastActivityRef.current > ZOMBIE_TIMEOUT_MS) { + socketRef.current?.close() + } + }, ZOMBIE_CHECK_INTERVAL_MS) + return () => { stopped = true if (reconnectTimer) clearTimeout(reconnectTimer) + clearInterval(zombieCheckTimer) socketRef.current?.close() socketRef.current = null }