Private
Public Access
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 <noreply@anthropic.com>
This commit is contained in:
1 parent
b1c9a85ab3
commit
6f67f36578
5 files changed
+137
-16
No files matched your search
@@ -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()
|
||||
+33
-1
@@ -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"})
|
||||
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user