default
11281e5be2
- Multi API Key + permission system (read/write/admin) with SHA-256 hash - WebSocket real-time push (location, alarm, device_status, attendance, bluetooth) - Geocoding proxy endpoints for Amap POI search and reverse geocode - Beacon modal map-based location picker with search and click-to-select - GCJ-02 ↔ WGS-84 bidirectional coordinate conversion - Data cleanup scheduler (configurable retention days) - Fix GPS longitude sign inversion (course_status bit 11: 0=East, 1=West) - Fix 2G CellID 2→3 bytes across all protocols (0x28, 0x2C, parser.py) - Fix parser loop guards, alarm_source field length, CommandLog.sent_at - Fix geocoding IMEI parameterization, require_admin import - Improve API error messages for 422 validation errors - Remove beacon floor/area fields (consolidated into name) via [HAPI](https://hapi.run) Co-Authored-By: HAPI <noreply@hapi.run>
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
"""
|
|
WebSocket Manager - WebSocket 连接管理器
|
|
Manages client connections, topic subscriptions, and broadcasting.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import WebSocket
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Maximum concurrent WebSocket connections
|
|
MAX_CONNECTIONS = 100
|
|
|
|
# Valid topics
|
|
VALID_TOPICS = {"location", "alarm", "device_status", "attendance", "bluetooth"}
|
|
|
|
|
|
class WebSocketManager:
|
|
"""Manages WebSocket connections with topic-based subscriptions."""
|
|
|
|
def __init__(self):
|
|
# {websocket: set_of_topics}
|
|
self.active_connections: dict[WebSocket, set[str]] = {}
|
|
|
|
@property
|
|
def connection_count(self) -> int:
|
|
return len(self.active_connections)
|
|
|
|
async def connect(self, websocket: WebSocket, topics: set[str]) -> bool:
|
|
"""Accept and register a WebSocket connection. Returns False if limit reached."""
|
|
if self.connection_count >= MAX_CONNECTIONS:
|
|
await websocket.close(code=1013, reason="Max connections reached")
|
|
return False
|
|
|
|
await websocket.accept()
|
|
filtered = topics & VALID_TOPICS
|
|
self.active_connections[websocket] = filtered if filtered else VALID_TOPICS
|
|
logger.info(
|
|
"WebSocket connected (%d total), topics: %s",
|
|
self.connection_count,
|
|
self.active_connections[websocket],
|
|
)
|
|
return True
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
"""Remove a WebSocket connection."""
|
|
self.active_connections.pop(websocket, None)
|
|
logger.info("WebSocket disconnected (%d remaining)", self.connection_count)
|
|
|
|
async def broadcast(self, topic: str, data: dict):
|
|
"""Broadcast a message to all subscribers of the given topic."""
|
|
if topic not in VALID_TOPICS:
|
|
return
|
|
|
|
message = json.dumps(
|
|
{"topic": topic, "data": data, "timestamp": datetime.now(timezone.utc).isoformat()},
|
|
default=str,
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
disconnected = []
|
|
# Snapshot dict to avoid RuntimeError from concurrent modification
|
|
for ws, topics in list(self.active_connections.items()):
|
|
if topic in topics:
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
disconnected.append(ws)
|
|
|
|
for ws in disconnected:
|
|
self.active_connections.pop(ws, None)
|
|
|
|
def broadcast_nonblocking(self, topic: str, data: dict):
|
|
"""Fire-and-forget broadcast (used from TCP handler context)."""
|
|
asyncio.create_task(self._safe_broadcast(topic, data))
|
|
|
|
async def _safe_broadcast(self, topic: str, data: dict):
|
|
try:
|
|
await self.broadcast(topic, data)
|
|
except Exception:
|
|
logger.exception("WebSocket broadcast error for topic %s", topic)
|
|
|
|
|
|
# Singleton instance
|
|
ws_manager = WebSocketManager()
|