"""WebSocket Connection Manager""" from datetime import UTC, datetime from typing import Any, Dict, Set, Optional from fastapi import WebSocket import redis.asyncio as redis from app.core.config import settings MAX_VESSEL_SUBSCRIPTION_LIMIT = 5000 MAX_VESSEL_WS_MESSAGE_ITEMS = 1000 MAX_VESSEL_BBOX_AREA = 2500.0 class ConnectionManager: """Manages WebSocket connections""" def __init__(self): self.active_connections: Dict[str, Set[WebSocket]] = {} # user_id -> connections self.channel_subscriptions: Dict[str, Set[WebSocket]] = {} self.websocket_channels: Dict[WebSocket, Set[str]] = {} self.vessel_subscriptions: Dict[WebSocket, dict[str, Any]] = {} self.redis_client: Optional[redis.Redis] = None async def connect(self, websocket: WebSocket, user_id: str): await websocket.accept() if user_id not in self.active_connections: self.active_connections[user_id] = set() self.active_connections[user_id].add(websocket) if self.redis_client is None: redis_url = settings.REDIS_URL if redis_url.startswith("redis://"): self.redis_client = redis.from_url(redis_url, decode_responses=True) else: self.redis_client = redis.Redis( host=settings.REDIS_SERVER, port=settings.REDIS_PORT, db=settings.REDIS_DB, decode_responses=True, ) def disconnect(self, websocket: WebSocket, user_id: str): if user_id in self.active_connections: self.active_connections[user_id].discard(websocket) if not self.active_connections[user_id]: del self.active_connections[user_id] self.unsubscribe_all(websocket) def subscribe(self, websocket: WebSocket, channels: list[str]): normalized_channels = { str(channel).strip() for channel in channels if str(channel).strip() } if not normalized_channels: return socket_channels = self.websocket_channels.setdefault(websocket, set()) for channel in normalized_channels: self.channel_subscriptions.setdefault(channel, set()).add(websocket) socket_channels.add(channel) def unsubscribe(self, websocket: WebSocket, channels: list[str]): for channel in {str(channel).strip() for channel in channels if str(channel).strip()}: subscribers = self.channel_subscriptions.get(channel) if subscribers is not None: subscribers.discard(websocket) if not subscribers: del self.channel_subscriptions[channel] socket_channels = self.websocket_channels.get(websocket) if socket_channels is not None: socket_channels.discard(channel) if not socket_channels: del self.websocket_channels[websocket] def unsubscribe_all(self, websocket: WebSocket): channels = list(self.websocket_channels.get(websocket, set())) if channels: self.unsubscribe(websocket, channels) self.vessel_subscriptions.pop(websocket, None) def subscribe_vessels(self, websocket: WebSocket, config: dict[str, Any]) -> dict[str, Any]: subscription = self._normalize_vessel_subscription(config) self.channel_subscriptions.setdefault("vessels", set()).add(websocket) self.websocket_channels.setdefault(websocket, set()).add("vessels") self.vessel_subscriptions[websocket] = subscription return subscription def _normalize_vessel_subscription(self, config: dict[str, Any]) -> dict[str, Any]: bbox = config.get("bbox") if not isinstance(bbox, (list, tuple)) or len(bbox) != 4: raise ValueError("vessels subscription requires bbox=[lon_min,lat_min,lon_max,lat_max]") try: lon_min, lat_min, lon_max, lat_max = [float(value) for value in bbox] except (TypeError, ValueError) as exc: raise ValueError("bbox values must be numbers") from exc if lat_min > lat_max: lat_min, lat_max = lat_max, lat_min if lon_min > lon_max: lon_min, lon_max = lon_max, lon_min if not (-180 <= lon_min <= 180 and -180 <= lon_max <= 180): raise ValueError("bbox longitude values must be between -180 and 180") if not (-90 <= lat_min <= 90 and -90 <= lat_max <= 90): raise ValueError("bbox latitude values must be between -90 and 90") if (lon_max - lon_min) * (lat_max - lat_min) > MAX_VESSEL_BBOX_AREA: raise ValueError("bbox is too large; zoom in or request a smaller viewport") zoom = int(config.get("zoom") or 1) if zoom < 1 or zoom > 20: raise ValueError("zoom must be between 1 and 20") limit = min(max(int(config.get("limit") or 1000), 1), MAX_VESSEL_SUBSCRIPTION_LIMIT) vessel_types = { str(item).strip().lower() for item in str(config.get("type") or "").split(",") if str(item).strip() } return { "bbox": (lon_min, lat_min, lon_max, lat_max), "zoom": zoom, "limit": limit, "type": vessel_types, "last_sent_at": None, } async def send_personal_message(self, message: dict, user_id: str): if user_id in self.active_connections: for connection in self.active_connections[user_id]: try: await connection.send_json(message) except Exception: pass async def broadcast(self, message: dict, channel: str = "all"): if channel == "all": for user_id in self.active_connections: await self.send_personal_message(message, user_id) else: for connection in list(self.channel_subscriptions.get(channel, set())): try: await connection.send_json(message) except Exception: self.unsubscribe_all(connection) async def broadcast_vessels(self, data: dict[str, Any]): vessels = data.get("vessels") if isinstance(data, dict) else None if not isinstance(vessels, list) or not vessels: return for connection, subscription in list(self.vessel_subscriptions.items()): matched = [ vessel for vessel in vessels if self._vessel_matches_subscription(vessel, subscription) ][: min(subscription["limit"], MAX_VESSEL_WS_MESSAGE_ITEMS)] if not matched: continue subscription["last_sent_at"] = datetime.now(UTC) message = { "type": "data_frame", "channel": "vessels", "timestamp": subscription["last_sent_at"].isoformat(), "payload": { **data, "vessels": matched, "subscription": { "bbox": list(subscription["bbox"]), "zoom": subscription["zoom"], "limit": subscription["limit"], }, }, } try: await connection.send_json(message) except Exception: self.unsubscribe_all(connection) def _vessel_matches_subscription( self, vessel: dict[str, Any], subscription: dict[str, Any], ) -> bool: try: lon = float(vessel.get("lon")) lat = float(vessel.get("lat")) except (TypeError, ValueError): return False lon_min, lat_min, lon_max, lat_max = subscription["bbox"] if not (lon_min <= lon <= lon_max and lat_min <= lat <= lat_max): return False requested_types = subscription.get("type") or set() if not requested_types: return True type_name = str(vessel.get("vessel_type_name") or "").lower() return any(requested_type in type_name for requested_type in requested_types) async def close_all(self): for user_id in self.active_connections: for connection in self.active_connections[user_id]: await connection.close() self.active_connections.clear() self.channel_subscriptions.clear() self.websocket_channels.clear() self.vessel_subscriptions.clear() manager = ConnectionManager() async def get_websocket_manager() -> ConnectionManager: return manager