File size: 1,353 Bytes
201b13c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import asyncio
from typing import Any
from fastapi import WebSocket
from core.logger import logger
class ConnectionManager:
def __init__(self):
self._connections: set[WebSocket] = set()
self._lock = asyncio.Lock()
async def connect(self, websocket: WebSocket) -> None:
await websocket.accept()
async with self._lock:
self._connections.add(websocket)
async def disconnect(self, websocket: WebSocket) -> None:
async with self._lock:
self._connections.discard(websocket)
async def send_json(self, websocket: WebSocket, payload: dict[str, Any]) -> None:
await websocket.send_json(payload)
async def broadcast(self, payload: dict[str, Any]) -> None:
async with self._lock:
connections = list(self._connections)
stale_connections: list[WebSocket] = []
for connection in connections:
try:
await connection.send_json(payload)
except Exception as exc:
logger.warning(f"WebSocket broadcast failed: {exc}")
stale_connections.append(connection)
if stale_connections:
async with self._lock:
for connection in stale_connections:
self._connections.discard(connection)
manager = ConnectionManager()
|