FastAPI WebSocket Broadcast Manager

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)