51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
import asyncio
|
|
from collections import defaultdict
|
|
|
|
from fastapi import WebSocket
|
|
|
|
from led_platform.domain import CommandEnvelope
|
|
|
|
|
|
class WebSocketHub:
|
|
def __init__(self) -> None:
|
|
self._lock = asyncio.Lock()
|
|
self._groups: dict[str, set[WebSocket]] = defaultdict(set)
|
|
|
|
async def connect(self, group: str, websocket: WebSocket) -> None:
|
|
await websocket.accept()
|
|
async with self._lock:
|
|
self._groups[group].add(websocket)
|
|
|
|
async def disconnect(self, group: str, websocket: WebSocket) -> None:
|
|
async with self._lock:
|
|
if group in self._groups:
|
|
self._groups[group].discard(websocket)
|
|
if not self._groups[group]:
|
|
self._groups.pop(group, None)
|
|
|
|
async def send(self, websocket: WebSocket, message: CommandEnvelope) -> None:
|
|
await websocket.send_json(message.model_dump(mode="json"))
|
|
|
|
async def broadcast(self, message: CommandEnvelope, *groups: str) -> None:
|
|
async with self._lock:
|
|
targets: list[tuple[str, WebSocket]] = []
|
|
for group in groups:
|
|
targets.extend((group, websocket) for websocket in self._groups.get(group, set()))
|
|
|
|
stale: list[tuple[str, WebSocket]] = []
|
|
for group, websocket in targets:
|
|
try:
|
|
await self.send(websocket, message)
|
|
except RuntimeError:
|
|
stale.append((group, websocket))
|
|
|
|
if stale:
|
|
async with self._lock:
|
|
for group, websocket in stale:
|
|
self._groups[group].discard(websocket)
|
|
|
|
async def counts(self) -> dict[str, int]:
|
|
async with self._lock:
|
|
return {group: len(websockets) for group, websockets in self._groups.items()}
|
|
|