453c18f192
Some checks failed
Backend Tests / pytest (Python 3.12) (push) Has been cancelled
Backend: - servers.py: instant connectivity check + VM cache sync after create/update; broadcast server_added/updated/deleted events via WebSocket - websocket_manager: add broadcast_event() helper for global fire-and-forget events Frontend: - use-nodes: invalidate topology, VM list and resources after server create/update/delete - use-realtime-sync: new hook with useGlobalRealtimeSync subscribing to server_* WS events and invalidating relevant React Query caches - AppLayout: mount useGlobalRealtimeSync globally Deploy: - deploy.sh: auto-install pve CLI during main/quick deploy (idempotent, --auto mode)
180 lines
7.3 KiB
Python
180 lines
7.3 KiB
Python
"""
|
|
WebSocket Connection Manager for real-time task updates and metrics push.
|
|
Maintains per-user WebSocket connections, broadcasts task events,
|
|
and supports channel-based subscriptions for metrics streaming.
|
|
"""
|
|
import asyncio
|
|
import json
|
|
from typing import Dict, List, Set
|
|
|
|
from fastapi import WebSocket
|
|
from loguru import logger
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections grouped by user_id and by channel."""
|
|
|
|
def __init__(self):
|
|
# user_id -> set of active WebSocket connections
|
|
self._connections: Dict[int, Set[WebSocket]] = {}
|
|
# channel_name -> set of subscribed WebSocket connections
|
|
self._channel_subscribers: Dict[str, Set[WebSocket]] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, user_id: int) -> None:
|
|
await websocket.accept()
|
|
if user_id not in self._connections:
|
|
self._connections[user_id] = set()
|
|
self._connections[user_id].add(websocket)
|
|
logger.debug(f"[WS TASKS] User {user_id} connected ({len(self._connections[user_id])} socket(s))")
|
|
|
|
def disconnect(self, websocket: WebSocket, user_id: int) -> None:
|
|
if user_id in self._connections:
|
|
self._connections[user_id].discard(websocket)
|
|
if not self._connections[user_id]:
|
|
del self._connections[user_id]
|
|
self.unsubscribe_all(websocket)
|
|
logger.debug(f"[WS TASKS] User {user_id} disconnected")
|
|
|
|
# ── Channel subscription helpers ──────────────────────────────────
|
|
|
|
def subscribe(self, websocket: WebSocket, channel: str) -> None:
|
|
"""Subscribe a WebSocket connection to a named channel."""
|
|
if channel not in self._channel_subscribers:
|
|
self._channel_subscribers[channel] = set()
|
|
self._channel_subscribers[channel].add(websocket)
|
|
logger.debug(f"[WS] Subscribed to channel '{channel}'")
|
|
|
|
def unsubscribe(self, websocket: WebSocket, channel: str) -> None:
|
|
"""Unsubscribe a WebSocket connection from a named channel."""
|
|
if channel in self._channel_subscribers:
|
|
self._channel_subscribers[channel].discard(websocket)
|
|
if not self._channel_subscribers[channel]:
|
|
del self._channel_subscribers[channel]
|
|
logger.debug(f"[WS] Unsubscribed from channel '{channel}'")
|
|
|
|
def unsubscribe_all(self, websocket: WebSocket) -> None:
|
|
"""Remove a WebSocket from every channel it is subscribed to."""
|
|
for channel in list(self._channel_subscribers):
|
|
self._channel_subscribers[channel].discard(websocket)
|
|
if not self._channel_subscribers[channel]:
|
|
del self._channel_subscribers[channel]
|
|
|
|
def has_channel_subscribers(self, channel: str) -> bool:
|
|
"""Return True if at least one connection is subscribed to *channel*."""
|
|
return bool(self._channel_subscribers.get(channel))
|
|
|
|
def get_channels_matching(self, prefix: str) -> List[str]:
|
|
"""Return all channel names that start with *prefix*."""
|
|
return [ch for ch in list(self._channel_subscribers) if ch.startswith(prefix)]
|
|
|
|
async def broadcast_to_channel(self, channel: str, data: dict) -> None:
|
|
"""Send JSON data to all connections subscribed to *channel*."""
|
|
connections = self._channel_subscribers.get(channel)
|
|
if not connections:
|
|
return
|
|
message = json.dumps(data, default=str)
|
|
dead: Set[WebSocket] = set()
|
|
for ws in list(connections):
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
dead.add(ws)
|
|
for ws in dead:
|
|
self._channel_subscribers[channel].discard(ws)
|
|
if not self._channel_subscribers.get(channel):
|
|
self._channel_subscribers.pop(channel, None)
|
|
|
|
@property
|
|
def active_user_count(self) -> int:
|
|
return len(self._connections)
|
|
|
|
async def send_to_user(self, user_id: int, data: dict) -> None:
|
|
"""Send a JSON message to all WebSocket connections of a specific user."""
|
|
connections = self._connections.get(user_id)
|
|
if not connections:
|
|
return
|
|
message = json.dumps(data, default=str)
|
|
dead: Set[WebSocket] = set()
|
|
for ws in list(connections):
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
dead.add(ws)
|
|
for ws in dead:
|
|
self._connections[user_id].discard(ws)
|
|
if not self._connections.get(user_id):
|
|
self._connections.pop(user_id, None)
|
|
|
|
async def broadcast(self, data: dict) -> None:
|
|
"""Send a JSON message to all connected users."""
|
|
message = json.dumps(data, default=str)
|
|
dead_pairs = []
|
|
for user_id, connections in list(self._connections.items()):
|
|
for ws in list(connections):
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
dead_pairs.append((user_id, ws))
|
|
for user_id, ws in dead_pairs:
|
|
if user_id in self._connections:
|
|
self._connections[user_id].discard(ws)
|
|
|
|
|
|
# ──────────────────────────────────────────────
|
|
# Module-level singleton used across the app
|
|
# ──────────────────────────────────────────────
|
|
ws_manager = ConnectionManager()
|
|
|
|
|
|
# Stores the main asyncio event loop set at application startup.
|
|
# Background threads (APScheduler) use this to schedule coroutines safely.
|
|
_main_loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
|
|
def set_main_loop(loop: asyncio.AbstractEventLoop) -> None:
|
|
"""Call this once at application startup to register the running event loop."""
|
|
global _main_loop
|
|
_main_loop = loop
|
|
|
|
|
|
def run_async_safe(coro) -> None:
|
|
"""
|
|
Schedule an async coroutine from a synchronous background thread.
|
|
Used by APScheduler jobs to push WS updates without blocking.
|
|
"""
|
|
global _main_loop
|
|
try:
|
|
# Prefer the stored main loop (set at startup) for cross-thread scheduling
|
|
if _main_loop is not None and _main_loop.is_running():
|
|
asyncio.run_coroutine_threadsafe(coro, _main_loop)
|
|
return
|
|
# Fallback: try to get/run the current thread's loop
|
|
loop = asyncio.get_event_loop()
|
|
if loop.is_running():
|
|
asyncio.run_coroutine_threadsafe(coro, loop)
|
|
else:
|
|
loop.run_until_complete(coro)
|
|
except RuntimeError:
|
|
# No event loop yet (e.g. during startup); silently skip
|
|
pass
|
|
|
|
|
|
def broadcast_task_update(user_id: int, event_type: str, task_data: dict) -> None:
|
|
"""
|
|
Fire-and-forget WebSocket broadcast from a sync context (e.g. APScheduler).
|
|
|
|
Wraps `ws_manager.send_to_user` and schedules it on the running asyncio loop.
|
|
"""
|
|
payload = {"type": event_type, "task": task_data}
|
|
run_async_safe(ws_manager.send_to_user(user_id, payload))
|
|
|
|
|
|
def broadcast_event(event_type: str, **fields) -> None:
|
|
"""
|
|
Fire-and-forget global broadcast of an event to all connected WebSocket
|
|
clients from a sync context. Used to push real-time invalidations such as
|
|
`server_added`, `server_updated`, `server_deleted`.
|
|
"""
|
|
payload = {"type": event_type, **fields}
|
|
run_async_safe(ws_manager.broadcast(payload))
|