Init

led_platform/services/ws_hub.py

@@ -0,0 +1,50 @@
+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()}
+