Phase 5: Redis pub/sub for horizontal scaling

Splits the WebSocket layer into three pieces so one app instance and many
behave identically: ConnectionManager stays a purely local socket registry;
RoomBroadcaster publishes chat messages to a per-room Redis channel and
every instance (including the publisher) forwards received messages to its
own local sockets via a single psubscribe("room:*") listener started in
main.py's lifespan; Presence is a Redis-backed refcounted hash per room
tracking who's connected across all instances.

Presence replaces the old process-local connected_user_ids check that
Phase 4's offline-push logic used -- without it, a user connected on a
different instance would look offline and get a redundant push. Fixing
this was scoped in beyond the issue's literal ask (message fan-out only)
since it's a real correctness gap in a phase specifically about running
more than one instance; a known limitation (no heartbeat/TTL, so a hard
crash leaks a presence increment) is documented in the README instead of
solved here.

New tests/test_broadcast.py spins up two independent app instances sharing
one Postgres + Redis to prove delivery and presence both actually cross
the Redis boundary, not just work in-process. Manually verified the same
thing against two real uvicorn processes on different ports.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-14 07:13:06 -06:00
co-authored by Claude Sonnet 5
parent d09bf4a30a
commit 0b995ef75f
11 changed files with 328 additions and 57 deletions
+4
View File
@@ -16,5 +16,9 @@ class Settings(BaseSettings):
vapid_private_key: str | None = None
vapid_subject: str = "mailto:admin@example.com"
# Cross-instance WebSocket fan-out + presence (ARCHITECTURE.md phase 5).
# No credentials in a local-dev default, unlike database_url.
redis_url: str = "redis://localhost:6379/0"
settings = Settings()
+25 -1
View File
@@ -1,14 +1,38 @@
import asyncio
import contextlib
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastapi import FastAPI
from redis.asyncio import Redis
from starlette.middleware.sessions import SessionMiddleware
from app.config import settings
from app.routers import auth, health, invites, push, rooms
from app.ws.broadcaster import RoomBroadcaster
from app.ws.chat import router as ws_router
from app.ws.connection_manager import ConnectionManager
from app.ws.presence import Presence
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
redis = Redis.from_url(settings.redis_url, decode_responses=True)
app.state.presence = Presence(redis)
broadcaster = RoomBroadcaster(redis, app.state.connection_manager)
app.state.broadcaster = broadcaster
listener_task = asyncio.create_task(broadcaster.listen())
yield
listener_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await listener_task
await redis.aclose()
def create_app() -> FastAPI:
app = FastAPI(title="KeepItTalking")
app = FastAPI(title="KeepItTalking", lifespan=lifespan)
app.add_middleware(
SessionMiddleware,
+41
View File
@@ -0,0 +1,41 @@
import json
import uuid
from redis.asyncio import Redis
from app.ws.connection_manager import ConnectionManager
ROOM_CHANNEL_PREFIX = "room:"
class RoomBroadcaster:
"""Cross-instance message fan-out (ARCHITECTURE.md phase 5).
Publishes to a per-room Redis channel; every app instance -- including
the one that published -- subscribes via a single pattern subscription
and forwards to its own locally connected WebSocket clients via
ConnectionManager. A single instance just talks to itself through Redis,
so there's no separate code path for the 1-instance vs N-instance case.
"""
def __init__(self, redis: Redis, manager: ConnectionManager) -> None:
self._redis = redis
self._manager = manager
async def publish(self, room_id: uuid.UUID, payload: dict) -> None:
await self._redis.publish(f"{ROOM_CHANNEL_PREFIX}{room_id}", json.dumps(payload))
async def listen(self) -> None:
pubsub = self._redis.pubsub()
await pubsub.psubscribe(f"{ROOM_CHANNEL_PREFIX}*")
try:
async for message in pubsub.listen():
if message["type"] != "pmessage":
continue
channel = message["channel"]
room_id = uuid.UUID(channel.removeprefix(ROOM_CHANNEL_PREFIX))
payload = json.loads(message["data"])
await self._manager.broadcast(room_id, payload)
finally:
await pubsub.punsubscribe(f"{ROOM_CHANNEL_PREFIX}*")
await pubsub.aclose()
+12 -6
View File
@@ -9,7 +9,7 @@ from app.database import get_db
from app.models import Room, RoomMembership, User
from app.services.message_service import create_message
from app.services.push_service import send_push_to_user
from app.ws.connection_manager import ConnectionManager
from app.ws.presence import Presence
router = APIRouter(tags=["ws"])
@@ -33,7 +33,7 @@ async def _is_room_member(db: AsyncSession, room_id: uuid.UUID, user_id: uuid.UU
async def _notify_offline_members(
db: AsyncSession,
manager: ConnectionManager,
presence: Presence,
room_id: uuid.UUID,
sender: User,
content: str,
@@ -42,7 +42,7 @@ async def _notify_offline_members(
select(RoomMembership.user_id).where(RoomMembership.room_id == room_id)
)
member_ids = {row[0] for row in result.all()}
offline_ids = member_ids - manager.connected_user_ids(room_id)
offline_ids = member_ids - await presence.connected_user_ids(room_id)
if not offline_ids:
return
@@ -70,6 +70,8 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
await websocket.accept()
manager = websocket.app.state.connection_manager
presence: Presence = websocket.app.state.presence
broadcaster = websocket.app.state.broadcaster
joined_rooms: set[uuid.UUID] = set()
try:
@@ -90,7 +92,8 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
{"type": "error", "detail": "Not a member of this room"}
)
continue
manager.join(envelope.room_id, websocket, user.id)
manager.join(envelope.room_id, websocket)
await presence.join(envelope.room_id, user.id)
joined_rooms.add(envelope.room_id)
await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)})
@@ -99,6 +102,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
await websocket.send_json({"type": "error", "detail": "room_id required"})
continue
manager.leave(envelope.room_id, websocket)
await presence.leave(envelope.room_id, user.id)
joined_rooms.discard(envelope.room_id)
elif envelope.type == "message":
@@ -115,7 +119,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
)
continue
message = await create_message(db, envelope.room_id, user.id, envelope.content)
await manager.broadcast(
await broadcaster.publish(
envelope.room_id,
{
"type": "message",
@@ -128,7 +132,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
},
)
await _notify_offline_members(
db, manager, envelope.room_id, user, envelope.content
db, presence, envelope.room_id, user, envelope.content
)
else:
@@ -140,3 +144,5 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
pass
finally:
manager.leave_all(websocket)
for room_id in joined_rooms:
await presence.leave(room_id, user.id)
+6 -18
View File
@@ -5,22 +5,19 @@ from fastapi import WebSocket
class ConnectionManager:
"""In-memory, single-process WebSocket registry.
"""Local, single-process WebSocket socket registry.
Correct for a single app-server instance only; cross-instance fan-out via
Redis pub/sub is a later phase (ARCHITECTURE.md phase 5).
Purely about delivering to sockets connected to *this* process --
cross-instance fan-out lives in RoomBroadcaster, and cross-instance
"who's connected" for push lives in Presence, both backed by Redis
(ARCHITECTURE.md phase 5).
"""
def __init__(self) -> None:
self._rooms: dict[uuid.UUID, set[WebSocket]] = defaultdict(set)
# A single connection can be joined to multiple rooms at once (one
# `join` message per room over the same socket), so this is keyed on
# the socket alone, not per-room.
self._ws_user: dict[WebSocket, uuid.UUID] = {}
def join(self, room_id: uuid.UUID, websocket: WebSocket, user_id: uuid.UUID) -> None:
def join(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
self._rooms[room_id].add(websocket)
self._ws_user[websocket] = user_id
def leave(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
self._rooms[room_id].discard(websocket)
@@ -30,15 +27,6 @@ class ConnectionManager:
def leave_all(self, websocket: WebSocket) -> None:
for room_id in list(self._rooms.keys()):
self.leave(room_id, websocket)
self._ws_user.pop(websocket, None)
def connected_user_ids(self, room_id: uuid.UUID) -> set[uuid.UUID]:
"""Users (not just sockets) with an active connection to this room --
used to skip push notifications for anyone already watching, per
ARCHITECTURE.md's "members with no active connection" push flow."""
return {
self._ws_user[ws] for ws in self._rooms.get(room_id, ()) if ws in self._ws_user
}
async def broadcast(self, room_id: uuid.UUID, payload: dict) -> None:
for websocket in list(self._rooms.get(room_id, ())):
+39
View File
@@ -0,0 +1,39 @@
import uuid
from redis.asyncio import Redis
class Presence:
"""Cross-instance "who's connected to this room," backed by a Redis hash
per room (field = user_id, value = connection refcount).
Refcounted rather than a plain set so a user with two connections to the
same room -- two tabs, or one per app instance -- doesn't get marked
offline when only one of those connections closes.
Known limitation: a hard crash (not a clean disconnect) leaks that
connection's increment forever, since there's no heartbeat/TTL here to
reclaim it -- out of scope for this phase, same category of
simplification as the "no server-side session revocation" note in the
README.
"""
def __init__(self, redis: Redis) -> None:
self._redis = redis
def _key(self, room_id: uuid.UUID) -> str:
return f"presence:{room_id}"
async def join(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
await self._redis.hincrby(self._key(room_id), str(user_id), 1)
async def leave(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
key = self._key(room_id)
field = str(user_id)
remaining = await self._redis.hincrby(key, field, -1)
if remaining <= 0:
await self._redis.hdel(key, field)
async def connected_user_ids(self, room_id: uuid.UUID) -> set[uuid.UUID]:
fields = await self._redis.hkeys(self._key(room_id))
return {uuid.UUID(f) for f in fields}