FastAPI + WebSocket 实战:从 0 构建支持 10 万在线的实时通信服务

阿里云上云采购季 2核2G云服务器99元/年 续费同价

你有没有遇到过这种情况:花了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 连接接收消息、发消息”这么简单。你需要考虑:

  1. 连接管理:谁连进来了,谁断开了,怎么追踪
  2. 房间/频道机制:不同用户只接收自己关注的消息
  3. 水平扩展:多实例部署时,消息如何跨进程广播
  4. 断线重连:客户端重连后如何恢复状态

架构图(文字版):

客户端 ←→ 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


发表评论