38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket
|
|
|
|
|
|
class WebSocketManager:
|
|
def __init__(self) -> None:
|
|
self._clients: set[WebSocket] = set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def connect(self, websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
async with self._lock:
|
|
self._clients.add(websocket)
|
|
|
|
async def disconnect(self, websocket: WebSocket) -> None:
|
|
async with self._lock:
|
|
self._clients.discard(websocket)
|
|
|
|
async def broadcast(self, event_type: str, payload: dict[str, Any] | None = None) -> None:
|
|
message = {"type": event_type, "payload": payload or {}}
|
|
async with self._lock:
|
|
clients = list(self._clients)
|
|
stale: list[WebSocket] = []
|
|
for client in clients:
|
|
try:
|
|
await client.send_json(message)
|
|
except Exception:
|
|
stale.append(client)
|
|
if stale:
|
|
async with self._lock:
|
|
for client in stale:
|
|
self._clients.discard(client)
|
|
|