40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self) -> None:
|
|
self._connections: set[WebSocket] = set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def connect(self, ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
async with self._lock:
|
|
self._connections.add(ws)
|
|
|
|
async def disconnect(self, ws: WebSocket) -> None:
|
|
async with self._lock:
|
|
self._connections.discard(ws)
|
|
|
|
async def broadcast_json(self, payload: dict[str, Any]) -> None:
|
|
async with self._lock:
|
|
conns = list(self._connections)
|
|
|
|
to_remove: list[WebSocket] = []
|
|
for ws in conns:
|
|
try:
|
|
await ws.send_json(payload)
|
|
except Exception:
|
|
to_remove.append(ws)
|
|
|
|
if to_remove:
|
|
async with self._lock:
|
|
for ws in to_remove:
|
|
self._connections.discard(ws)
|
|
|
|
|
|
manager = ConnectionManager()
|