Spaces:
Running
Running
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import logging | |
| from contextlib import suppress | |
| from typing import Any | |
| import redis.asyncio as redis | |
| from fastapi import WebSocket | |
| from redis.exceptions import RedisError | |
| from app.core.config import settings | |
| logger = logging.getLogger(__name__) | |
| class ConnectionManager: | |
| """Project-scoped WebSocket fan-out backed by a resilient Redis listener.""" | |
| def __init__(self) -> None: | |
| self.active_connections: dict[str, set[WebSocket]] = {} | |
| self._pubsub_task: asyncio.Task[None] | None = None | |
| async def connect( | |
| self, | |
| websocket: WebSocket, | |
| project_id: str, | |
| *, | |
| subprotocol: str | None = None, | |
| ) -> None: | |
| if subprotocol: | |
| await websocket.accept(subprotocol=subprotocol) | |
| else: | |
| await websocket.accept() | |
| self.active_connections.setdefault(project_id, set()).add(websocket) | |
| self._ensure_listener() | |
| logger.info( | |
| "WebSocket connected project_id=%s project_connections=%d total_connections=%d", | |
| project_id, | |
| len(self.active_connections[project_id]), | |
| self.connection_count, | |
| ) | |
| async def disconnect(self, websocket: WebSocket, project_id: str) -> None: | |
| connections = self.active_connections.get(project_id) | |
| if connections is not None: | |
| connections.discard(websocket) | |
| if not connections: | |
| self.active_connections.pop(project_id, None) | |
| logger.info( | |
| "WebSocket disconnected project_id=%s total_connections=%d", | |
| project_id, | |
| self.connection_count, | |
| ) | |
| if not self.active_connections: | |
| await self._stop_listener() | |
| # A new connect may arrive while the cancelled listener is unwinding. | |
| # Re-check after the await so no live socket is left without a relay. | |
| if self.active_connections: | |
| self._ensure_listener() | |
| def connection_count(self) -> int: | |
| return sum(len(connections) for connections in self.active_connections.values()) | |
| def _ensure_listener(self) -> None: | |
| if self._pubsub_task is None or self._pubsub_task.done(): | |
| self._pubsub_task = asyncio.create_task( | |
| self._listen_to_redis(), name="backend-project-events" | |
| ) | |
| async def _stop_listener(self) -> None: | |
| task = self._pubsub_task | |
| self._pubsub_task = None | |
| if task is None or task.done() or task is asyncio.current_task(): | |
| return | |
| task.cancel() | |
| with suppress(asyncio.CancelledError): | |
| await task | |
| async def _listen_to_redis(self) -> None: | |
| reconnect_delay = 0.25 | |
| try: | |
| while self.active_connections: | |
| client: redis.Redis | None = None | |
| pubsub: redis.client.PubSub | None = None | |
| try: | |
| client = redis.from_url( | |
| settings.redis_url, | |
| decode_responses=True, | |
| max_connections=settings.redis_max_connections, | |
| health_check_interval=30, | |
| ) | |
| pubsub = client.pubsub() | |
| await pubsub.subscribe(f"{settings.redis_prefix}:events") | |
| reconnect_delay = 0.25 | |
| logger.info("WebSocket Redis event listener subscribed") | |
| async for message in pubsub.listen(): | |
| if not self.active_connections: | |
| break | |
| if message.get("type") != "message": | |
| continue | |
| try: | |
| payload = json.loads(message["data"]) | |
| project_id = payload.get("project_id") | |
| if isinstance(project_id, str): | |
| await self.broadcast(project_id, payload) | |
| if not self.active_connections: | |
| break | |
| except (TypeError, ValueError, json.JSONDecodeError): | |
| logger.exception("Unable to decode project event") | |
| except asyncio.CancelledError: | |
| raise | |
| except Exception: # noqa: BLE001 | |
| logger.exception( | |
| "WebSocket Redis listener disconnected; retrying in %.2fs", | |
| reconnect_delay, | |
| ) | |
| if self.active_connections: | |
| await asyncio.sleep(reconnect_delay) | |
| reconnect_delay = min( | |
| reconnect_delay * 2, | |
| settings.websocket_redis_reconnect_max_seconds, | |
| ) | |
| finally: | |
| if pubsub is not None: | |
| with suppress(RedisError, OSError): | |
| await pubsub.unsubscribe() | |
| with suppress(RedisError, OSError): | |
| await pubsub.aclose() | |
| if client is not None: | |
| with suppress(RedisError, OSError): | |
| await client.aclose() | |
| finally: | |
| if self._pubsub_task is asyncio.current_task(): | |
| self._pubsub_task = None | |
| logger.info("WebSocket Redis event listener stopped") | |
| async def broadcast(self, project_id: str, message: Any) -> None: | |
| connections = tuple(self.active_connections.get(project_id, set())) | |
| if not connections: | |
| return | |
| async def send(connection: WebSocket) -> WebSocket | None: | |
| try: | |
| await asyncio.wait_for(connection.send_json(message), timeout=5.0) | |
| return None | |
| except Exception: # noqa: BLE001 | |
| logger.warning("Dropping an unresponsive WebSocket project_id=%s", project_id) | |
| return connection | |
| dead = [item for item in await asyncio.gather(*(send(item) for item in connections)) if item] | |
| for connection in dead: | |
| current = self.active_connections.get(project_id) | |
| if current is not None: | |
| current.discard(connection) | |
| if not current: | |
| self.active_connections.pop(project_id, None) | |
| async def shutdown(self) -> None: | |
| connections = [ | |
| connection | |
| for project_connections in self.active_connections.values() | |
| for connection in project_connections | |
| ] | |
| self.active_connections.clear() | |
| await self._stop_listener() | |
| if connections: | |
| await asyncio.gather( | |
| *(connection.close(code=1001, reason="Backend shutting down") for connection in connections), | |
| return_exceptions=True, | |
| ) | |
| manager = ConnectionManager() | |