228 lines
9.4 KiB
Python
228 lines
9.4 KiB
Python
"""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()}:
|
|
if channel == "vessels":
|
|
self.vessel_subscriptions.pop(websocket, None)
|
|
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]:
|
|
global_scope = config.get("scope") == "global"
|
|
bbox = [-180, -90, 180, 90] if global_scope else 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 not global_scope and (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 {
|
|
"scope": "global" if global_scope else "viewport",
|
|
"bbox": (lon_min, lat_min, lon_max, lat_max),
|
|
"zoom": zoom,
|
|
"limit": limit,
|
|
"type": sorted(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)
|
|
]
|
|
if subscription.get("scope") != "global":
|
|
matched = matched[:subscription["limit"]]
|
|
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": [],
|
|
"subscription": {
|
|
"bbox": list(subscription["bbox"]),
|
|
"zoom": subscription["zoom"],
|
|
"limit": subscription["limit"],
|
|
"scope": subscription.get("scope", "viewport"),
|
|
},
|
|
},
|
|
}
|
|
try:
|
|
for offset in range(0, len(matched), MAX_VESSEL_WS_MESSAGE_ITEMS):
|
|
await connection.send_json({
|
|
**message,
|
|
"payload": {
|
|
**message["payload"],
|
|
"vessels": matched[offset:offset + MAX_VESSEL_WS_MESSAGE_ITEMS],
|
|
},
|
|
})
|
|
except Exception:
|
|
self.unsubscribe_all(connection)
|
|
|
|
def _vessel_matches_subscription(
|
|
self,
|
|
vessel: dict[str, Any],
|
|
subscription: dict[str, Any],
|
|
) -> bool:
|
|
if vessel.get("action") == "remove" and subscription.get("scope") == "global":
|
|
return True
|
|
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
|