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
+1
View File
@@ -1,6 +1,7 @@
DATABASE_URL=postgresql+asyncpg://chatapp:chatapp@localhost:5432/chatapp
SESSION_SECRET=change-me-to-a-long-random-string
SESSION_HTTPS_ONLY=false
REDIS_URL=redis://localhost:6379/0
# Optional: push notifications are skipped if unset. Generate with:
# python -m app.cli generate-vapid-keys
+58 -17
View File
@@ -1,10 +1,10 @@
# KeepItTalking backend (Phase 1 + 2 + 4)
# KeepItTalking backend (Phase 1 + 2 + 4 + 5)
FastAPI + SQLAlchemy 2.0 (async) + PostgreSQL. Implements auth, room CRUD
(open and private), room roles (owner/admin/member) and invites, a
single-instance WebSocket chat endpoint, and Web Push notifications for
offline room members. See `../ARCHITECTURE.md` for the full system design
and the phased build plan.
FastAPI + SQLAlchemy 2.0 (async) + PostgreSQL + Redis. Implements auth, room
CRUD (open and private), room roles (owner/admin/member) and invites, a
WebSocket chat endpoint that fans out across multiple app-server instances
via Redis pub/sub, and Web Push notifications for offline room members. See
`../ARCHITECTURE.md` for the full system design and the phased build plan.
This is an **invite-only site**: there is no public registration endpoint.
Accounts are created by an operator on the app server — see step 4 below.
@@ -30,7 +30,16 @@ docker exec chatapp-postgres psql -U chatapp -d chatapp -c "CREATE DATABASE chat
(Docker here is purely a local-dev convenience for standing up Postgres quickly —
the actual deployment target has no containers at all, see `ARCHITECTURE.md` §9.)
### 2. Python environment
### 2. Redis
Used for cross-instance WebSocket fan-out and presence (see the section
below). Required — there's no in-memory fallback.
```bash
docker run -d --name chatapp-redis -p 6379:6379 redis:7-alpine
```
### 3. Python environment
```bash
cd backend
@@ -41,13 +50,13 @@ cp .env.example .env
# python3 -c "import secrets; print(secrets.token_urlsafe(32))"
```
### 3. Migrations
### 4. Migrations
```bash
.venv/bin/alembic upgrade head
```
### 4. Create a user
### 5. Create a user
There's no public sign-up. Create accounts directly with the CLI (add
`--admin` to grant `is_site_admin`, useful ahead of the phase-6 admin portal):
@@ -56,7 +65,7 @@ There's no public sign-up. Create accounts directly with the CLI (add
.venv/bin/python -m app.cli create-user alice alice@example.com "some-password"
```
### 5. (Optional) Set up push notifications
### 6. (Optional) Set up push notifications
Push works without any setup — `VAPID_PUBLIC_KEY`/`VAPID_PRIVATE_KEY` are
unset by default and push delivery is silently skipped. To enable it:
@@ -66,7 +75,7 @@ unset by default and push delivery is silently skipped. To enable it:
# paste the three printed lines into backend/.env
```
### 6. Run the dev server
### 7. Run the dev server
```bash
.venv/bin/uvicorn app.main:app --reload
@@ -74,10 +83,16 @@ unset by default and push delivery is silently skipped. To enable it:
API docs: http://localhost:8000/docs. WebSocket chat endpoint: `ws://localhost:8000/ws/chat`.
### 7. Run tests
To try horizontal scaling locally, run a second instance on another port
against the same Postgres + Redis (`.venv/bin/uvicorn app.main:app --port 8001`)
— a message sent through one instance's WebSocket is delivered to clients
connected to the other, purely via Redis.
### 8. Run tests
Tests run against a real Postgres database (`chatapp_test` by default — native
`ENUM`/`UUID` types aren't faithfully reproduced by SQLite), with each test
`ENUM`/`UUID` types aren't faithfully reproduced by SQLite) and a real Redis
(db 15 by default, kept separate from dev use of db 0), with each test
wrapped in a transaction that's rolled back afterward:
```bash
@@ -99,20 +114,46 @@ app/
schemas/ Pydantic request/response models
routers/ auth, rooms, invites, push, health
services/ business logic called by routers
ws/ WebSocket connection manager + /ws/chat handler
ws/ connection_manager (local sockets), presence +
broadcaster (Redis), /ws/chat handler
alembic/ migrations
tests/ pytest + httpx/TestClient tests
```
## Cross-instance broadcast (Phase 5)
The WebSocket layer is split into three pieces so that running one app
instance and running many behave identically:
- `app/ws/connection_manager.py` — purely local: which sockets on *this*
process are in which room, used only to actually `send_json` to them.
- `app/ws/broadcaster.py` (`RoomBroadcaster`) — on a chat message,
`publish()`s it to a Redis channel scoped to the room (`room:{id}`).
Every app instance, including the publisher, runs a single background
`listen()` task (started in `app/main.py`'s lifespan) pattern-subscribed
to `room:*`; each message it receives is handed to its own local
`ConnectionManager.broadcast()`. One instance just talks to itself
through Redis, so there's no separate single-instance code path.
- `app/ws/presence.py` (`Presence`) — a Redis hash per room
(`presence:{room_id}`, field = user ID, value = connection refcount) is
the cross-instance answer to "is this member connected *anywhere* right
now," which is what the Phase 4 offline-push check uses instead of the
local `ConnectionManager`. Refcounted so a user connected from two tabs
(or two instances) isn't marked offline until every connection closes.
Known limitation: `Presence` has no heartbeat/TTL, so a hard process crash
(not a clean disconnect) leaks that connection's increment forever — same
category of simplification as the "no server-side session revocation" note
below.
## Push notifications (Phase 4)
`POST /api/push/subscribe` (upserts by `endpoint`) / `DELETE /api/push/subscribe`
manage a user's `push_subscriptions` rows; `GET /api/push/vapid-public-key` gives
the frontend the key it needs for `PushManager.subscribe()`. On every chat
message, `app/ws/chat.py` computes `room members - ConnectionManager.
message, `app/ws/chat.py` computes `room members - Presence.
connected_user_ids(room_id)` (who's actually connected to *that room* right
now, tracked alongside the existing WebSocket registry) and sends each
offline member a push via `pywebpush`, awaited inline against the same
now, across every app instance — see Phase 5 below) and sends each offline
member a push via `pywebpush`, awaited inline against the same
request-scoped session rather than fired as a background task — the
broadcast to online members already happened by that point, so nothing
online-facing is delayed, and it sidesteps `asyncio.create_task()`s outliving
+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}
+1
View File
@@ -15,6 +15,7 @@ dependencies = [
"argon2-cffi>=23.1",
"itsdangerous>=2.2",
"pywebpush>=2.0",
"redis>=5.0",
]
[project.scripts]
+31 -8
View File
@@ -1,3 +1,4 @@
import contextlib
import os
from pathlib import Path
@@ -6,6 +7,9 @@ os.environ.setdefault(
)
os.environ.setdefault("SESSION_SECRET", "test-secret")
os.environ.setdefault("SESSION_HTTPS_ONLY", "false")
# DB index 15 keeps test presence/pub-sub state separate from whatever a
# developer's local Redis is doing on db 0.
os.environ.setdefault("REDIS_URL", "redis://localhost:6379/15")
import pytest
import pytest_asyncio
@@ -76,15 +80,15 @@ async def client(app):
@pytest.fixture
def ws_client():
def ws_client_factory():
# Starlette's TestClient (needed for websocket_connect, which httpx's
# async client doesn't support) runs the ASGI app on a background thread
# with its own event loop via anyio's BlockingPortal. asyncpg connections
# are bound to the loop they're opened on, so this app gets its own
# engine created here (no connections opened yet) rather than reusing
# the `db_session`/`app` fixtures' engine, which belongs to pytest's
# loop. No per-test rollback here (see test_ws_chat.py for the
# unique-name convention that keeps tests independent without it).
# are bound to the loop they're opened on, so each app built here gets
# its own engine (no connections opened yet) rather than reusing the
# `db_session`/`app` fixtures' engine, which belongs to pytest's loop.
# No per-test rollback here (see test_ws_chat.py for the unique-name
# convention that keeps tests independent without it).
#
# poolclass=NullPool: with pooling, a WS test that does more than one
# DB round trip per message (e.g. the offline-push lookup) can hit a
@@ -94,6 +98,17 @@ def ws_client():
# surfaces as "connection is closed". A fresh connection per session
# sidesteps it; fine for tests, not something prod needs (prod isn't
# juggling a background portal thread against the main test thread).
#
# A factory (not a single client) so tests can spin up more than one
# independent app instance -- sharing the same test Postgres and Redis,
# like separate app-server processes behind Nginx would -- to exercise
# cross-instance broadcast/presence (see test_broadcast.py). Each
# TestClient is entered via an ExitStack so its lifespan (which opens
# the Redis connection/pubsub listener) starts immediately and all of
# them get torn down together at fixture teardown.
stack = contextlib.ExitStack()
def _make() -> TestClient:
application = create_app()
test_engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool)
test_session_factory = async_sessionmaker(test_engine, expire_on_commit=False)
@@ -104,9 +119,17 @@ def ws_client():
application.dependency_overrides[get_db] = _get_db
with TestClient(application) as tc:
tc = stack.enter_context(TestClient(application))
tc.session_factory = test_session_factory # type: ignore[attr-defined]
yield tc
return tc
yield _make
stack.close()
@pytest.fixture
def ws_client(ws_client_factory):
return ws_client_factory()
async def register_and_login(
+103
View File
@@ -0,0 +1,103 @@
import uuid
from app.schemas.user import UserCreate
from app.services.auth_service import register_user
def _unique(prefix: str) -> str:
return f"{prefix}-{uuid.uuid4().hex[:8]}"
def _register_ws(ws_client, username: str) -> dict:
async def _seed():
async with ws_client.session_factory() as session:
await register_user(
session,
UserCreate(username=username, email=f"{username}@example.com", password="password123"),
)
ws_client.portal.call(_seed)
resp = ws_client.post(
"/api/auth/login", json={"username_or_email": username, "password": "password123"}
)
assert resp.status_code == 200, resp.text
return resp.json()
def test_message_fans_out_across_instances(ws_client_factory):
# Two independent app instances -- separate ConnectionManager, separate
# Redis pubsub subscription, separate everything except the Postgres and
# Redis they're both pointed at -- the same way two app-server processes
# behind Nginx would be. Proves delivery actually crosses Redis, not
# just in-process delivery within a single ConnectionManager.
instance1 = ws_client_factory()
instance2 = ws_client_factory()
alice = _register_ws(instance1, _unique("alice"))
room = instance1.post("/api/rooms", json={"name": _unique("general")}).json()
bob = _register_ws(instance2, _unique("bob"))
instance2.post(f"/api/rooms/{room['id']}/join")
with instance2.websocket_connect("/ws/chat") as bob_ws:
bob_ws.send_json({"type": "join", "room_id": room["id"]})
assert bob_ws.receive_json()["type"] == "joined"
with instance1.websocket_connect("/ws/chat") as alice_ws:
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
alice_ws.send_json(
{"type": "message", "room_id": room["id"], "content": "hi from instance 1"}
)
assert alice_ws.receive_json()["type"] == "message"
received = bob_ws.receive_json()
assert received["type"] == "message"
assert received["content"] == "hi from instance 1"
assert received["username"] == alice["username"]
def test_presence_is_shared_across_instances(ws_client_factory, monkeypatch):
calls = []
monkeypatch.setattr("app.services.push_service.webpush", lambda **kw: calls.append(kw))
instance1 = ws_client_factory()
instance2 = ws_client_factory()
alice = _register_ws(instance1, _unique("alice"))
room = instance1.post("/api/rooms", json={"name": _unique("general")}).json()
bob = _register_ws(instance2, _unique("bob"))
instance2.post(f"/api/rooms/{room['id']}/join")
instance2.post(
"/api/push/subscribe",
json={
"endpoint": f"https://push.example.com/ep-{bob['id']}",
"keys": {"p256dh": "p256dh-bob", "auth": "auth-bob"},
},
)
with instance2.websocket_connect("/ws/chat") as bob_ws:
bob_ws.send_json({"type": "join", "room_id": room["id"]})
assert bob_ws.receive_json()["type"] == "joined"
with instance1.websocket_connect("/ws/chat") as alice_ws:
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
alice_ws.send_json({"type": "message", "room_id": room["id"], "content": "hi"})
assert alice_ws.receive_json()["type"] == "message"
# bob is connected -- just on the other instance -- so he should
# get the broadcast via Redis, not a push notification. If
# presence were still process-local (pre-phase-5 behavior) he'd
# look offline to instance1 and get a redundant push.
assert bob_ws.receive_json()["type"] == "message"
# Sync barrier: the handler processes frames strictly
# sequentially, so a second (idempotent) join only acks once the
# "message" frame's full handling -- including the offline-push
# step -- has completed on instance1.
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
assert calls == []