Private
Public Access
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:
@@ -1,6 +1,7 @@
|
|||||||
DATABASE_URL=postgresql+asyncpg://chatapp:chatapp@localhost:5432/chatapp
|
DATABASE_URL=postgresql+asyncpg://chatapp:chatapp@localhost:5432/chatapp
|
||||||
SESSION_SECRET=change-me-to-a-long-random-string
|
SESSION_SECRET=change-me-to-a-long-random-string
|
||||||
SESSION_HTTPS_ONLY=false
|
SESSION_HTTPS_ONLY=false
|
||||||
|
REDIS_URL=redis://localhost:6379/0
|
||||||
|
|
||||||
# Optional: push notifications are skipped if unset. Generate with:
|
# Optional: push notifications are skipped if unset. Generate with:
|
||||||
# python -m app.cli generate-vapid-keys
|
# python -m app.cli generate-vapid-keys
|
||||||
|
|||||||
+58
-17
@@ -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
|
FastAPI + SQLAlchemy 2.0 (async) + PostgreSQL + Redis. Implements auth, room
|
||||||
(open and private), room roles (owner/admin/member) and invites, a
|
CRUD (open and private), room roles (owner/admin/member) and invites, a
|
||||||
single-instance WebSocket chat endpoint, and Web Push notifications for
|
WebSocket chat endpoint that fans out across multiple app-server instances
|
||||||
offline room members. See `../ARCHITECTURE.md` for the full system design
|
via Redis pub/sub, and Web Push notifications for offline room members. See
|
||||||
and the phased build plan.
|
`../ARCHITECTURE.md` for the full system design and the phased build plan.
|
||||||
|
|
||||||
This is an **invite-only site**: there is no public registration endpoint.
|
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.
|
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 —
|
(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.)
|
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
|
```bash
|
||||||
cd backend
|
cd backend
|
||||||
@@ -41,13 +50,13 @@ cp .env.example .env
|
|||||||
# python3 -c "import secrets; print(secrets.token_urlsafe(32))"
|
# python3 -c "import secrets; print(secrets.token_urlsafe(32))"
|
||||||
```
|
```
|
||||||
|
|
||||||
### 3. Migrations
|
### 4. Migrations
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
.venv/bin/alembic upgrade head
|
.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
|
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):
|
`--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"
|
.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
|
Push works without any setup — `VAPID_PUBLIC_KEY`/`VAPID_PRIVATE_KEY` are
|
||||||
unset by default and push delivery is silently skipped. To enable it:
|
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
|
# paste the three printed lines into backend/.env
|
||||||
```
|
```
|
||||||
|
|
||||||
### 6. Run the dev server
|
### 7. Run the dev server
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
.venv/bin/uvicorn app.main:app --reload
|
.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`.
|
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
|
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:
|
wrapped in a transaction that's rolled back afterward:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
@@ -99,20 +114,46 @@ app/
|
|||||||
schemas/ Pydantic request/response models
|
schemas/ Pydantic request/response models
|
||||||
routers/ auth, rooms, invites, push, health
|
routers/ auth, rooms, invites, push, health
|
||||||
services/ business logic called by routers
|
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
|
alembic/ migrations
|
||||||
tests/ pytest + httpx/TestClient tests
|
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)
|
## Push notifications (Phase 4)
|
||||||
|
|
||||||
`POST /api/push/subscribe` (upserts by `endpoint`) / `DELETE /api/push/subscribe`
|
`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
|
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
|
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
|
connected_user_ids(room_id)` (who's actually connected to *that room* right
|
||||||
now, tracked alongside the existing WebSocket registry) and sends each
|
now, across every app instance — see Phase 5 below) and sends each offline
|
||||||
offline member a push via `pywebpush`, awaited inline against the same
|
member a push via `pywebpush`, awaited inline against the same
|
||||||
request-scoped session rather than fired as a background task — the
|
request-scoped session rather than fired as a background task — the
|
||||||
broadcast to online members already happened by that point, so nothing
|
broadcast to online members already happened by that point, so nothing
|
||||||
online-facing is delayed, and it sidesteps `asyncio.create_task()`s outliving
|
online-facing is delayed, and it sidesteps `asyncio.create_task()`s outliving
|
||||||
|
|||||||
@@ -16,5 +16,9 @@ class Settings(BaseSettings):
|
|||||||
vapid_private_key: str | None = None
|
vapid_private_key: str | None = None
|
||||||
vapid_subject: str = "mailto:admin@example.com"
|
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()
|
settings = Settings()
|
||||||
|
|||||||
+25
-1
@@ -1,14 +1,38 @@
|
|||||||
|
import asyncio
|
||||||
|
import contextlib
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
|
from redis.asyncio import Redis
|
||||||
from starlette.middleware.sessions import SessionMiddleware
|
from starlette.middleware.sessions import SessionMiddleware
|
||||||
|
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
from app.routers import auth, health, invites, push, rooms
|
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.chat import router as ws_router
|
||||||
from app.ws.connection_manager import ConnectionManager
|
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:
|
def create_app() -> FastAPI:
|
||||||
app = FastAPI(title="KeepItTalking")
|
app = FastAPI(title="KeepItTalking", lifespan=lifespan)
|
||||||
|
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
SessionMiddleware,
|
SessionMiddleware,
|
||||||
|
|||||||
@@ -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
@@ -9,7 +9,7 @@ from app.database import get_db
|
|||||||
from app.models import Room, RoomMembership, User
|
from app.models import Room, RoomMembership, User
|
||||||
from app.services.message_service import create_message
|
from app.services.message_service import create_message
|
||||||
from app.services.push_service import send_push_to_user
|
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"])
|
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(
|
async def _notify_offline_members(
|
||||||
db: AsyncSession,
|
db: AsyncSession,
|
||||||
manager: ConnectionManager,
|
presence: Presence,
|
||||||
room_id: uuid.UUID,
|
room_id: uuid.UUID,
|
||||||
sender: User,
|
sender: User,
|
||||||
content: str,
|
content: str,
|
||||||
@@ -42,7 +42,7 @@ async def _notify_offline_members(
|
|||||||
select(RoomMembership.user_id).where(RoomMembership.room_id == room_id)
|
select(RoomMembership.user_id).where(RoomMembership.room_id == room_id)
|
||||||
)
|
)
|
||||||
member_ids = {row[0] for row in result.all()}
|
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:
|
if not offline_ids:
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -70,6 +70,8 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
|||||||
|
|
||||||
await websocket.accept()
|
await websocket.accept()
|
||||||
manager = websocket.app.state.connection_manager
|
manager = websocket.app.state.connection_manager
|
||||||
|
presence: Presence = websocket.app.state.presence
|
||||||
|
broadcaster = websocket.app.state.broadcaster
|
||||||
joined_rooms: set[uuid.UUID] = set()
|
joined_rooms: set[uuid.UUID] = set()
|
||||||
|
|
||||||
try:
|
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"}
|
{"type": "error", "detail": "Not a member of this room"}
|
||||||
)
|
)
|
||||||
continue
|
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)
|
joined_rooms.add(envelope.room_id)
|
||||||
await websocket.send_json({"type": "joined", "room_id": str(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"})
|
await websocket.send_json({"type": "error", "detail": "room_id required"})
|
||||||
continue
|
continue
|
||||||
manager.leave(envelope.room_id, websocket)
|
manager.leave(envelope.room_id, websocket)
|
||||||
|
await presence.leave(envelope.room_id, user.id)
|
||||||
joined_rooms.discard(envelope.room_id)
|
joined_rooms.discard(envelope.room_id)
|
||||||
|
|
||||||
elif envelope.type == "message":
|
elif envelope.type == "message":
|
||||||
@@ -115,7 +119,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
|||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
message = await create_message(db, envelope.room_id, user.id, envelope.content)
|
message = await create_message(db, envelope.room_id, user.id, envelope.content)
|
||||||
await manager.broadcast(
|
await broadcaster.publish(
|
||||||
envelope.room_id,
|
envelope.room_id,
|
||||||
{
|
{
|
||||||
"type": "message",
|
"type": "message",
|
||||||
@@ -128,7 +132,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
await _notify_offline_members(
|
await _notify_offline_members(
|
||||||
db, manager, envelope.room_id, user, envelope.content
|
db, presence, envelope.room_id, user, envelope.content
|
||||||
)
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
@@ -140,3 +144,5 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
|
|||||||
pass
|
pass
|
||||||
finally:
|
finally:
|
||||||
manager.leave_all(websocket)
|
manager.leave_all(websocket)
|
||||||
|
for room_id in joined_rooms:
|
||||||
|
await presence.leave(room_id, user.id)
|
||||||
|
|||||||
@@ -5,22 +5,19 @@ from fastapi import WebSocket
|
|||||||
|
|
||||||
|
|
||||||
class ConnectionManager:
|
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
|
Purely about delivering to sockets connected to *this* process --
|
||||||
Redis pub/sub is a later phase (ARCHITECTURE.md phase 5).
|
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:
|
def __init__(self) -> None:
|
||||||
self._rooms: dict[uuid.UUID, set[WebSocket]] = defaultdict(set)
|
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._rooms[room_id].add(websocket)
|
||||||
self._ws_user[websocket] = user_id
|
|
||||||
|
|
||||||
def leave(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
|
def leave(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
|
||||||
self._rooms[room_id].discard(websocket)
|
self._rooms[room_id].discard(websocket)
|
||||||
@@ -30,15 +27,6 @@ class ConnectionManager:
|
|||||||
def leave_all(self, websocket: WebSocket) -> None:
|
def leave_all(self, websocket: WebSocket) -> None:
|
||||||
for room_id in list(self._rooms.keys()):
|
for room_id in list(self._rooms.keys()):
|
||||||
self.leave(room_id, websocket)
|
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:
|
async def broadcast(self, room_id: uuid.UUID, payload: dict) -> None:
|
||||||
for websocket in list(self._rooms.get(room_id, ())):
|
for websocket in list(self._rooms.get(room_id, ())):
|
||||||
|
|||||||
@@ -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}
|
||||||
@@ -15,6 +15,7 @@ dependencies = [
|
|||||||
"argon2-cffi>=23.1",
|
"argon2-cffi>=23.1",
|
||||||
"itsdangerous>=2.2",
|
"itsdangerous>=2.2",
|
||||||
"pywebpush>=2.0",
|
"pywebpush>=2.0",
|
||||||
|
"redis>=5.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
|
|||||||
+38
-15
@@ -1,3 +1,4 @@
|
|||||||
|
import contextlib
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -6,6 +7,9 @@ os.environ.setdefault(
|
|||||||
)
|
)
|
||||||
os.environ.setdefault("SESSION_SECRET", "test-secret")
|
os.environ.setdefault("SESSION_SECRET", "test-secret")
|
||||||
os.environ.setdefault("SESSION_HTTPS_ONLY", "false")
|
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
|
||||||
import pytest_asyncio
|
import pytest_asyncio
|
||||||
@@ -76,15 +80,15 @@ async def client(app):
|
|||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def ws_client():
|
def ws_client_factory():
|
||||||
# Starlette's TestClient (needed for websocket_connect, which httpx's
|
# Starlette's TestClient (needed for websocket_connect, which httpx's
|
||||||
# async client doesn't support) runs the ASGI app on a background thread
|
# async client doesn't support) runs the ASGI app on a background thread
|
||||||
# with its own event loop via anyio's BlockingPortal. asyncpg connections
|
# 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
|
# are bound to the loop they're opened on, so each app built here gets
|
||||||
# engine created here (no connections opened yet) rather than reusing
|
# its own engine (no connections opened yet) rather than reusing the
|
||||||
# the `db_session`/`app` fixtures' engine, which belongs to pytest's
|
# `db_session`/`app` fixtures' engine, which belongs to pytest's loop.
|
||||||
# loop. No per-test rollback here (see test_ws_chat.py for the
|
# No per-test rollback here (see test_ws_chat.py for the unique-name
|
||||||
# unique-name convention that keeps tests independent without it).
|
# convention that keeps tests independent without it).
|
||||||
#
|
#
|
||||||
# poolclass=NullPool: with pooling, a WS test that does more than one
|
# 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
|
# DB round trip per message (e.g. the offline-push lookup) can hit a
|
||||||
@@ -94,19 +98,38 @@ def ws_client():
|
|||||||
# surfaces as "connection is closed". A fresh connection per session
|
# surfaces as "connection is closed". A fresh connection per session
|
||||||
# sidesteps it; fine for tests, not something prod needs (prod isn't
|
# sidesteps it; fine for tests, not something prod needs (prod isn't
|
||||||
# juggling a background portal thread against the main test thread).
|
# juggling a background portal thread against the main test thread).
|
||||||
application = create_app()
|
#
|
||||||
test_engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool)
|
# A factory (not a single client) so tests can spin up more than one
|
||||||
test_session_factory = async_sessionmaker(test_engine, expire_on_commit=False)
|
# 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()
|
||||||
|
|
||||||
async def _get_db():
|
def _make() -> TestClient:
|
||||||
async with test_session_factory() as session:
|
application = create_app()
|
||||||
yield session
|
test_engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool)
|
||||||
|
test_session_factory = async_sessionmaker(test_engine, expire_on_commit=False)
|
||||||
|
|
||||||
application.dependency_overrides[get_db] = _get_db
|
async def _get_db():
|
||||||
|
async with test_session_factory() as session:
|
||||||
|
yield session
|
||||||
|
|
||||||
with TestClient(application) as tc:
|
application.dependency_overrides[get_db] = _get_db
|
||||||
|
|
||||||
|
tc = stack.enter_context(TestClient(application))
|
||||||
tc.session_factory = test_session_factory # type: ignore[attr-defined]
|
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(
|
async def register_and_login(
|
||||||
|
|||||||
@@ -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 == []
|
||||||
Reference in New Issue
Block a user