【python】FastAPI 实时推送：从 SSE 到 WebSocket

前言

构建量化策略系统时，我们需要后端向前端实时推送任务执行结果（如数据更新完成、策略信号生成等）。起初选择了 SSE (Server-Sent Events) 作为推送方案，却在开发环境中屡次遭遇 uvicorn 重载/关闭时的 CancelledError 异常，尝试了多种优化仍无法根治。最终改用 FastAPI 原生 WebSocket，问题迎刃而解。本文记录了这一过程的探索与总结。

项目背景

后端采用 FastAPI，前端 Vue 3，需要一种轻量级的服务端推送机制，用于通知前端后台任务（如数据采集、策略预处理）的完成状态。初期考虑 SSE 因为其单向推送、基于 HTTP、实现简单，且前端 EventSource 接口方便。但上线开发后，却陷入了与 CancelledError 的苦战。

SSE 的实现与问题

1.基于 asyncio.Queue 的 SSE 管理

我们实现了一个 SSEManager，内部维护一个字典 Dict $str, Set\[asyncio.Queue$ ]，每个事件类型对应一组订阅者队列。前端通过 EventSource 连接到 /sse/subscribe/{event_type}，后端返回 StreamingResponse（或后来使用的 EventSourceResponse），生成器不断从队列中获取消息并 yield。当任务完成时，调用 send_event 将消息放入所有订阅者的队列。

python 复制代码

# sse_manager.py (简化)
class SSEManager:
    def __init__(self):
        self._clients: Dict[str, Set[asyncio.Queue]] = {}
    async def subscribe(self, event_type: str) -> asyncio.Queue:
        queue = asyncio.Queue()
        self._clients.setdefault(event_type, set()).add(queue)
        return queue
    async def send_event(self, event_type: str, data: dict):
        msg = f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
        for q in self._clients.get(event_type, []):
            await q.put(msg)

路由：

python 复制代码

@router.get("/subscribe/{event_type}")
async def subscribe(event_type: str):
    queue = await sse_manager.subscribe(event_type)
    async def event_generator():
        try:
            while True:
                message = await queue.get()
                if message is None: break
                yield message
        except asyncio.CancelledError:
            pass
        finally:
            await sse_manager.unsubscribe(event_type, queue)
    return EventSourceResponse(event_generator())

2. 令人头疼的关闭错误

在开发过程中，每次按 Ctrl+C 停止 uvicorn 或修改代码触发自动重载时，控制台总会抛出大量 CancelledError，并伴有 timeout graceful shutdown exceeded 的错误：

bash 复制代码

ERROR:    Cancel 1 running task(s), timeout graceful shutdown exceeded
asyncio.exceptions.CancelledError

尽管功能上不影响前端使用（前端有自动重连），但开发体验极差，而且担心生产环境存在资源泄漏隐患。

3. 原因分析

经过调试和查阅 uvicorn 源码，我们逐渐厘清了根源：

uvicorn 在关闭时会先等待活跃连接关闭（--timeout-graceful-shutdown），然后才执行应用的 shutdown 事件（即 lifespan 的 yield 之后部分）。
SSE 底层是长连接 HTTP 响应，uvicorn 在等待超时后强制取消了生成器任务，导致 CancelledError。
因为生命周期顺序问题，我们无论在 lifespan 中如何提前发送关闭信号，都来不及在强制取消之前优雅关闭连接。

尝试过的优化（均未彻底解决）

延长超时时间 ：将 --timeout-graceful-shutdown 设大（如 30s），希望 lifespan

中的清理能赶上，但发现 lifespan 的清理总是延迟到强制取消之后。
信号处理器：在 lifespan 中注册 SIGINT/SIGTERM 信号处理器，试图在收到信号后立即清理，但在 Windows上部分信号不支持，且仍无法保证在 uvicorn 开始关闭连接前执行完毕。
使用 EventSourceResponse：替换手写的 StreamingResponse，只是对报文做了封装，并不能修正生命周期问题。
心跳与优雅退出：向队列发送 None 信号，生成器正常退出，但依然错过时机。

最终，我们意识到 SSE 在 FastAPI 中的关闭行为受限于 uvicorn 的架构，应用层很难完美控制。

转向 WebSocket：柳暗花明

既然 SSE 长连接存在协议层面的关闭劣势，我们决定试试 WebSocket ------ FastAPI 对 WebSocket 的支持非常成熟，而且 WebSocket 协议本身就有明确的关闭握手。

实现 WebSocket 管理器

使用原生 WebSocket 替代 asyncio.Queue 和 StreamingResponse，管理器变得更简洁：

python 复制代码

# ws_manager.py
class WSManager:
    def __init__(self):
        self.connections: dict[str, dict[str, WebSocket]] = {}  # event_type → {conn_id: ws}
    async def connect(self, event_type: str, ws: WebSocket) -> str:
        await ws.accept()
        conn_id = uuid.uuid4().hex[:8]
        self.connections.setdefault(event_type, {})[conn_id] = ws
        return conn_id
    def disconnect(self, event_type: str, conn_id: str):
        ...
    async def broadcast(self, event_type: str, data: dict):
        msg = {"event": event_type, "data": data}
        for ws in self.connections.get(event_type, {}).values():
            await ws.send_json(msg)

路由：

python 复制代码

@router.websocket("/subscribe/{event_type}")
async def websocket_subscribe(websocket: WebSocket, event_type: str):
    conn_id = await ws_manager.connect(event_type, websocket)
    try:
        while True:
            await websocket.receive_text()   # 保持连接
    except WebSocketDisconnect:
        pass
    finally:
        ws_manager.disconnect(event_type, conn_id)

前端改用 WebSocket 替代 EventSource，增加自动重连逻辑。

测试效果

关闭 uvicorn 后，日志输出清爽无比

python 复制代码

INFO:     Shutting down
WebSocket 断开: event=task_completed, id=2ff6ce24
INFO:     connection closed
INFO:     Application shutdown complete.

无需任何额外清理代码，WebSocket 连接自动正常关闭，再也没有 CancelledError 骚扰。

为何 WebSocket 能优雅自动关闭？

协议层面有关闭握手：WebSocket 关闭时，客户端和服务器会交换关闭帧，确保双方都正常关闭。
uvicorn/Starlette 原生支持：FastAPI 对 WebSocket 的生命周期管理非常完善，关闭时会主动发送关闭帧并等待客户端响应，然后才释放资源。
异常处理清晰：WebSocketDisconnect 异常能准确捕获连接断开，方便执行清理逻辑。

而 SSE 只是单向流，没有关闭信号，底层只能靠断开 TCP 连接来终止，在 uvicorn 中表现为强制取消任务，导致错误。

总结与建议

特性	SSE	WebSocket
协议基础	单向 HTTP 流	全双工协议
浏览器支持	EventSource 简单	WebSocket 标准
FastAPI 集成	StreamingResponse / EventSourceResponse	WebSocket 原生支持
关闭行为	依赖 uvicorn 强制取消，易出错	正常握手关闭，无异常
适用场景	只需服务端推送，客户端不发送数据	需要双向通信，或对关闭稳定性要求高

建议：如果你的场景是纯服务端推送，且能接受开发时偶尔的错误日志，SSE 仍可选用；但如果追求稳定、干净的关闭体验，或需要双向通信，WebSocket 是更优解。尤其在 FastAPI 中，WebSocket 的集成几乎零成本，值得优先考虑。当然也可以通过临时建立SSE，完成推送后自动断开的方式来避免SSE长连接带来的关闭错误。