80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
# --- routers/logs.py ---
|
||
import asyncio
|
||
import logging
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||
from logger import ColoredFormatter, register_websocket_handler
|
||
|
||
router = APIRouter()
|
||
|
||
class LogBroadcaster:
|
||
"""管理所有活躍 WebSocket 連線的廣播器 (Connection Manager)"""
|
||
def __init__(self):
|
||
self.active_connections: set[WebSocket] = set()
|
||
self.loop: asyncio.AbstractEventLoop = None
|
||
|
||
async def connect(self, websocket: WebSocket):
|
||
await websocket.accept()
|
||
self.active_connections.add(websocket)
|
||
|
||
def disconnect(self, websocket: WebSocket):
|
||
self.active_connections.discard(websocket)
|
||
|
||
async def broadcast(self, message: str):
|
||
"""非同步推播日誌給所有連線中的客戶端"""
|
||
if not self.active_connections:
|
||
return
|
||
|
||
# 併發發送，並透過 return_exceptions=True 確保單一連線異常不影響其他客戶端
|
||
tasks = [self._safe_send(conn, message) for conn in self.active_connections]
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
async def _safe_send(self, websocket: WebSocket, message: str):
|
||
try:
|
||
await websocket.send_text(message)
|
||
except Exception:
|
||
self.disconnect(websocket)
|
||
|
||
# 建立全域廣播器實例
|
||
log_broadcaster = LogBroadcaster()
|
||
|
||
class WebSocketLogHandler(logging.Handler):
|
||
"""自訂 Log Handler，攔截系統日誌並安全地派發至非同步廣播器"""
|
||
def __init__(self, broadcaster: LogBroadcaster):
|
||
super().__init__()
|
||
self.broadcaster = broadcaster
|
||
# 沿用系統 ColoredFormatter，完美保留 ANSI 色碼
|
||
self.setFormatter(ColoredFormatter())
|
||
|
||
def emit(self, record):
|
||
try:
|
||
msg = self.format(record)
|
||
loop = self.broadcaster.loop
|
||
# 確保在 Event Loop 處於執行狀態時，安全地跨執行緒派發任務
|
||
if loop and loop.is_running():
|
||
loop.call_soon_threadsafe(
|
||
lambda: asyncio.create_task(self.broadcaster.broadcast(msg))
|
||
)
|
||
except Exception:
|
||
self.handleError(record)
|
||
|
||
# 建立 Handler 實例並註冊至所有受管控的 Logger
|
||
websocket_log_handler = WebSocketLogHandler(log_broadcaster)
|
||
register_websocket_handler(websocket_log_handler)
|
||
|
||
@router.websocket("/ws/logs")
|
||
async def websocket_logs(websocket: WebSocket):
|
||
"""WebSocket 實時日誌串流端點"""
|
||
# 若尚未綁定 Event Loop，則於首次連線時動態綁定
|
||
if not log_broadcaster.loop:
|
||
log_broadcaster.loop = asyncio.get_running_loop()
|
||
|
||
await log_broadcaster.connect(websocket)
|
||
try:
|
||
# 保持連線，監聽客戶端斷線狀態
|
||
while True:
|
||
await websocket.receive_text()
|
||
except WebSocketDisconnect:
|
||
log_broadcaster.disconnect(websocket)
|
||
except Exception:
|
||
log_broadcaster.disconnect(websocket)
|