46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
|
import asyncio
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
from typing import Set
|
||
|
|
|
||
|
|
from fastapi import WebSocket
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class MissionConnectionManager:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self._lock = asyncio.Lock()
|
||
|
|
self._connections: Set[WebSocket] = set()
|
||
|
|
|
||
|
|
async def connect(self, websocket: WebSocket) -> None:
|
||
|
|
await websocket.accept()
|
||
|
|
async with self._lock:
|
||
|
|
self._connections.add(websocket)
|
||
|
|
logger.info("📡 Mission WS connected (%s active)", len(self._connections))
|
||
|
|
|
||
|
|
async def disconnect(self, websocket: WebSocket) -> None:
|
||
|
|
async with self._lock:
|
||
|
|
self._connections.discard(websocket)
|
||
|
|
logger.info("📡 Mission WS disconnected (%s active)", len(self._connections))
|
||
|
|
|
||
|
|
async def broadcast(self, event: str, payload: dict) -> None:
|
||
|
|
message = json.dumps({"event": event, "data": payload}, default=str)
|
||
|
|
async with self._lock:
|
||
|
|
targets = list(self._connections)
|
||
|
|
|
||
|
|
dead: list[WebSocket] = []
|
||
|
|
for websocket in targets:
|
||
|
|
try:
|
||
|
|
await websocket.send_text(message)
|
||
|
|
except Exception:
|
||
|
|
dead.append(websocket)
|
||
|
|
|
||
|
|
if dead:
|
||
|
|
async with self._lock:
|
||
|
|
for websocket in dead:
|
||
|
|
self._connections.discard(websocket)
|
||
|
|
|
||
|
|
|
||
|
|
mission_ws_manager = MissionConnectionManager()
|