| import asyncio |
| import logging |
| import time |
| from typing import Dict, Set, Optional |
| from fastapi import WebSocket |
| from dataclasses import dataclass, field |
|
|
| from ..config import settings |
| from .protocol import ( |
| MessageType, |
| ReadyMessage, |
| StatsMessage, |
| ) |
|
|
|
|
| @dataclass |
| class ConnectionInfo: |
| websocket: WebSocket |
| connected_at_ms: int |
| last_active_ms: int = 0 |
| chunks_received: int = 0 |
| bytes_received: int = 0 |
| client_info: dict = field(default_factory=dict) |
|
|
|
|
| class ConnectionManager: |
| """Manages WebSocket connections with per-connection isolation.""" |
|
|
| def __init__(self, max_connections: int = settings.max_connections): |
| self._connections: Dict[str, ConnectionInfo] = {} |
| self._active_ids: Set[str] = set() |
| self._max_connections = max_connections |
| self._logger = logging.getLogger(__name__) |
| self._lock = asyncio.Lock() |
|
|
| async def connect(self, websocket: WebSocket, connection_id: str, **kwargs) -> None: |
| """Accept connection and register it.""" |
| async with self._lock: |
| if len(self._connections) >= self._max_connections: |
| await websocket.close(code=1013, reason="Too many connections") |
| raise RuntimeError(f"Max connections ({self._max_connections}) reached") |
|
|
| await websocket.accept() |
| info = ConnectionInfo( |
| websocket=websocket, |
| connected_at_ms=int(time.time() * 1000), |
| client_info=kwargs, |
| ) |
| self._connections[connection_id] = info |
| self._active_ids.add(connection_id) |
| self._logger.info(f"Connection established: {connection_id}") |
|
|
| async def disconnect(self, connection_id: str) -> None: |
| """Remove connection and cleanup resources.""" |
| async with self._lock: |
| self._connections.pop(connection_id, None) |
| self._active_ids.discard(connection_id) |
| self._logger.info(f"Connection closed: {connection_id}") |
|
|
| def is_active(self, connection_id: str) -> bool: |
| return connection_id in self._connections |
|
|
| def get_connection(self, connection_id: str) -> Optional[ConnectionInfo]: |
| return self._connections.get(connection_id) |
|
|
| def touch(self, connection_id: str, n_chunks: int = 0, n_bytes: int = 0) -> None: |
| """Update connection activity timestamps and counters.""" |
| info = self._connections.get(connection_id) |
| if info: |
| info.last_active_ms = int(time.time() * 1000) |
| info.chunks_received += n_chunks |
| info.bytes_received += n_bytes |
|
|
| async def send_text(self, connection_id: str, message: str) -> bool: |
| """Send text message to a specific connection.""" |
| info = self._connections.get(connection_id) |
| if not info: |
| return False |
| try: |
| await info.websocket.send_text(message) |
| return True |
| except Exception as e: |
| self._logger.error(f"Failed to send to {connection_id}: {e}") |
| return False |
|
|
| async def send_json(self, connection_id: str, message: dict) -> bool: |
| """Send JSON message to a specific connection.""" |
| info = self._connections.get(connection_id) |
| if not info: |
| return False |
| try: |
| await info.websocket.send_json(message) |
| return True |
| except Exception as e: |
| self._logger.error(f"Failed to send JSON to {connection_id}: {e}") |
| return False |
|
|
| async def broadcast_json(self, message: dict) -> int: |
| """Broadcast message to all connections. Returns count of successes.""" |
| successes = 0 |
| for cid in list(self._connections.keys()): |
| if await self.send_json(cid, message): |
| successes += 1 |
| return successes |
|
|
| def get_active_connections(self) -> Dict[str, ConnectionInfo]: |
| return dict(self._connections) |
|
|
| def get_connection_count(self) -> int: |
| return len(self._connections) |
|
|
| def get_connection_summary(self) -> dict: |
| """Get summary of all connections for monitoring.""" |
| now = int(time.time() * 1000) |
| return { |
| cid: { |
| "connected_at_ms": info.connected_at_ms, |
| "uptime_ms": now - info.connected_at_ms, |
| "last_active_ms": info.last_active_ms, |
| "chunks_received": info.chunks_received, |
| "bytes_received": info.bytes_received, |
| } |
| for cid, info in self._connections.items() |
| } |
|
|
| async def close_all(self) -> None: |
| """Close all connections gracefully.""" |
| for cid in list(self._connections.keys()): |
| try: |
| await self._connections[cid].websocket.close(code=1001, reason="Server shutting down") |
| except Exception: |
| pass |
| self._connections.clear() |
| self._active_ids.clear() |
|
|
| async def session_cleanup(self) -> None: |
| """Periodic cleanup of inactive connections.""" |
| now = int(time.time() * 1000) |
| for cid in list(self._connections.keys()): |
| info = self._connections[cid] |
| if now - info.last_active_ms > settings.ws_ping_timeout * 1000 * 2: |
| self._logger.warning(f"Cleaning up inactive connection: {cid}") |
| await self.disconnect(cid) |
