WebSockets — Real-Time, Two-Way Connections
AdvancedA @app.websocket endpoint accepts a persistent two-way connection: await receive/send in a loop, a ConnectionManager broadcasts to rooms, auth happens at the handshake, and multi-instance fan-out needs Redis pub/sub.
Overview
HTTP is request-reply; a chat, a live scoreboard, or a collaborative editor needs the server to push whenever it likes. A WebSocket starts as an HTTP request that upgrades into a persistent, full-duplex TCP channel — either side sends at any time until someone disconnects. FastAPI (via Starlette) makes the endpoint an async loop: accept(), then receive/send until WebSocketDisconnect. Real apps immediately need a ConnectionManager — a registry of live sockets grouped into rooms so one message fans out to every participant. Auth changes shape: browsers cannot set an Authorization header on the WS handshake, so tokens arrive via query param or first-message auth, validated before accept() completes. And the architecture question that separates toy from production: sockets live on ONE instance, so broadcasting across multiple pods requires a shared bus — Redis pub/sub — carrying messages between instances.
The Endpoint + a ConnectionManager for Rooms
The endpoint is a loop: receive, act, send. The manager tracks who is in which room and broadcasts — with the disconnect cleanup that prevents ghost sockets from accumulating.
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
class ConnectionManager:
def __init__(self):
self.rooms: dict[str, set[WebSocket]] = {} # room → live sockets
async def connect(self, room: str, ws: WebSocket):
await ws.accept() # completes the handshake
self.rooms.setdefault(room, set()).add(ws)
def disconnect(self, room: str, ws: WebSocket):
self.rooms.get(room, set()).discard(ws) # ALWAYS clean up
async def broadcast(self, room: str, message: dict):
dead = []
for ws in self.rooms.get(room, set()):
try:
await ws.send_json(message)
except Exception:
dead.append(ws) # died mid-broadcast
for ws in dead:
self.disconnect(room, ws)
manager = ConnectionManager()
@app.websocket("/ws/drives/{drive_id}/chat")
async def drive_chat(ws: WebSocket, drive_id: str):
await manager.connect(drive_id, ws)
await manager.broadcast(drive_id, {"sys": "someone joined"})
try:
while True: # the connection's life
data = await ws.receive_json() # blocks until a message
await manager.broadcast(drive_id, {
"from": data.get("name", "anon"),
"text": data["text"],
})
except WebSocketDisconnect: # tab closed, network died
manager.disconnect(drive_id, ws)
await manager.broadcast(drive_id, {"sys": "someone left"})
# Browser side:
# const ws = new WebSocket("wss://api.example.com/ws/drives/7/chat")
# ws.onmessage = (e) => render(JSON.parse(e.data))
# ws.send(JSON.stringify({name: "Asha", text: "When is the TCS drive?"}))Auth at the Handshake + Scaling Past One Instance
Validate the token BEFORE accept() finishes — reject with the WS policy-violation code. Then the multi-pod problem: your in-memory manager only knows ITS sockets; Redis pub/sub is the bridge.
from fastapi import WebSocket, WebSocketException, status
import jwt
@app.websocket("/ws/notifications")
async def notifications(ws: WebSocket, token: str = ""): # ?token=... query param
# No Authorization header on browser WS handshakes — token arrives as
# a query param (or as the first message before any data flows).
try:
claims = jwt.decode(token, SECRET, algorithms=["HS256"])
except jwt.InvalidTokenError:
# reject BEFORE accepting — 1008 = policy violation
raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
await ws.accept()
await ws.send_json({"hello": claims["sub"]})
# ... loop ...
# Query-param tokens land in access logs — use short-lived, WS-only
# tickets minted by an authenticated HTTP call, not your main JWT.
# ── The multi-instance problem ──
# pod A holds Asha's socket; pod B processes the event that concerns her.
# B's ConnectionManager has no idea Asha exists. In-memory fan-out is
# single-instance only.
#
# The standard fix — Redis pub/sub as the message bus:
# every pod: subscriber task (started in lifespan) listens to
# channel "room:{id}", forwards to ITS local sockets
# any pod: publishes events to "room:{id}" instead of broadcasting
#
# async def redis_listener(app):
# pubsub = app.state.redis.pubsub()
# await pubsub.psubscribe("room:*")
# async for msg in pubsub.listen():
# if msg["type"] == "pmessage":
# room = msg["channel"].split(":", 1)[1]
# await manager.broadcast(room, json.loads(msg["data"]))
#
# Also on the production checklist: heartbeats (ping/pong) to detect
# half-dead connections, client auto-reconnect with backoff, and a
# LB/proxy timeout raised for long-lived connections (nginx
# proxy_read_timeout — default 60s kills quiet sockets).Key Points to Remember
- 1WS = HTTP upgrade to persistent full-duplex; endpoint is an accept + receive/send loop
- 2ConnectionManager tracks rooms; always clean up on WebSocketDisconnect
- 3Auth at the handshake (query-param ticket / first message) — reject with 1008 before accept
- 4Sockets are per-instance: multi-pod broadcast requires Redis pub/sub between pods
Interview Questions
Sign in to ask AriaDesign live seat-availability for a booking page — WebSocket, SSE, or polling, and why?
Your chat works on one pod but messages vanish across pods — explain and fix.
How do you authenticate a browser WebSocket when you cannot set headers?
Ask Aria about WebSockets — Real-Time, Two-Way Connections
Your personal AI tutor — ask anything about this concept
Revision Status
Personal Notes
Sign in to save personal notes for this topic.
Discussion
Sign in to join the discussion.