Your Name
32feb1b6db
Fixes: - F821: install.py site_url 未定义(NameError)、sync_v2.py os 未导入、script_jobs.py LOG f-string - F401: 移除 22 个未使用导入（models/__init__.py、install.py、auth.py 等） - B904: 20 个 except 块 raise 添加 from e/from None - S110: 20 个有意的 try-except-pass 添加 noqa 注释（SSH清理/DDL幂等/WS断连） - B007: 3 个未使用循环变量重命名（dirs→_dirs、server_id→_server_id） - F541: 3 个无占位符 f-string 修正 - F841: auth.py 未使用 ip_address 变量移除 - S105: auth_service.py Redis key prefix 误报 noqa - B023: script_execution_flush.py 循环变量绑定修复 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
501 lines
18 KiB
Python
501 lines
18 KiB
Python
"""Nexus — WebSocket Alert Push (Two-Layer Architecture)
|
|
|
|
ADR-010: Memory ConnectionManager + Redis Pub/Sub two-layer architecture.
|
|
- Layer 1: In-memory ConnectionManager tracks all local WebSocket connections
|
|
- Layer 2: Redis Pub/Sub relays messages across multiple uvicorn workers
|
|
- Alert/recovery/system events published to Redis channel `nexus:alerts`
|
|
- Each worker subscribes to `nexus:alerts` and pushes to local clients
|
|
|
|
ADR-011: WebSocket JWT authentication via query parameter `?token=xxx`
|
|
|
|
R4: Ping/pong heartbeat every 30s + zombie connection cleanup (60s idle → disconnect)
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from typing import Dict, Optional
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
|
|
|
from server.config import settings
|
|
|
|
logger = logging.getLogger("nexus.websocket")
|
|
|
|
router = APIRouter()
|
|
|
|
# Redis Pub/Sub channel name
|
|
REDIS_CHANNEL = "nexus:alerts"
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Layer 1: In-memory WebSocket connection manager.
|
|
|
|
Tracks all connected clients for this worker process.
|
|
Handles ping/pong heartbeat and zombie connection cleanup.
|
|
"""
|
|
|
|
MAX_CONNECTIONS = 500
|
|
|
|
def __init__(self):
|
|
# client_id → {websocket, last_pong_time}
|
|
self._connections: Dict[str, dict] = {}
|
|
self._heartbeat_task: Optional[asyncio.Task] = None
|
|
|
|
async def connect(self, client_id: str, websocket: WebSocket):
|
|
"""Accept and register a new WebSocket connection"""
|
|
if len(self._connections) >= self.MAX_CONNECTIONS:
|
|
logger.warning(f"WebSocket connection limit reached ({self.MAX_CONNECTIONS}), rejecting {client_id}")
|
|
await websocket.close(code=4003, reason="Too many connections")
|
|
return
|
|
await websocket.accept()
|
|
self._connections[client_id] = {
|
|
"websocket": websocket,
|
|
"last_pong": time.monotonic(),
|
|
}
|
|
logger.info(f"WebSocket connected: {client_id}, total: {len(self._connections)}")
|
|
|
|
# Start heartbeat task if first connection
|
|
if len(self._connections) == 1 and self._heartbeat_task is None:
|
|
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="ws_heartbeat")
|
|
|
|
def disconnect(self, client_id: str):
|
|
"""Remove a WebSocket connection and close the underlying TCP connection"""
|
|
conn = self._connections.pop(client_id, None)
|
|
if conn:
|
|
# Close the WebSocket to release the TCP connection
|
|
try:
|
|
asyncio.get_running_loop().create_task(
|
|
conn["websocket"].close(code=1000, reason="cleanup")
|
|
)
|
|
except RuntimeError:
|
|
pass # No running event loop (e.g. during shutdown)
|
|
except Exception: # noqa: S110 — client disconnected, best-effort cleanup
|
|
pass
|
|
logger.info(f"WebSocket disconnected: {client_id}, total: {len(self._connections)}")
|
|
|
|
# Stop heartbeat task if no connections
|
|
if not self._connections and self._heartbeat_task:
|
|
self._heartbeat_task.cancel()
|
|
self._heartbeat_task = None
|
|
|
|
async def send_to_client(self, client_id: str, message: dict) -> bool:
|
|
"""Send JSON message to a specific client. Returns False if send failed."""
|
|
conn = self._connections.get(client_id)
|
|
if not conn:
|
|
return False
|
|
try:
|
|
await conn["websocket"].send_json(message)
|
|
return True
|
|
except Exception:
|
|
self.disconnect(client_id)
|
|
return False
|
|
|
|
async def broadcast_local(self, message: dict):
|
|
"""Broadcast message to all locally connected clients (Layer 1 only)"""
|
|
stale = []
|
|
for client_id, conn in self._connections.items():
|
|
try:
|
|
await conn["websocket"].send_json(message)
|
|
except Exception:
|
|
stale.append(client_id)
|
|
for client_id in stale:
|
|
self.disconnect(client_id)
|
|
|
|
def update_pong(self, client_id: str):
|
|
"""Update last pong time for a client (received pong response)"""
|
|
conn = self._connections.get(client_id)
|
|
if conn:
|
|
conn["last_pong"] = time.monotonic()
|
|
|
|
@property
|
|
def client_count(self) -> int:
|
|
"""Local client count (this worker only)"""
|
|
return len(self._connections)
|
|
|
|
async def global_client_count(self) -> int:
|
|
"""Global client count across all workers (via Redis).
|
|
|
|
In multi-worker deployments, each worker only knows its local connections.
|
|
This method queries Redis for the aggregate count.
|
|
"""
|
|
try:
|
|
from server.infrastructure.redis.client import get_redis
|
|
redis = get_redis()
|
|
# Each worker periodically writes its count to Redis with TTL
|
|
await redis.set(f"ws:count:{id(self)}", len(self._connections), ex=90)
|
|
# Sum all worker counts
|
|
keys = await redis.keys("ws:count:*")
|
|
if not keys:
|
|
return len(self._connections)
|
|
total = 0
|
|
for key in keys:
|
|
val = await redis.get(key)
|
|
if val:
|
|
total += int(val)
|
|
return total
|
|
except Exception:
|
|
return len(self._connections)
|
|
|
|
async def _heartbeat_loop(self):
|
|
"""R4: Send ping every 30s, cleanup connections idle > 60s"""
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(30)
|
|
now = time.monotonic()
|
|
stale = []
|
|
|
|
for client_id, conn in self._connections.items():
|
|
# Check for zombie connections (no pong for 60s)
|
|
if now - conn["last_pong"] > 60:
|
|
stale.append(client_id)
|
|
continue
|
|
|
|
# Send ping
|
|
try:
|
|
await conn["websocket"].send_json({"type": "ping"})
|
|
except Exception:
|
|
stale.append(client_id)
|
|
|
|
# Cleanup stale connections
|
|
for client_id in stale:
|
|
self.disconnect(client_id)
|
|
logger.info(f"WebSocket zombie cleanup: {client_id}")
|
|
|
|
if stale:
|
|
logger.info(f"Heartbeat cleanup: removed {len(stale)} stale connections")
|
|
except asyncio.CancelledError:
|
|
logger.debug("WebSocket heartbeat task cancelled")
|
|
raise
|
|
|
|
def cancel_heartbeat(self):
|
|
"""Cancel the heartbeat task — called during app shutdown (H11)"""
|
|
if self._heartbeat_task:
|
|
self._heartbeat_task.cancel()
|
|
self._heartbeat_task = None
|
|
|
|
|
|
# Global connection manager for this worker
|
|
manager = ConnectionManager()
|
|
|
|
|
|
# ── Redis Pub/Sub Layer 2 ──
|
|
|
|
_pubsub_task: Optional[asyncio.Task] = None
|
|
_redis_subscriber = None
|
|
|
|
|
|
async def start_redis_subscriber():
|
|
"""Start Redis Pub/Sub subscriber (Layer 2) — called during app startup"""
|
|
global _pubsub_task, _redis_subscriber
|
|
|
|
try:
|
|
from server.infrastructure.redis.client import get_redis
|
|
redis = get_redis()
|
|
_redis_subscriber = redis.pubsub()
|
|
await _redis_subscriber.subscribe(REDIS_CHANNEL)
|
|
|
|
_pubsub_task = asyncio.create_task(_redis_listen_loop(), name="ws_redis_sub")
|
|
logger.info(f"Redis Pub/Sub subscriber started on channel: {REDIS_CHANNEL}")
|
|
except Exception as e:
|
|
logger.error(f"Redis Pub/Sub subscriber failed to start: {e}")
|
|
# Non-fatal — WebSocket still works locally
|
|
|
|
|
|
async def stop_redis_subscriber():
|
|
"""Stop Redis Pub/Sub subscriber — called during app shutdown"""
|
|
global _pubsub_task, _redis_subscriber
|
|
|
|
if _pubsub_task:
|
|
_pubsub_task.cancel()
|
|
try:
|
|
await _pubsub_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
_pubsub_task = None
|
|
|
|
if _redis_subscriber:
|
|
try:
|
|
await _redis_subscriber.unsubscribe(REDIS_CHANNEL)
|
|
await _redis_subscriber.aclose()
|
|
except Exception:
|
|
logger.debug("Failed to unsubscribe Redis Pub/Sub during shutdown", exc_info=True)
|
|
_redis_subscriber = None
|
|
|
|
logger.info("Redis Pub/Sub subscriber stopped")
|
|
|
|
|
|
async def _redis_listen_loop():
|
|
"""Listen for messages on Redis Pub/Sub channel and broadcast to local clients"""
|
|
while True:
|
|
try:
|
|
message = await _redis_subscriber.get_message(
|
|
ignore_subscribe_messages=True, timeout=1.0
|
|
)
|
|
if message and message.get("type") == "message":
|
|
data = message.get("data")
|
|
if isinstance(data, str):
|
|
msg = json.loads(data)
|
|
await manager.broadcast_local(msg)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Redis Pub/Sub listen error: {e}")
|
|
await asyncio.sleep(1) # Back off on error
|
|
|
|
|
|
async def _publish_to_redis(message: dict) -> bool:
|
|
"""Publish message to Redis channel (Layer 2) — other workers will receive it"""
|
|
try:
|
|
from server.infrastructure.redis.client import get_redis
|
|
redis = get_redis()
|
|
await redis.publish(REDIS_CHANNEL, json.dumps(message))
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Redis Pub/Sub publish failed: {e}")
|
|
return False
|
|
|
|
|
|
async def _dispatch_ws_message(message: dict) -> None:
|
|
"""Deliver alert to local WebSocket clients without duplicate delivery.
|
|
|
|
When Pub/Sub is active, only publish — this worker's subscriber calls broadcast_local once.
|
|
Fallback to direct broadcast if Redis subscriber is unavailable.
|
|
"""
|
|
if _pubsub_task is not None and _redis_subscriber is not None:
|
|
if not await _publish_to_redis(message):
|
|
await manager.broadcast_local(message)
|
|
else:
|
|
await manager.broadcast_local(message)
|
|
|
|
|
|
# ── WebSocket Endpoint ──
|
|
|
|
@router.websocket("/ws/alerts")
|
|
async def alert_ws(
|
|
websocket: WebSocket,
|
|
token: Optional[str] = Query(None, description="JWT access token"),
|
|
):
|
|
"""WebSocket endpoint for real-time alert push to frontend
|
|
|
|
ADR-011: JWT authentication via query parameter `?token=xxx`
|
|
R4: Ping/pong heartbeat + zombie connection cleanup
|
|
"""
|
|
# ── JWT Authentication (ADR-011) ──
|
|
if not token:
|
|
await websocket.close(code=4001, reason="Missing JWT token")
|
|
return
|
|
|
|
admin = await _verify_ws_token(token)
|
|
if not admin:
|
|
await websocket.close(code=4001, reason="Invalid or expired JWT token")
|
|
return
|
|
|
|
# ── Accept and register ──
|
|
client_id = f"{admin.id}:{id(websocket)}"
|
|
await manager.connect(client_id, websocket)
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
if data == "pong":
|
|
manager.update_pong(client_id)
|
|
elif data == "ping":
|
|
await websocket.send_json({"type": "pong"})
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(client_id)
|
|
except Exception as e:
|
|
logger.warning(f"WebSocket error for {client_id}: {e}")
|
|
manager.disconnect(client_id)
|
|
|
|
|
|
async def _verify_ws_token(token: str):
|
|
"""Verify JWT token for WebSocket connection (same logic as get_current_admin)
|
|
|
|
Checks: valid signature, not expired, sub exists, admin active, and
|
|
updated_at matches (invalidates token after password change).
|
|
"""
|
|
try:
|
|
import jwt as pyjwt
|
|
from server.infrastructure.database.session import AsyncSessionLocal
|
|
from server.infrastructure.database.admin_repo import AdminRepositoryImpl
|
|
|
|
payload = pyjwt.decode(token, settings.SECRET_KEY, algorithms=["HS256"],
|
|
options={"require": ["exp", "sub"]})
|
|
admin_id = payload.get("sub")
|
|
if not admin_id:
|
|
return None
|
|
|
|
try:
|
|
admin_id = int(admin_id)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
|
|
async with AsyncSessionLocal() as session:
|
|
admin_repo = AdminRepositoryImpl(session)
|
|
admin = await admin_repo.get_by_id(admin_id)
|
|
if not admin or not admin.is_active:
|
|
return None
|
|
|
|
# Primary: token_version must match (immediate invalidation on password change)
|
|
token_tv = payload.get("tv", -1)
|
|
if token_tv != admin.token_version:
|
|
return None
|
|
|
|
# Secondary: updated_at timestamp (defense in depth)
|
|
token_updated = payload.get("updated", 0)
|
|
if token_updated and admin.updated_at:
|
|
admin_ts = int(admin.updated_at.timestamp())
|
|
if admin_ts > token_updated + 5:
|
|
return None
|
|
|
|
return admin
|
|
except Exception:
|
|
logger.debug("WebSocket JWT verification failed", exc_info=True)
|
|
return None
|
|
|
|
|
|
# ── Alert Aggregation / Dedup ──
|
|
# Track last alert time per (server_id, alert_type) to avoid Telegram alert storms.
|
|
# WebSocket broadcasts are always sent (real-time), but Telegram pushes are deduplicated.
|
|
_ALERT_COOLDOWN_SECONDS = 300 # 5 minutes — same server+metric won't re-push Telegram
|
|
_last_alert_time: Dict[str, float] = {} # key="alert:{server_id}:{type}" → monotonic timestamp
|
|
|
|
|
|
def _should_push_telegram(alert_key: str) -> bool:
|
|
"""Check if enough time has passed since last Telegram push for this alert key.
|
|
|
|
Returns True if we should send the Telegram notification, False if suppressed (cooldown).
|
|
Also prunes stale entries older than 1 hour to prevent unbounded memory growth.
|
|
"""
|
|
now = time.monotonic()
|
|
|
|
# Prune entries older than 1 hour (prevent unbounded growth)
|
|
stale_keys = [k for k, v in _last_alert_time.items() if now - v > 3600]
|
|
for k in stale_keys:
|
|
del _last_alert_time[k]
|
|
|
|
last = _last_alert_time.get(alert_key, 0)
|
|
if now - last < _ALERT_COOLDOWN_SECONDS:
|
|
return False # Suppressed — too soon
|
|
_last_alert_time[alert_key] = now
|
|
return True
|
|
|
|
|
|
# ── Broadcast Functions (called from other modules) ──
|
|
|
|
async def _save_alert_log(server_id: int, server_name: str, alert_type: str, value: str, is_recovery: bool):
|
|
"""Persist alert / recovery event to MySQL alert_logs table."""
|
|
try:
|
|
from server.infrastructure.database.session import AsyncSessionLocal
|
|
from server.domain.models import AlertLog
|
|
async with AsyncSessionLocal() as session:
|
|
session.add(AlertLog(
|
|
server_id=server_id,
|
|
server_name=server_name or None,
|
|
alert_type=alert_type,
|
|
value=value,
|
|
is_recovery=is_recovery,
|
|
))
|
|
await session.commit()
|
|
except Exception as e:
|
|
logger.debug("Failed to save alert log: %s", e)
|
|
|
|
|
|
async def broadcast_alert(server_id: int, alert_type: str, alert_value: float, server_name: str = ""):
|
|
"""Broadcast alert to all WebSocket clients + Telegram
|
|
|
|
Called when Agent heartbeat detects CPU/memory/disk > threshold.
|
|
Publishes to Redis Pub/Sub (multi-worker) + Telegram.
|
|
|
|
WebSocket refreshes dashboard stats only (no browser sound/list UI).
|
|
Telegram push is deduplicated (5min cooldown per server+metric).
|
|
"""
|
|
msg = {
|
|
"type": "alert",
|
|
"server_id": server_id,
|
|
"server_name": server_name,
|
|
"alert_type": alert_type,
|
|
"alert_value": alert_value,
|
|
}
|
|
|
|
# Redis Pub/Sub (multi-worker) — local clients via subscriber only (C9)
|
|
await _dispatch_ws_message(msg)
|
|
|
|
# Persist to alert_logs
|
|
await _save_alert_log(server_id, server_name, alert_type, f"{alert_value:.1f}%", is_recovery=False)
|
|
|
|
# Telegram push — deduplicated per (server_id, alert_type) with cooldown
|
|
alert_key = f"alert:{server_id}:{alert_type}"
|
|
if _should_push_telegram(alert_key):
|
|
from server.infrastructure.telegram import send_telegram_alert
|
|
await send_telegram_alert(server_name or str(server_id), alert_type, alert_value)
|
|
else:
|
|
logger.debug(f"Alert suppressed (cooldown): {alert_key}")
|
|
|
|
|
|
async def broadcast_recovery(server_id: int, metric: str, value: float, server_name: str = ""):
|
|
"""Broadcast recovery notification when alert metric returns to normal
|
|
|
|
Recovery alerts are always pushed (clears the alert state).
|
|
Also clears the alert cooldown so next alert will be sent immediately.
|
|
"""
|
|
msg = {
|
|
"type": "recovery",
|
|
"server_id": server_id,
|
|
"server_name": server_name,
|
|
"metric": metric,
|
|
"value": value,
|
|
}
|
|
|
|
await _dispatch_ws_message(msg)
|
|
|
|
# Persist recovery to alert_logs
|
|
await _save_alert_log(server_id, server_name, metric, f"{value:.1f}%", is_recovery=True)
|
|
|
|
# Recovery always sends Telegram + clears cooldown for next alert cycle
|
|
alert_key = f"alert:{server_id}:{metric}"
|
|
_last_alert_time.pop(alert_key, None) # Clear cooldown
|
|
|
|
from server.infrastructure.telegram import send_telegram_recovery
|
|
await send_telegram_recovery(server_name or str(server_id), metric, value)
|
|
|
|
|
|
async def broadcast_system_event(event_type: str, message: str):
|
|
"""Broadcast system-level events (Python restart, Redis reconnect, etc.)"""
|
|
msg = {"type": "system", "event_type": event_type, "message": message}
|
|
|
|
await _dispatch_ws_message(msg)
|
|
|
|
|
|
async def broadcast_sync_progress(
|
|
batch_id: str, server_id: int, server_name: str,
|
|
status: str, completed: int, failed: int, total: int,
|
|
error_message: str = None, duration_seconds: int = None,
|
|
retry_job_id: int = None, cancelled: int = 0,
|
|
):
|
|
"""Broadcast per-server push progress to all WebSocket clients.
|
|
|
|
Called from sync_engine_v2._sync_one() after each server completes.
|
|
Frontend uses batch_id to filter stale events from previous pushes.
|
|
retry_job_id is set when a failed push auto-creates a retry job.
|
|
cancelled is the count of servers skipped due to user cancellation.
|
|
"""
|
|
msg = {
|
|
"type": "sync_progress",
|
|
"batch_id": batch_id,
|
|
"server_id": server_id,
|
|
"server_name": server_name,
|
|
"status": status,
|
|
"completed": completed,
|
|
"failed": failed,
|
|
"total": total,
|
|
"error_message": error_message,
|
|
"duration_seconds": duration_seconds,
|
|
"retry_job_id": retry_job_id,
|
|
"cancelled": cancelled,
|
|
}
|
|
await _dispatch_ws_message(msg)
|
|
|