| """Realtime fan-out over WebSocket. |
| |
| A single in-process connection manager broadcasts events/alerts to every |
| connected browser. Run the API with a single worker so all sockets share this |
| manager; for multi-worker scale put Redis pub/sub in front (see docs). |
| """ |
| 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) |
|
|
| @property |
| def count(self) -> int: |
| return len(self._connections) |
|
|
| async def broadcast(self, message: dict[str, Any]) -> None: |
| if not self._connections: |
| return |
| async with self._lock: |
| targets = list(self._connections) |
| dead: list[WebSocket] = [] |
| for ws in targets: |
| try: |
| await ws.send_json(message) |
| except Exception: |
| dead.append(ws) |
| if dead: |
| async with self._lock: |
| for ws in dead: |
| self._connections.discard(ws) |
|
|
|
|
| manager = ConnectionManager() |
|
|
|
|
| class StreamController: |
| """Background task that emits mock logs at an interval until stopped.""" |
|
|
| def __init__(self) -> None: |
| self._task: asyncio.Task | None = None |
| self.running = False |
| self.rate = 3 |
| self.interval = 1.0 |
| self.profile = "mixed" |
|
|
| def is_running(self) -> bool: |
| return self.running and self._task is not None and not self._task.done() |
|
|
| async def start(self, emit_coro, rate: int = 3, interval: float = 1.0, profile: str = "mixed") -> None: |
| if self.is_running(): |
| return |
| self.rate, self.interval, self.profile = rate, interval, profile |
| self.running = True |
| self._task = asyncio.create_task(self._run(emit_coro)) |
|
|
| async def _run(self, emit_coro) -> None: |
| try: |
| while self.running: |
| await emit_coro(self.rate, self.profile) |
| await asyncio.sleep(self.interval) |
| except asyncio.CancelledError: |
| pass |
| finally: |
| self.running = False |
|
|
| async def stop(self) -> None: |
| self.running = False |
| if self._task: |
| self._task.cancel() |
| try: |
| await self._task |
| except asyncio.CancelledError: |
| pass |
| self._task = None |
|
|
|
|
| stream = StreamController() |
|
|
