FastAPI WebSocket 基础指南
文章目录
什么是 WebSocket
WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议。与 HTTP 的「请求-响应」模式不同,WebSocket 连接一旦建立,客户端和服务端可以随时互相发送消息,非常适合聊天应用、实时通知、股票行情、协作编辑等场景。
FastAPI 对 WebSocket 提供了原生支持,使用起来非常简洁。
最简单的 WebSocket 端点
1 | from fastapi import FastAPI, WebSocket |
要点说明:
- 用
@app.websocket()装饰器定义 WebSocket 路由 - 必须先调用
await websocket.accept()接受连接,否则无法通信 receive_text()接收文本消息,send_text()发送文本消息
WebSocket 对象的核心方法
| 方法 | 说明 |
|---|---|
await websocket.accept() |
接受连接,可传入 headers、subprotocol 等参数 |
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 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect |
FastAPI 内置了 WebSocketDisconnect 异常,当客户端断开或网络异常时会自动抛出,无需手动判断。
连接管理:携带查询参数
WebSocket 连接建立时可以携带查询参数,常用于传递 token 或房间号:
1 |
|
访问时使用 ws://localhost:8000/ws/alice?token=abc123,FastAPI 会自动解析路径参数和查询参数。
广播消息:多客户端聊天室
实际场景中常常需要将消息推送给所有连接的客户端,需要维护一个连接池:
1 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect |
这个 ConnectionManager 封装了连接管理逻辑:连接时注册、断开时移除、广播时遍历所有连接发送消息。
在生产环境中,如果有多个 worker 进程,消息广播需要借助 Redis Pub/Sub 或消息队列来实现跨进程通信。
依赖注入与 WebSocket
WebSocket 端点同样可以使用 Depends 进行依赖注入,比如验证 token:
1 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, Query |
注意:依赖在
accept()之前执行,如果认证失败抛出异常,连接不会被接受,客户端会收到 HTTP 403 响应。
生产环境部署注意事项
部署 WebSocket 应用时需要注意以下几点:
- 反向代理配置:Nginx 需要升级 HTTP 连接为 WebSocket,添加
proxy_set_header Upgrade $http_upgrade和proxy_set_header Connection "upgrade" - 超时设置:WebSocket 是长连接,需要调整代理和负载均衡器的超时时间
- 多进程问题:使用 ASGI 服务器(如 Uvicorn)时,多 worker 下连接分布在不同进程,广播需借助 Redis 等外部工具
- 心跳检测:长时间无消息时连接可能被中间件断开,建议定期发送 ping/pong 保活
小结
FastAPI 的 WebSocket 支持非常直观:用装饰器定义端点,用 accept() 建立连接,用 receive_* 和 send_* 收发消息。配合依赖注入和连接管理器,可以轻松构建实时通信功能。核心要记住:先 accept,再收发,最后处理 WebSocketDisconnect 异常即可。