FastAPI WebSocket 广播管理器

示例说明

广播的朴素写法是遍历连接逐个 await send——一个慢客户端就拖住所有人。改进方案是每客户端一条有界队列作为发送缓冲：broadcast 只做非阻塞入队立即返回；真正的网络发送在各连接自己的事件循环回合里完成。队列满说明该消费者持续掉队，直接断开隔离，保护其他客户端的实时性。

技术要点

• 每客户端有界队列做发送缓冲
• put_nowait 非阻塞广播零等待
• 队列满即断开隔离慢消费者

代码解读

ConnectionManager 维护连接与队列两张表；relay 端点用 asyncio.wait 同时监听队列出口与客户端入站消息，先完成者先处理——既能推送广播也能接收消息，取消另一个挂起任务防止泄漏。断连异常统一走 disconnect 清理两张表。

使用建议

断开慢消费者前可先发警告帧给客户端自查。心跳帧穿插防代理层掐空闲连接。横向扩展时广播要走 Redis 发布订阅。

代码示例

import asyncio

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()


class ConnectionManager:
    def __init__(self, queue_size: int = 64):
        self.active: dict[int, WebSocket] = {}
        self.pending: dict[int, asyncio.Queue] = {}
        self.queue_size = queue_size

    async def connect(self, client_id: int, socket: WebSocket) -> None:
        await socket.accept()
        self.active[client_id] = socket
        self.pending[client_id] = asyncio.Queue(maxsize=self.queue_size)

    def disconnect(self, client_id: int) -> None:
        self.active.pop(client_id, None)
        self.pending.pop(client_id, None)

    async def broadcast(self, message: dict) -> int:
        delivered = 0
        for client_id, queue in list(self.pending.items()):
            try:
                queue.put_nowait(message)
                delivered += 1
            except asyncio.QueueFull:
                self.disconnect(client_id)  # 慢消费者被隔离
        return delivered


manager = ConnectionManager()


@app.websocket('/ws/{client_id}')
async def relay(socket: WebSocket, client_id: int):
    await manager.connect(client_id, socket)
    queue = manager.pending[client_id]
    try:
        while True:
            send_task = asyncio.create_task(queue.get())
            recv_task = asyncio.create_task(socket.receive_text())
            done, pending = await asyncio.wait(
                {send_task, recv_task}, return_when=asyncio.FIRST_COMPLETED
            )
            for task in pending:
                task.cancel()
            if recv_task in done:
                await manager.broadcast({'from': client_id, 'text': recv_task.result()})
            if send_task in done:
                await socket.send_json(send_task.result())
    except WebSocketDisconnect:
        manager.disconnect(client_id)