WebSocket 实时通信与 任务调度 应用

先说背景。前阵子给内部做了个任务平台，用户可以提交一批批处理任务，后台跑完再把结果推回前端。一开始图省事，前端直接 setInterval 每 5 秒拉一次 /api/tasks，接口压力倒还好，毕竟用的人不多，但是那个体验是真的难受——任务一多，表格里的状态就一卡一卡地跳，明明任务已经跑完了，前端还得等下一次轮询才能看见。更要命的是，任务执行过程中我还想给它加个实时日志，用轮询那基本就没法看了。所以最后还是老老实实换成了 WebSocket。

下面是做完之后的一点记录，代码是能跑的，但踩过的坑我也一并写下来，免得你跟我一样在同一个地方蹲半天。

先把 WebSocket 这套通道搭起来

后端是 FastAPI，WebSocket 支持开箱即用。我的思路很简单：维护一个 task_id -> 连接集合 的映射，任务状态一变，就往对应的连接里推。先来个最基础的连接管理：

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from collections import defaultdict

app = FastAPI()
subscribers: dict[str, set[WebSocket]] = defaultdict(set)

@app.websocket("/ws/tasks/{task_id}")
async def task_ws(ws: WebSocket, task_id: str):
    await ws.accept()
    subscribers[task_id].add(ws)
    try:
        while True:
            # 这里其实不用收消息，但得挂着，不然连接会被判定为空闲
            await ws.receive_text()
    except WebSocketDisconnect:
        subscribers[task_id].discard(ws)
        if not subscribers[task_id]:
            subscribers.pop(task_id, None)

async def broadcast(task_id: str, payload: dict):
    dead = set()
    for ws in subscribers.get(task_id, set()):
        try:
            await ws.send_json(payload)
        except Exception:
            dead.add(ws)
    for ws in dead:
        subscribers[task_id].discard(ws)

前端就一句 new WebSocket(\/ws/tasks/${taskId}`)，然后 onmessage` 里更新状态，完事。到这一步，页面上的状态刷新就已经是"推送式"的了，那种点了提交、状态立刻从 pending 跳到 running 的感觉，确实比轮询爽。

调度器这边，别想得太复杂

任务调度我没上 Celery，因为我们的任务是用户提交、后台按顺序跑，量也不大，直接 asyncio + 一个简单的队列就够了。真正需要注意的是：调度器不是"调完就完事"，它得在每一个状态节点通知外面。

import asyncio

queue: asyncio.Queue[str] = asyncio.Queue()

async def worker():
    while True:
        task_id = await queue.get()
        try:
            await broadcast(task_id, {"status": "running"})
            result = await run_task(task_id)          # 真正的业务
            await broadcast(task_id, {"status": "done", "result": result})
        except Exception as e:
            await broadcast(task_id, {"status": "failed", "error": str(e)})
        finally:
            queue.task_done()

@app.on_event("startup")
async def startup():
    for _ in range(4):                            # 起几个并发 worker
        asyncio.create_task(worker())

run_task 里如果还要推中间进度，比如处理到第几批，直接再 await broadcast(task_id, {"progress": i / total}) 就行，前端拿到就更新进度条。这块比轮询干净太多了，不用再在接口里塞什么 progress 字段让前端去猜。

真正开始踩坑

第一坑：多进程。我一开始 uvicorn main:app --workers 4 起服务，结果任务在 worker 2 里跑，用户的 WebSocket 连接落在 worker 1 上，广播当然是发不出去的，前端就干等着。这个不怪 WebSocket，是我自己没想清楚——多进程下内存里的 subscribers 根本不共享。解决办法要么换成单 worker（我们量小，够用），要么上 Redis 的 pub/sub，各进程订阅同一个频道，把 broadcast 改成往 Redis 里 publish。我最后是先用单 worker 顶着，等真有压力了再说，别一上来就过度设计。

第二坑：连接会静默死掉。Nginx 默认 60 秒没数据就断，用户切到别的标签页回来，连接早就没了，但他自己不知道。所以必须加心跳，前端定时发个 ping，后端回 pong，简单粗暴：

@app.websocket("/ws/tasks/{task_id}")
async def task_ws(ws: WebSocket, task_id: str):
    await ws.accept()
    subscribers[task_id].add(ws)
    try:
        while True:
            msg = await ws.receive_text()
            if msg == "ping":
                await ws.send_text("pong")
    except WebSocketDisconnect:
        ...

前端再配一个自动重连，onclose 里 setTimeout 重新建连，重连后重新发一次 GET /api/tasks/{id} 把当前状态补齐——别指望重连能"续上"之前的推送，该重新拉一次就重新拉。

第三坑：广播别在业务逻辑里乱 await。我一开始在 run_task 里每处理一条数据就 broadcast 一次，数据一多，全卡在 send_json 上，任务本身反而变慢了。后来是攒一批再推，或者把推送丢进另一个队列由单独的协程去发。实时归实时，别让它拖累真正干活的线程。

最后

这套东西跑下来，代码量其实不大，麻烦的地方全在那些边角：进程模型、心跳、重连、推送频率。WebSocket 本身好写，难的是你默认它"一直连着"，然后忘了网络这玩意儿随时会断。想清楚这两点，剩下的基本就是业务代码了。