105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
"""WebSocket endpoint for real-time log streaming.
|
|
|
|
Clients connect to /ws/logs and receive live log lines from all modules.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from collections import deque
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
from Project.Sanad.core.logger import set_ws_push
|
|
|
|
router = APIRouter()
|
|
|
|
MAX_WATCHERS = 50
|
|
|
|
# Ring buffer of recent log lines (shared across connections).
|
|
_recent: deque[str] = deque(maxlen=500)
|
|
# Each watcher is an (event_loop, queue) pair. We keep the loop so cross-thread
|
|
# producers can schedule the enqueue on the consumer's loop (asyncio.Queue is
|
|
# NOT thread-safe — calling put_nowait off-loop neither wakes the parked
|
|
# `await queue.get()` nor safely mutates the queue's internals).
|
|
_watchers: set[tuple[asyncio.AbstractEventLoop, asyncio.Queue]] = set()
|
|
_watchers_lock = threading.Lock()
|
|
|
|
|
|
def push_log_line(line: str):
|
|
"""Called from the logging system to feed new lines.
|
|
|
|
May be called from ANY thread (logging is multi-threaded), so the append
|
|
to _recent and the per-watcher enqueue are done together under the same
|
|
lock that log_ws holds while snapshotting history + registering — that
|
|
closes the history/live overlap window so a connecting client can't see a
|
|
line both in its history replay and again live. The enqueue itself is
|
|
marshalled onto each watcher's loop via call_soon_threadsafe because
|
|
asyncio.Queue.put_nowait is not safe to call from a foreign thread.
|
|
"""
|
|
with _watchers_lock:
|
|
_recent.append(line)
|
|
snapshot = list(_watchers)
|
|
for loop, q in snapshot:
|
|
try:
|
|
loop.call_soon_threadsafe(_safe_put, q, line)
|
|
except RuntimeError:
|
|
# Loop already closed — watcher is going away; skip it.
|
|
pass
|
|
|
|
|
|
def _safe_put(q: asyncio.Queue, line: str) -> None:
|
|
"""Enqueue on the consumer's own loop thread (so it's safe)."""
|
|
try:
|
|
q.put_nowait(line)
|
|
except asyncio.QueueFull:
|
|
# Drop on overflow rather than block — logs are not critical data
|
|
pass
|
|
|
|
|
|
# Register with the logger so all log records are pushed to WS clients.
|
|
# Wrap so a logger registration failure doesn't break Dashboard import.
|
|
try:
|
|
set_ws_push(push_log_line)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@router.websocket("/ws/logs")
|
|
async def log_ws(ws: WebSocket):
|
|
await ws.accept()
|
|
|
|
loop = asyncio.get_running_loop()
|
|
queue: asyncio.Queue[str] = asyncio.Queue(maxsize=200)
|
|
watcher = (loop, queue)
|
|
with _watchers_lock:
|
|
if len(_watchers) >= MAX_WATCHERS:
|
|
await ws.close(code=1013, reason="Too many log watchers")
|
|
return
|
|
# Register the live queue and snapshot history under the SAME lock that
|
|
# push_log_line holds — so every line is either in this history
|
|
# snapshot or arrives on the queue, never both (no replay duplicates).
|
|
_watchers.add(watcher)
|
|
history = list(_recent)
|
|
|
|
try:
|
|
# Send recent history
|
|
for line in history:
|
|
await ws.send_text(line)
|
|
|
|
while True:
|
|
line = await queue.get()
|
|
await ws.send_text(line)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception:
|
|
# Any other error closes the connection cleanly
|
|
try:
|
|
await ws.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
with _watchers_lock:
|
|
_watchers.discard(watcher)
|