Yei.J. (AskaEth)
85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
from server.telemetry.listener import telemetry_listener
|
|
from server.telemetry.data import TelemetryData
|
|
from utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self._connections: dict[str, WebSocket] = {}
|
|
self._counter = 0
|
|
self._broadcast_task: asyncio.Task | None = None
|
|
self._pending_data: TelemetryData | None = None
|
|
|
|
async def connect(self, ws: WebSocket) -> str:
|
|
await ws.accept()
|
|
self._counter += 1
|
|
cid = f"client_{self._counter}"
|
|
self._connections[cid] = ws
|
|
logger.info("WS client connected: %s (total: %d)", cid, len(self._connections))
|
|
if not self._broadcast_task or self._broadcast_task.done():
|
|
self._broadcast_task = asyncio.create_task(self._broadcast_loop())
|
|
return cid
|
|
|
|
def disconnect(self, cid: str):
|
|
self._connections.pop(cid, None)
|
|
logger.info("WS client disconnected: %s (total: %d)", cid, len(self._connections))
|
|
if not self._connections and self._broadcast_task:
|
|
self._broadcast_task.cancel()
|
|
self._broadcast_task = None
|
|
|
|
def push_telemetry(self, data: TelemetryData):
|
|
self._pending_data = data
|
|
|
|
async def broadcast(self, message: dict[str, Any]):
|
|
dead = []
|
|
for cid, ws in self._connections.items():
|
|
try:
|
|
await ws.send_json(message)
|
|
except Exception:
|
|
dead.append(cid)
|
|
for cid in dead:
|
|
self.disconnect(cid)
|
|
|
|
async def _broadcast_loop(self):
|
|
last_sent = 0.0
|
|
throttle_interval = 1.0 / 30.0
|
|
try:
|
|
while self._connections:
|
|
now = time.time()
|
|
if now - last_sent >= throttle_interval and self._pending_data:
|
|
data = self._pending_data
|
|
self._pending_data = None
|
|
last_sent = now
|
|
await self.broadcast({
|
|
"type": "telemetry",
|
|
"data": data.to_dict(),
|
|
})
|
|
await asyncio.sleep(0.01)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
@property
|
|
def client_count(self) -> int:
|
|
return len(self._connections)
|
|
|
|
|
|
ws_manager = ConnectionManager()
|
|
|
|
|
|
def on_telemetry(data: TelemetryData):
|
|
ws_manager.push_telemetry(data)
|
|
|
|
|
|
telemetry_listener.on_data(on_telemetry)
|