你有没有遇到过这种情况:花了3天搭了个 WebSocket 服务,上线后发现一超过 1000 个连接就开始丢消息?这篇文章带你用 FastAPI 从零做一个真正能扛住并发的实时通信架构,附完整可运行代码。
一、为什么选 FastAPI 做 WebSocket?
先简单对比一下 Python 主流框架的 WebSocket 支持:
| 框架 | WebSocket 支持 | 异步原生 | 性能(连接数/核) | 适合场景 |
|---|---|---|---|---|
| FastAPI | ✅ 原生支持 | ✅ | 5000+ | API + WS 混合 |
| Django Channels | ✅ 需要 ASGI 插件 | ⚠️ 部分 | 2000+ | Django 项目扩展 |
| Flask-SocketIO | ✅ 依赖 eventlet | ❌ | 1500+ | 简单场景 |
| aiohttp | ✅ 原生 | ✅ | 8000+ | 纯 WS 服务 |
| Tornado | ✅ 原生 | ✅ | 6000+ | 老项目 |
FastAPI 基于 Starlette,天然支持 ASGI,WebSocket 性能出色,又能复用你的 HTTP API 代码库——对于需要同时提供 REST API 和实时推送的项目,是最优解。
二、核心架构设计
一个生产级实时通信服务,不是”一个 WebSocket 连接接收消息、发消息”这么简单。你需要考虑:
- 连接管理:谁连进来了,谁断开了,怎么追踪
- 房间/频道机制:不同用户只接收自己关注的消息
- 水平扩展:多实例部署时,消息如何跨进程广播
- 断线重连:客户端重连后如何恢复状态
架构图(文字版):
客户端 ←→ FastAPI (WebSocket) ←→ Redis Pub/Sub ←→ FastAPI (WebSocket) ←→ 客户端
↕
消息持久化 (可选)
用 Redis Pub/Sub 作为消息总线,每个 FastAPI 实例订阅同一个频道,实现多实例间的消息广播。
三、完整代码实现
3.1 安装依赖
pip install fastapi uvicorn redis asyncio-redis
# 或者用 aioredis(asyncio 原生)
pip install aioredis
3.2 连接管理器(单机版)
# connection_manager.py
from fastapi import WebSocket
from typing import Dict, Set
import asyncio
import logging
logger = logging.getLogger(__name__)
class ConnectionManager:
"""WebSocket 连接管理器,支持房间/频道机制"""
def __init__(self):
# 全局连接池:{client_id: WebSocket}
self.active_connections: Dict[str, WebSocket] = {}
# 房间映射:{room_id: Set[client_id]}
self.rooms: Dict[str, Set[str]] = {}
async def connect(self, websocket: WebSocket, client_id: str):
"""接受新连接"""
await websocket.accept()
self.active_connections[client_id] = websocket
logger.info(f"客户端 {client_id} 已连接,当前连接数: {len(self.active_connections)}")
def disconnect(self, client_id: str):
"""断开连接,清理所有房间"""
if client_id in self.active_connections:
del self.active_connections[client_id]
# 从所有房间移除
for room_id, members in self.rooms.items():
members.discard(client_id)
logger.info(f"客户端 {client_id} 已断开,当前连接数: {len(self.active_connections)}")
async def join_room(self, client_id: str, room_id: str):
"""加入房间"""
if room_id not in self.rooms:
self.rooms[room_id] = set()
self.rooms[room_id].add(client_id)
async def leave_room(self, client_id: str, room_id: str):
"""离开房间"""
if room_id in self.rooms:
self.rooms[room_id].discard(client_id)
async def send_personal(self, message: str, client_id: str):
"""发送私信"""
ws = self.active_connections.get(client_id)
if ws:
try:
await ws.send_text(message)
except Exception as e:
logger.error(f"发送给 {client_id} 失败: {e}")
self.disconnect(client_id)
async def broadcast_to_room(self, message: str, room_id: str, exclude: str = None):
"""广播到房间内所有成员"""
if room_id not in self.rooms:
return
members = self.rooms[room_id].copy()
tasks = []
for client_id in members:
if client_id == exclude:
continue
tasks.append(self.send_personal(message, client_id))
await asyncio.gather(*tasks, return_exceptions=True)
async def broadcast_all(self, message: str):
"""广播给所有连接"""
tasks = [
self.send_personal(message, cid)
for cid in list(self.active_connections.keys())
]
await asyncio.gather(*tasks, return_exceptions=True)
def get_stats(self):
return {
"total_connections": len(self.active_connections),
"rooms": {
room_id: len(members)
for room_id, members in self.rooms.items()
}
}
manager = ConnectionManager()
3.3 Redis 跨实例广播(生产必备)
# redis_pubsub.py
import aioredis
import asyncio
import json
import logging
from typing import Optional
logger = logging.getLogger(__name__)
REDIS_URL = "redis://localhost:6379"
class RedisPubSubManager:
"""基于 Redis Pub/Sub 的跨实例消息广播"""
def __init__(self, redis_url: str = REDIS_URL):
self.redis_url = redis_url
self.pub_client: Optional[aioredis.Redis] = None
self.sub_client: Optional[aioredis.Redis] = None
self._subscriber = None
async def connect(self):
self.pub_client = await aioredis.from_url(self.redis_url, encoding="utf-8", decode_responses=True)
self.sub_client = await aioredis.from_url(self.redis_url, encoding="utf-8", decode_responses=True)
logger.info("Redis Pub/Sub 连接成功")
async def publish(self, channel: str, message: dict):
"""发布消息到 Redis 频道"""
if self.pub_client:
await self.pub_client.publish(channel, json.dumps(message, ensure_ascii=False))
async def subscribe(self, channel: str, callback):
"""订阅频道,收到消息后调用 callback"""
async with self.sub_client.pubsub() as pubsub:
await pubsub.subscribe(channel)
logger.info(f"已订阅 Redis 频道: {channel}")
async for message in pubsub.listen():
if message["type"] == "message":
try:
data = json.loads(message["data"])
await callback(data)
except Exception as e:
logger.error(f"处理 Redis 消息失败: {e}")
async def disconnect(self):
if self.pub_client:
await self.pub_client.close()
if self.sub_client:
await self.sub_client.close()
redis_manager = RedisPubSubManager()
3.4 FastAPI 主应用
# main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from fastapi.responses import HTMLResponse
import json
import uuid
import asyncio
from connection_manager import manager
from redis_pubsub import redis_manager
app = FastAPI(title="实时通信服务", version="1.0.0")
BROADCAST_CHANNEL = "global_broadcast"
@app.on_event("startup")
async def startup():
"""启动时连接 Redis 并开始监听广播频道"""
await redis_manager.connect()
# 后台任务:监听 Redis 广播,转发给本实例的 WebSocket 连接
asyncio.create_task(
redis_manager.subscribe(BROADCAST_CHANNEL, handle_redis_message)
)
@app.on_event("shutdown")
async def shutdown():
await redis_manager.disconnect()
async def handle_redis_message(data: dict):
"""收到 Redis 广播时,推送给本实例的 WebSocket 连接"""
room_id = data.get("room_id")
message = data.get("message")
sender = data.get("sender_id")
if room_id:
await manager.broadcast_to_room(
json.dumps(data, ensure_ascii=False),
room_id,
exclude=sender # 不重复发给发送者
)
else:
await manager.broadcast_all(json.dumps(data, ensure_ascii=False))
@app.websocket("/ws/{room_id}")
async def websocket_endpoint(
websocket: WebSocket,
room_id: str,
client_id: str = Query(default=None)
):
"""WebSocket 入口,支持房间机制"""
if not client_id:
client_id = str(uuid.uuid4())[:8]
await manager.connect(websocket, client_id)
await manager.join_room(client_id, room_id)
# 通知房间内其他人有新成员加入
join_msg = json.dumps({
"type": "system",
"room_id": room_id,
"sender_id": "system",
"message": f"用户 {client_id} 加入了房间",
"user_count": len(manager.rooms.get(room_id, set()))
}, ensure_ascii=False)
await manager.broadcast_to_room(join_msg, room_id, exclude=client_id)
try:
while True:
raw = await websocket.receive_text()
try:
data = json.loads(raw)
except json.JSONDecodeError:
data = {"text": raw}
payload = {
"type": "message",
"room_id": room_id,
"sender_id": client_id,
"message": data.get("message", raw),
"timestamp": asyncio.get_event_loop().time()
}
# 通过 Redis 广播(所有实例都会收到)
await redis_manager.publish(BROADCAST_CHANNEL, payload)
except WebSocketDisconnect:
manager.disconnect(client_id)
await manager.leave_room(client_id, room_id)
leave_msg = json.dumps({
"type": "system",
"room_id": room_id,
"sender_id": "system",
"message": f"用户 {client_id} 离开了房间"
}, ensure_ascii=False)
await manager.broadcast_to_room(leave_msg, room_id)
@app.get("/stats")
async def get_stats():
"""查看连接统计"""
return manager.get_stats()
@app.get("/", response_class=HTMLResponse)
async def test_page():
"""简单测试页面"""
return """
<html>
<body>
<input id="msg" placeholder="输入消息"><button onclick="send()">发送</button>
<div id="log"></div>
<script>
const ws = new WebSocket(`ws://localhost:8000/ws/room1`);
ws.onmessage = e => {
document.getElementById('log').innerHTML += '<p>' + e.data + '</p>';
};
function send() {
const msg = document.getElementById('msg').value;
ws.send(JSON.stringify({message: msg}));
}
</script>
</body>
</html>
"""
3.5 启动服务
# 单机启动
uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1
# 生产多进程(注意:多 worker 时必须依赖 Redis 广播)
uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4
四、性能调优 5 个关键点
4.1 调整 WebSocket ping/pong 心跳
# 防止连接被中间件超时断开
@app.websocket("/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str):
await websocket.accept()
# 设置 ping 间隔为 20 秒(Uvicorn 默认 20s)
# 客户端需在 30s 内响应 pong,否则断开
...
Nginx 配置也需要对应调整:
location /ws/ {
proxy_pass http://backend;
proxy_http_version 1.1;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection "upgrade";
proxy_read_timeout 3600s; # 长连接不超时
proxy_send_timeout 3600s;
}
4.2 使用 asyncio.gather 并发发送
# ✅ 并发发送,速度是串行的 N 倍
await asyncio.gather(*[ws.send_text(msg) for ws in connections])
# ❌ 串行发送(慢,不要用)
for ws in connections:
await ws.send_text(msg)
4.3 消息压缩
import zlib
import base64
def compress_message(msg: str) -> str:
"""消息超过 1KB 时压缩发送"""
if len(msg) > 1024:
compressed = zlib.compress(msg.encode())
return base64.b64encode(compressed).decode()
return msg
4.4 连接限速(防刷)
from collections import defaultdict
import time
class RateLimiter:
def __init__(self, max_messages: int = 30, window: int = 10):
self.max_messages = max_messages
self.window = window
self.records: defaultdict = defaultdict(list)
def is_allowed(self, client_id: str) -> bool:
now = time.time()
records = self.records[client_id]
# 清除窗口外的记录
self.records[client_id] = [t for t in records if now - t < self.window]
if len(self.records[client_id]) >= self.max_messages:
return False
self.records[client_id].append(now)
return True
rate_limiter = RateLimiter()
4.5 生产部署配置参考
| 配置项 | 推荐值 | 说明 |
|---|---|---|
| Uvicorn workers | CPU 核数 × 2 | 充分利用多核 |
| Redis 连接池大小 | 50~100 | 避免频繁建连 |
| WebSocket 消息大小限制 | 1MB | 防止大消息攻击 |
| 最大连接数/实例 | 10000 | 超过需水平扩展 |
| Nginx worker_connections | 65535 | 操作系统 ulimit 需同步调整 |
五、常见踩坑
坑1:多 worker 下消息只发给了部分客户端
原因:没有 Redis Pub/Sub,消息只在当前进程内广播。解决:接入 Redis 消息总线。
坑2:大量连接断开时服务卡顿
原因:disconnect() 逻辑是同步的,触发大量字典操作。解决:用 asyncio.Lock 保护连接池。
坑3:前端收不到消息,但后端日志显示已发送
原因:Nginx 缓冲区未关闭,消息被缓存。解决:proxy_buffering off;
坑4:WebSocket 连接数上限
Linux 默认文件描述符限制 1024,需修改:
# /etc/security/limits.conf
* soft nofile 1000000
* hard nofile 1000000
# 临时修改(当前会话)
ulimit -n 1000000
总结
| 层次 | 技术方案 | 作用 |
|---|---|---|
| 连接管理 | ConnectionManager + asyncio | 高效追踪/广播 |
| 跨实例广播 | Redis Pub/Sub | 水平扩展 |
| 反向代理 | Nginx | 负载均衡 + 长连接 |
| 限速 | RateLimiter | 防刷保护 |
| 监控 | /stats 接口 | 实时连接数查看 |
FastAPI WebSocket 的上限取决于你的 Redis 和 Nginx 配置,单机跑 10 万在线完全可行。下次有人问你”WebSocket 怎么做高并发”,直接把这篇甩给他。
👤 作者简介
一枚在大中原腹地(河南)卖公有云的从业者,主营腾讯云/阿里云/华为云,曾踩坑无数,现专注AI大模型应用落地。关注公众号「AI热点日报」,围观AI前沿动态~
博客:y1cloud.cn
