什么是 WebSocket

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议。与 HTTP 的「请求-响应」模式不同,WebSocket 连接一旦建立,客户端和服务端可以随时互相发送消息,非常适合聊天应用、实时通知、股票行情、协作编辑等场景。

FastAPI 对 WebSocket 提供了原生支持,使用起来非常简洁。

最简单的 WebSocket 端点

1
2
3
4
5
6
7
8
9
10
11
from fastapi import FastAPI, WebSocket

app = FastAPI()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
while True:
data = await websocket.receive_text()
await websocket.send_text(f"收到消息: {data}")

要点说明:

  • @app.websocket() 装饰器定义 WebSocket 路由
  • 必须先调用 await websocket.accept() 接受连接,否则无法通信
  • receive_text() 接收文本消息,send_text() 发送文本消息

WebSocket 对象的核心方法

方法 说明
await websocket.accept() 接受连接,可传入 headerssubprotocol 等参数
await websocket.receive_text() 接收文本消息
await websocket.receive_bytes() 接收二进制消息
await websocket.receive_json() 接收 JSON 并自动解析为 dict
await websocket.send_text(data) 发送文本消息
await websocket.send_bytes(data) 发送二进制消息
await websocket.send_json(data) 将 dict 序列化为 JSON 发送
await websocket.close(code) 关闭连接,code 为关闭状态码(如 1000 表示正常关闭)

处理异常断开

客户端可能随时断开连接,捕获异常是必要的:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"回复: {data}")
except WebSocketDisconnect:
print("客户端已断开连接")

FastAPI 内置了 WebSocketDisconnect 异常,当客户端断开或网络异常时会自动抛出,无需手动判断。

连接管理:携带查询参数

WebSocket 连接建立时可以携带查询参数,常用于传递 token 或房间号:

1
2
3
4
5
6
7
8
9
10
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str, token: str = ""):
# 可以在此处验证 token
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"[{client_id}] 说: {data}")
except WebSocketDisconnect:
print(f"{client_id} 已离开")

访问时使用 ws://localhost:8000/ws/alice?token=abc123,FastAPI 会自动解析路径参数和查询参数。

广播消息:多客户端聊天室

实际场景中常常需要将消息推送给所有连接的客户端,需要维护一个连接池:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List

app = FastAPI()


class ConnectionManager:
def __init__(self):
self.active_connections: List[WebSocket] = []

async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)

def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)

async def broadcast(self, message: str):
for connection in self.active_connections:
await connection.send_text(message)


manager = ConnectionManager()


@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.broadcast(f"有人发来消息: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket)
await manager.broadcast("一位用户离开了聊天室")

这个 ConnectionManager 封装了连接管理逻辑:连接时注册、断开时移除、广播时遍历所有连接发送消息。

在生产环境中,如果有多个 worker 进程,消息广播需要借助 Redis Pub/Sub 或消息队列来实现跨进程通信。

依赖注入与 WebSocket

WebSocket 端点同样可以使用 Depends 进行依赖注入,比如验证 token:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, Query


async def get_token(token: str = Query(...)):
# 此处可以查询数据库验证 token 有效性
if token != "secret":
raise Exception("认证失败")
return token


@app.websocket("/ws/secure")
async def secure_websocket(
websocket: WebSocket,
token: str = Depends(get_token),
):
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"已认证用户: {data}")
except WebSocketDisconnect:
print("断开连接")

注意:依赖在 accept() 之前执行,如果认证失败抛出异常,连接不会被接受,客户端会收到 HTTP 403 响应。

生产环境部署注意事项

部署 WebSocket 应用时需要注意以下几点:

  1. 反向代理配置:Nginx 需要升级 HTTP 连接为 WebSocket,添加 proxy_set_header Upgrade $http_upgradeproxy_set_header Connection "upgrade"
  2. 超时设置:WebSocket 是长连接,需要调整代理和负载均衡器的超时时间
  3. 多进程问题:使用 ASGI 服务器(如 Uvicorn)时,多 worker 下连接分布在不同进程,广播需借助 Redis 等外部工具
  4. 心跳检测:长时间无消息时连接可能被中间件断开,建议定期发送 ping/pong 保活

小结

FastAPI 的 WebSocket 支持非常直观:用装饰器定义端点,用 accept() 建立连接,用 receive_*send_* 收发消息。配合依赖注入和连接管理器,可以轻松构建实时通信功能。核心要记住:先 accept,再收发,最后处理 WebSocketDisconnect 异常即可。