1. WebSocket Basics in FastAPI
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
"""WebSocket endpoint cơ bản."""
await websocket.accept()
try:
while True:
# Nhận message từ client
data = await websocket.receive_text()
# Gửi message lại client
await websocket.send_text(f"Echo: {data}")
except WebSocketDisconnect:
print("Client disconnected")
2. Connection Manager
# app/core/websocket.py
import json
from dataclasses import dataclass, field
from fastapi import WebSocket
@dataclass
class ConnectionManager:
"""Quản lý WebSocket connections."""
# Tất cả connections
active_connections: list[WebSocket] = field(default_factory=list)
# Connections theo room
rooms: dict[str, list[WebSocket]] = field(default_factory=dict)
# Connections theo user
user_connections: dict[int, list[WebSocket]] = field(default_factory=dict)
async def connect(
self, websocket: WebSocket, user_id: int | None = None, room: str | None = None
):
"""Kết nối client."""
await websocket.accept()
self.active_connections.append(websocket)
if user_id:
self.user_connections.setdefault(user_id, []).append(websocket)
if room:
self.rooms.setdefault(room, []).append(websocket)
def disconnect(
self, websocket: WebSocket, user_id: int | None = None, room: str | None = None
):
"""Ngắt kết nối client."""
if websocket in self.active_connections:
self.active_connections.remove(websocket)
if user_id and user_id in self.user_connections:
conns = self.user_connections[user_id]
if websocket in conns:
conns.remove(websocket)
if not conns:
del self.user_connections[user_id]
if room and room in self.rooms:
conns = self.rooms[room]
if websocket in conns:
conns.remove(websocket)
if not conns:
del self.rooms[room]
async def send_personal(self, message: dict, websocket: WebSocket):
"""Gửi message đến 1 client."""
await websocket.send_json(message)
async def send_to_user(self, message: dict, user_id: int):
"""Gửi message đến tất cả connections của 1 user."""
connections = self.user_connections.get(user_id, [])
for ws in connections:
try:
await ws.send_json(message)
except Exception:
pass
async def broadcast(self, message: dict, exclude: WebSocket | None = None):
"""Gửi message đến tất cả clients."""
for ws in self.active_connections:
if ws != exclude:
try:
await ws.send_json(message)
except Exception:
pass
async def broadcast_to_room(
self, message: dict, room: str, exclude: WebSocket | None = None
):
"""Gửi message đến tất cả clients trong room."""
connections = self.rooms.get(room, [])
for ws in connections:
if ws != exclude:
try:
await ws.send_json(message)
except Exception:
pass
@property
def connection_count(self) -> int:
return len(self.active_connections)
# Global instance
manager = ConnectionManager()
3. Real-time Chat Room
# app/api/v1/chat.py
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
from datetime import datetime
from app.core.websocket import manager
router = APIRouter(tags=["Chat"])
@router.websocket("/ws/chat/{room_id}")
async def chat_websocket(
websocket: WebSocket,
room_id: str,
username: str = Query(...),
):
"""WebSocket endpoint cho chat room."""
await manager.connect(websocket, room=room_id)
# Thông báo user joined
await manager.broadcast_to_room(
{
"type": "system",
"message": f"{username} joined the room",
"timestamp": datetime.utcnow().isoformat(),
},
room=room_id,
exclude=websocket,
)
try:
while True:
# Nhận message
data = await websocket.receive_json()
# Broadcast to room
await manager.broadcast_to_room(
{
"type": "message",
"username": username,
"message": data.get("message", ""),
"timestamp": datetime.utcnow().isoformat(),
},
room=room_id,
)
except WebSocketDisconnect:
manager.disconnect(websocket, room=room_id)
await manager.broadcast_to_room(
{
"type": "system",
"message": f"{username} left the room",
"timestamp": datetime.utcnow().isoformat(),
},
room=room_id,
)
4. Real-time Notifications
# app/api/v1/notifications.py
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from app.core.auth import get_current_user_ws
from app.core.websocket import manager
from app.models.user import User
router = APIRouter(tags=["Notifications"])
@router.websocket("/ws/notifications")
async def notifications_websocket(
websocket: WebSocket,
token: str, # JWT token qua query parameter
):
"""WebSocket cho real-time notifications."""
# Authenticate WebSocket connection
from app.core.security import decode_token
payload = decode_token(token)
if not payload:
await websocket.close(code=4001, reason="Unauthorized")
return
user_id = int(payload["sub"])
await manager.connect(websocket, user_id=user_id)
try:
while True:
# Keep connection alive, nhận ping/pong
data = await websocket.receive_json()
if data.get("type") == "ping":
await websocket.send_json({"type": "pong"})
except WebSocketDisconnect:
manager.disconnect(websocket, user_id=user_id)
# Service gửi notification
async def send_notification(user_id: int, notification: dict):
"""Gửi notification đến user qua WebSocket."""
await manager.send_to_user(
{
"type": "notification",
**notification,
},
user_id=user_id,
)
5. Server-Sent Events (SSE)
# app/api/v1/events.py
import asyncio
from datetime import datetime
from fastapi import APIRouter, Request
from fastapi.responses import StreamingResponse
router = APIRouter(tags=["Events"])
async def event_generator(request: Request):
"""Generate Server-Sent Events."""
while True:
# Check nếu client disconnect
if await request.is_disconnected():
break
# Gửi event
data = {
"timestamp": datetime.utcnow().isoformat(),
"connections": 42, # Example data
}
yield f"data: {data}\n\n"
await asyncio.sleep(1) # Gửi mỗi giây
@router.get("/events/stream")
async def stream_events(request: Request):
"""SSE endpoint - alternative cho WebSocket."""
return StreamingResponse(
event_generator(request),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no", # Disable nginx buffering
},
)
6. WebSocket vs SSE vs Long Polling
| Features | WebSockets | SSE | Long Polling |
|---|---|---|---|
| Direction | Bi-directional | Server → Client | Request-Response |
| Protocol | WS/WSS | HTTP | HTTP |
| Reconnect | Manual | Auto (browser) | Manual |
| Binary data | Yes | No (text only) | Yes |
| Proxy support | Sometimes issues | Full | Full |
| Use cases | Chat, Gaming | Notifications, Feeds | Simple updates |
Summary
In this article, we have implemented:
- WebSocket endpoints: Bi-directional real-time communication
- Connection Manager: Manage connections, rooms, users
- Chat Room: Real-time chat with rooms pattern
- Notifications: Push notifications via WebSocket
- SSE: Server-Sent Events for one-way streaming
The next article will learn about Background Tasks and Celery.