61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
"""WebSocket connection manager for broadcasting live audio data."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from starlette.websockets import WebSocket, WebSocketState
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections and broadcasts messages to all clients."""
|
|
|
|
def __init__(self) -> None:
|
|
self._clients: set[WebSocket] = set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def connect(self, ws: WebSocket) -> None:
|
|
"""Accept and register new WebSocket connection."""
|
|
await ws.accept()
|
|
async with self._lock:
|
|
self._clients.add(ws)
|
|
logger.info("WS client connected (total=%d)", len(self._clients))
|
|
|
|
async def disconnect(self, ws: WebSocket) -> None:
|
|
"""Remove WebSocket connection."""
|
|
async with self._lock:
|
|
self._clients.discard(ws)
|
|
logger.info("WS client disconnected (total=%d)", len(self._clients))
|
|
|
|
async def broadcast_json(self, message: dict[str, Any]) -> None:
|
|
"""Broadcast JSON message to all connected clients."""
|
|
if not self._clients:
|
|
return
|
|
|
|
payload = json.dumps(message, ensure_ascii=False, separators=(",", ":"))
|
|
|
|
async with self._lock:
|
|
clients = list(self._clients)
|
|
|
|
dead: list[WebSocket] = []
|
|
for ws in clients:
|
|
try:
|
|
if ws.client_state == WebSocketState.CONNECTED:
|
|
await ws.send_text(payload)
|
|
else:
|
|
dead.append(ws)
|
|
except Exception:
|
|
dead.append(ws)
|
|
|
|
if dead:
|
|
async with self._lock:
|
|
for ws in dead:
|
|
self._clients.discard(ws)
|
|
logger.debug("WS cleanup: removed %d dead clients", len(dead))
|