Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WebSocket 模块
|
||||
提供 WebSocket 连接和消息处理功能
|
||||
"""
|
||||
from core.websocket.router import router
|
||||
|
||||
__all__ = ['router']
|
||||
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WebSocket Consumers 模块
|
||||
"""
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer
|
||||
from core.websocket.consumers.test_consumer import TestWebSocketConsumer
|
||||
from core.websocket.consumers.notification_consumer import NotificationConsumer
|
||||
from core.websocket.consumers.server_monitor_consumer import ServerMonitorConsumer
|
||||
from core.websocket.consumers.redis_monitor_consumer import RedisMonitorConsumer
|
||||
from core.websocket.consumers.database_monitor_consumer import DatabaseMonitorConsumer
|
||||
|
||||
__all__ = [
|
||||
'TokenAuthWebSocketConsumer',
|
||||
'TestWebSocketConsumer',
|
||||
'NotificationConsumer',
|
||||
'ServerMonitorConsumer',
|
||||
'RedisMonitorConsumer',
|
||||
'DatabaseMonitorConsumer',
|
||||
]
|
||||
@@ -0,0 +1,240 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WebSocket 基础消费者类
|
||||
提供 Token 认证和基础消息处理功能
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Optional, Dict, Any, Set
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
|
||||
from app.config import settings
|
||||
from utils.security import verify_access_token
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConnectionManager:
|
||||
"""WebSocket 连接管理器"""
|
||||
|
||||
def __init__(self):
|
||||
# 活跃连接: {user_id: {websocket1, websocket2, ...}}
|
||||
self.active_connections: Dict[str, Set[WebSocket]] = {}
|
||||
# 组连接: {group_name: {websocket1, websocket2, ...}}
|
||||
self.groups: Dict[str, Set[WebSocket]] = {}
|
||||
|
||||
async def connect(self, websocket: WebSocket, user_id: str):
|
||||
"""添加连接"""
|
||||
await websocket.accept()
|
||||
if user_id not in self.active_connections:
|
||||
self.active_connections[user_id] = set()
|
||||
self.active_connections[user_id].add(websocket)
|
||||
|
||||
def disconnect(self, websocket: WebSocket, user_id: str):
|
||||
"""移除连接"""
|
||||
if user_id in self.active_connections:
|
||||
self.active_connections[user_id].discard(websocket)
|
||||
if not self.active_connections[user_id]:
|
||||
del self.active_connections[user_id]
|
||||
|
||||
# 从所有组中移除
|
||||
for group_name in list(self.groups.keys()):
|
||||
self.groups[group_name].discard(websocket)
|
||||
if not self.groups[group_name]:
|
||||
del self.groups[group_name]
|
||||
|
||||
async def group_add(self, group_name: str, websocket: WebSocket):
|
||||
"""将连接添加到组"""
|
||||
if group_name not in self.groups:
|
||||
self.groups[group_name] = set()
|
||||
self.groups[group_name].add(websocket)
|
||||
|
||||
async def group_discard(self, group_name: str, websocket: WebSocket):
|
||||
"""从组中移除连接"""
|
||||
if group_name in self.groups:
|
||||
self.groups[group_name].discard(websocket)
|
||||
if not self.groups[group_name]:
|
||||
del self.groups[group_name]
|
||||
|
||||
async def broadcast_to_group(self, group_name: str, message: dict):
|
||||
"""向组内所有连接广播消息"""
|
||||
if group_name in self.groups:
|
||||
message_text = json.dumps(message)
|
||||
for websocket in list(self.groups[group_name]):
|
||||
try:
|
||||
await websocket.send_text(message_text)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def send_to_user(self, user_id: str, message: dict):
|
||||
"""向指定用户的所有连接发送消息"""
|
||||
if user_id in self.active_connections:
|
||||
message_text = json.dumps(message)
|
||||
for websocket in list(self.active_connections[user_id]):
|
||||
try:
|
||||
await websocket.send_text(message_text)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def is_online(self, user_id: str) -> bool:
|
||||
"""判断用户是否在线（有活跃的WebSocket连接）"""
|
||||
return user_id in self.active_connections and len(self.active_connections[user_id]) > 0
|
||||
|
||||
def get_online_user_ids(self) -> Set[str]:
|
||||
"""获取所有在线用户ID"""
|
||||
return set(self.active_connections.keys())
|
||||
|
||||
|
||||
# 全局连接管理器实例
|
||||
manager = ConnectionManager()
|
||||
|
||||
|
||||
class TokenAuthWebSocketConsumer:
|
||||
"""基于Token认证的WebSocket消费者基类"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
self.websocket = websocket
|
||||
self.user_id: Optional[str] = None
|
||||
self.is_authenticated = False
|
||||
self._token: Optional[str] = None # 保存原始token用于心跳校验
|
||||
|
||||
async def authenticate(self) -> bool:
|
||||
"""
|
||||
进行Token认证
|
||||
从查询参数中获取token并验证
|
||||
"""
|
||||
# 获取查询参数中的token
|
||||
query_string = self.websocket.scope.get('query_string', b'').decode('utf-8')
|
||||
token = None
|
||||
|
||||
if query_string:
|
||||
query_params = parse_qs(query_string)
|
||||
token_list = query_params.get('token', [])
|
||||
if token_list:
|
||||
token = token_list[0]
|
||||
|
||||
if not token:
|
||||
logger.warning("WebSocket connection rejected: No token provided")
|
||||
# 必须先accept才能close
|
||||
await self.websocket.accept()
|
||||
await self.websocket.close(code=4001)
|
||||
return False
|
||||
|
||||
# 验证token
|
||||
try:
|
||||
payload = verify_access_token(token)
|
||||
|
||||
if not payload:
|
||||
logger.warning("WebSocket connection rejected: Invalid token")
|
||||
await self.websocket.accept()
|
||||
await self.websocket.close(code=4001)
|
||||
return False
|
||||
|
||||
user_id = payload.get('sub')
|
||||
if not user_id:
|
||||
logger.warning("WebSocket connection rejected: Invalid token payload")
|
||||
await self.websocket.accept()
|
||||
await self.websocket.close(code=4001)
|
||||
return False
|
||||
|
||||
self.user_id = user_id
|
||||
self.is_authenticated = True
|
||||
self._token = token # 保存token用于后续心跳校验
|
||||
logger.info(f"WebSocket connection accepted for user {user_id}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket authentication failed: {str(e)}")
|
||||
await self.websocket.accept()
|
||||
await self.websocket.close(code=4001)
|
||||
return False
|
||||
|
||||
async def connect(self):
|
||||
"""连接时进行Token认证"""
|
||||
if await self.authenticate():
|
||||
await manager.connect(self.websocket, self.user_id)
|
||||
|
||||
async def disconnect(self, close_code: int = 1000):
|
||||
"""断开连接"""
|
||||
if self.user_id:
|
||||
manager.disconnect(self.websocket, self.user_id)
|
||||
logger.info(f"WebSocket disconnected with code {close_code}")
|
||||
|
||||
async def receive(self, text_data: str):
|
||||
"""接收消息的基础处理"""
|
||||
try:
|
||||
data = json.loads(text_data)
|
||||
message_type = data.get('type', 'unknown')
|
||||
|
||||
# 根据消息类型处理
|
||||
if message_type == 'ping':
|
||||
await self._handle_ping(data)
|
||||
else:
|
||||
await self.handle_message(data)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
await self.send_error('Invalid JSON format')
|
||||
except Exception as e:
|
||||
logger.error(f"Error receiving message: {str(e)}")
|
||||
await self.send_error(f'处理消息时出错: {str(e)}')
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""子类需要实现的消息处理方法"""
|
||||
await self.send_error('Message type not supported')
|
||||
|
||||
async def send_message(self, message_type: str, message: str, data: Optional[Dict] = None):
|
||||
"""发送消息"""
|
||||
response = {
|
||||
'type': message_type,
|
||||
'message': message,
|
||||
'timestamp': datetime.now().isoformat()
|
||||
}
|
||||
if data:
|
||||
response['data'] = data
|
||||
|
||||
await self.websocket.send_text(json.dumps(response))
|
||||
|
||||
async def send_error(self, error_message: str):
|
||||
"""发送错误消息"""
|
||||
await self.send_message('error', error_message)
|
||||
|
||||
async def _handle_ping(self, data: Dict[str, Any]):
|
||||
"""
|
||||
处理心跳ping消息
|
||||
同时校验token是否仍然有效，过期则通知前端并关闭连接
|
||||
"""
|
||||
if not self._token:
|
||||
await self.send_message('pong', '心跳响应')
|
||||
return
|
||||
|
||||
# 重新验证token有效性
|
||||
payload = verify_access_token(self._token)
|
||||
if not payload:
|
||||
logger.warning(f"WebSocket token expired for user {self.user_id}")
|
||||
await self.send_message('token_expired', 'Access token已过期，请刷新token后重连')
|
||||
# 使用4002关闭码表示token过期（区别于4001认证失败）
|
||||
await self.websocket.close(code=4002)
|
||||
return
|
||||
|
||||
await self.send_message('pong', '心跳响应')
|
||||
|
||||
async def run(self):
|
||||
"""运行WebSocket消费者的主循环"""
|
||||
await self.connect()
|
||||
|
||||
if not self.is_authenticated:
|
||||
return
|
||||
|
||||
try:
|
||||
while True:
|
||||
text_data = await self.websocket.receive_text()
|
||||
await self.receive(text_data)
|
||||
except WebSocketDisconnect as e:
|
||||
await self.disconnect(e.code)
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket error: {str(e)}")
|
||||
await self.disconnect(1011)
|
||||
@@ -0,0 +1,213 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
数据库监控 WebSocket 消费者
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer, manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseMonitorConsumer(TokenAuthWebSocketConsumer):
|
||||
"""数据库监控WebSocket消费者"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
super().__init__(websocket)
|
||||
self.monitor_task: Optional[asyncio.Task] = None
|
||||
self.is_monitoring = False
|
||||
self.monitor_interval = 2 # 固定2秒更新一次
|
||||
self.current_db_name: Optional[str] = None
|
||||
|
||||
async def connect(self):
|
||||
"""连接并开始监控"""
|
||||
await super().connect()
|
||||
if self.is_authenticated and self.user_id:
|
||||
# 加入数据库监控组
|
||||
await manager.group_add(
|
||||
"database_monitor",
|
||||
self.websocket
|
||||
)
|
||||
|
||||
async def disconnect(self, close_code: int = 1000):
|
||||
"""断开连接并停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self.user_id:
|
||||
await manager.group_discard(
|
||||
"database_monitor",
|
||||
self.websocket
|
||||
)
|
||||
await super().disconnect(close_code)
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""处理数据库监控消息"""
|
||||
message_type = data.get('type', 'unknown')
|
||||
|
||||
if message_type == 'start_monitor':
|
||||
db_name = data.get('db_name')
|
||||
if not db_name:
|
||||
await self.send_error('缺少数据库名称参数')
|
||||
return
|
||||
await self.start_monitoring(db_name)
|
||||
elif message_type == 'stop_monitor':
|
||||
await self.stop_monitoring()
|
||||
elif message_type == 'get_overview':
|
||||
db_name = data.get('db_name')
|
||||
if not db_name:
|
||||
await self.send_error('缺少数据库名称参数')
|
||||
return
|
||||
await self.send_database_overview(db_name)
|
||||
elif message_type == 'get_realtime':
|
||||
db_name = data.get('db_name')
|
||||
if not db_name:
|
||||
await self.send_error('缺少数据库名称参数')
|
||||
return
|
||||
await self.send_realtime_stats(db_name)
|
||||
elif message_type == 'test_connection':
|
||||
db_name = data.get('db_name')
|
||||
if not db_name:
|
||||
await self.send_error('缺少数据库名称参数')
|
||||
return
|
||||
await self.test_database_connection(db_name)
|
||||
elif message_type == 'get_configs':
|
||||
await self.send_database_configs()
|
||||
else:
|
||||
await self.send_error(f'未知的数据库监控命令: {message_type}')
|
||||
|
||||
async def start_monitoring(self, db_name: str):
|
||||
"""开始监控"""
|
||||
if self.is_monitoring:
|
||||
await self.send_message('monitor_status', '数据库监控已在运行')
|
||||
return
|
||||
|
||||
self.current_db_name = db_name
|
||||
self.is_monitoring = True
|
||||
self.monitor_task = asyncio.create_task(self.monitor_loop())
|
||||
await self.send_message('monitor_started', f'开始数据库监控({db_name})，间隔{self.monitor_interval}秒')
|
||||
|
||||
async def stop_monitoring(self):
|
||||
"""停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self.monitor_task = None
|
||||
self.current_db_name = None
|
||||
await self.send_message('monitor_stopped', '数据库监控已停止')
|
||||
|
||||
async def restart_monitoring(self):
|
||||
"""重启监控"""
|
||||
if self.current_db_name:
|
||||
db_name = self.current_db_name
|
||||
await self.stop_monitoring()
|
||||
await asyncio.sleep(0.1) # 短暂延迟
|
||||
await self.start_monitoring(db_name)
|
||||
|
||||
async def monitor_loop(self):
|
||||
"""监控循环"""
|
||||
try:
|
||||
while self.is_monitoring and self.current_db_name:
|
||||
try:
|
||||
await self.send_realtime_stats(self.current_db_name)
|
||||
except Exception as e:
|
||||
logger.error(f"发送数据库实时数据失败: {str(e)}")
|
||||
# 发送错误消息但不停止监控循环
|
||||
try:
|
||||
await self.send_error(f'获取数据库监控数据失败: {str(e)}')
|
||||
except:
|
||||
pass
|
||||
|
||||
# 等待下一次监控间隔
|
||||
await asyncio.sleep(self.monitor_interval)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("数据库监控循环被取消")
|
||||
except Exception as e:
|
||||
logger.error(f"数据库监控循环严重错误: {str(e)}")
|
||||
self.is_monitoring = False
|
||||
|
||||
async def _resolve_db_config(self, db_name: str) -> Optional[Dict[str, Any]]:
|
||||
from app.database import AsyncSessionLocal
|
||||
from core.database_monitor.config_resolver import resolve_monitor_target
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
return await resolve_monitor_target(session, db_name)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
def _create_collector(self, db_config: Dict[str, Any]):
|
||||
from core.database_monitor.config_resolver import create_monitor_collector
|
||||
|
||||
return create_monitor_collector(db_config)
|
||||
|
||||
async def send_database_configs(self):
|
||||
"""发送数据库配置列表"""
|
||||
try:
|
||||
from app.database import AsyncSessionLocal
|
||||
from core.database_monitor.config_resolver import list_monitor_configs
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
configs = await list_monitor_configs(session)
|
||||
await self.send_message('database_configs', '数据库配置列表', configs)
|
||||
except Exception as e:
|
||||
logger.error(f"获取数据库配置失败: {str(e)}")
|
||||
await self.send_error(f'获取数据库配置失败: {str(e)}')
|
||||
|
||||
async def send_database_overview(self, db_name: str):
|
||||
"""发送数据库概览信息"""
|
||||
try:
|
||||
db_config = await self._resolve_db_config(db_name)
|
||||
if not db_config:
|
||||
await self.send_error(f'数据库 {db_name} 未找到')
|
||||
return
|
||||
|
||||
collector = self._create_collector(db_config)
|
||||
overview_data = await collector.get_all_info(db_name, db_config['name'])
|
||||
await self.send_message('database_overview', '数据库概览信息', overview_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取数据库概览失败: {str(e)}")
|
||||
await self.send_error(f'获取数据库概览失败: {str(e)}')
|
||||
|
||||
async def send_realtime_stats(self, db_name: str):
|
||||
"""发送数据库实时统计信息"""
|
||||
try:
|
||||
db_config = await self._resolve_db_config(db_name)
|
||||
if not db_config:
|
||||
await self.send_error(f'数据库 {db_name} 未找到')
|
||||
return
|
||||
|
||||
collector = self._create_collector(db_config)
|
||||
realtime_data = await collector.get_realtime_stats(db_name)
|
||||
await self.send_message('database_realtime', '数据库实时统计', realtime_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取数据库实时统计失败: {str(e)}")
|
||||
await self.send_error(f'获取数据库实时统计失败: {str(e)}')
|
||||
|
||||
async def test_database_connection(self, db_name: str):
|
||||
"""测试数据库连接"""
|
||||
try:
|
||||
db_config = await self._resolve_db_config(db_name)
|
||||
if not db_config:
|
||||
await self.send_error(f'数据库 {db_name} 未找到')
|
||||
return
|
||||
|
||||
collector = self._create_collector(db_config)
|
||||
test_result = await collector.test_connection()
|
||||
await self.send_message('connection_test', '数据库连接测试结果', test_result)
|
||||
except Exception as e:
|
||||
logger.error(f"数据库连接测试失败: {str(e)}")
|
||||
await self.send_error(f'数据库连接测试失败: {str(e)}')
|
||||
@@ -0,0 +1,48 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
通知 WebSocket 消费者
|
||||
"""
|
||||
from typing import Dict, Any
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer, manager
|
||||
|
||||
|
||||
class NotificationConsumer(TokenAuthWebSocketConsumer):
|
||||
"""通知WebSocket消费者"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
super().__init__(websocket)
|
||||
|
||||
async def connect(self):
|
||||
"""连接并加入通知组"""
|
||||
await super().connect()
|
||||
if self.is_authenticated and self.user_id:
|
||||
# 加入用户通知组
|
||||
await manager.group_add(
|
||||
f"notifications_user_{self.user_id}",
|
||||
self.websocket
|
||||
)
|
||||
|
||||
async def disconnect(self, close_code: int = 1000):
|
||||
"""断开连接并离开通知组"""
|
||||
if self.user_id:
|
||||
await manager.group_discard(
|
||||
f"notifications_user_{self.user_id}",
|
||||
self.websocket
|
||||
)
|
||||
await super().disconnect(close_code)
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""处理通知相关消息"""
|
||||
message_type = data.get('type', 'unknown')
|
||||
|
||||
if message_type == 'subscribe':
|
||||
await self.send_message('subscribe_response', '已订阅通知')
|
||||
else:
|
||||
await self.send_message('notification_response', f'通知消息处理: {message_type}')
|
||||
|
||||
async def notification_message(self, event: Dict[str, Any]):
|
||||
"""处理组广播的通知消息"""
|
||||
await self.send_message('notification', event['message'], event.get('data'))
|
||||
@@ -0,0 +1,202 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Redis 监控 WebSocket 消费者
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from app.config import settings
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer, manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RedisMonitorConsumer(TokenAuthWebSocketConsumer):
|
||||
"""Redis监控WebSocket消费者"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
super().__init__(websocket)
|
||||
self.monitor_task: Optional[asyncio.Task] = None
|
||||
self.is_monitoring = False
|
||||
self.monitor_interval = 2 # 固定2秒更新一次
|
||||
|
||||
async def connect(self):
|
||||
"""连接并开始监控"""
|
||||
await super().connect()
|
||||
if self.is_authenticated and self.user_id:
|
||||
# 加入Redis监控组
|
||||
await manager.group_add(
|
||||
"redis_monitor",
|
||||
self.websocket
|
||||
)
|
||||
|
||||
async def disconnect(self, close_code: int = 1000):
|
||||
"""断开连接并停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self.user_id:
|
||||
await manager.group_discard(
|
||||
"redis_monitor",
|
||||
self.websocket
|
||||
)
|
||||
await super().disconnect(close_code)
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""处理Redis监控消息"""
|
||||
message_type = data.get('type', 'unknown')
|
||||
|
||||
if message_type == 'start_monitor':
|
||||
await self.start_monitoring()
|
||||
elif message_type == 'stop_monitor':
|
||||
await self.stop_monitoring()
|
||||
elif message_type == 'get_overview':
|
||||
await self.send_redis_overview()
|
||||
elif message_type == 'get_realtime':
|
||||
await self.send_realtime_stats()
|
||||
elif message_type == 'test_connection':
|
||||
await self.test_redis_connection()
|
||||
else:
|
||||
await self.send_error(f'未知的Redis监控命令: {message_type}')
|
||||
|
||||
async def start_monitoring(self):
|
||||
"""开始监控"""
|
||||
if self.is_monitoring:
|
||||
await self.send_message('monitor_status', 'Redis监控已在运行')
|
||||
return
|
||||
|
||||
self.is_monitoring = True
|
||||
self.monitor_task = asyncio.create_task(self.monitor_loop())
|
||||
await self.send_message('monitor_started', f'开始Redis监控，间隔{self.monitor_interval}秒')
|
||||
|
||||
async def stop_monitoring(self):
|
||||
"""停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self.monitor_task = None
|
||||
await self.send_message('monitor_stopped', 'Redis监控已停止')
|
||||
|
||||
async def restart_monitoring(self):
|
||||
"""重启监控"""
|
||||
await self.stop_monitoring()
|
||||
await asyncio.sleep(0.1) # 短暂延迟
|
||||
await self.start_monitoring()
|
||||
|
||||
async def monitor_loop(self):
|
||||
"""监控循环"""
|
||||
try:
|
||||
while self.is_monitoring:
|
||||
try:
|
||||
await self.send_realtime_stats()
|
||||
except Exception as e:
|
||||
logger.error(f"发送Redis实时数据失败: {str(e)}")
|
||||
# 发送错误消息但不停止监控循环
|
||||
try:
|
||||
await self.send_error(f'获取Redis监控数据失败: {str(e)}')
|
||||
except:
|
||||
pass
|
||||
|
||||
# 等待下一次监控间隔
|
||||
await asyncio.sleep(self.monitor_interval)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Redis监控循环被取消")
|
||||
except Exception as e:
|
||||
logger.error(f"Redis监控循环严重错误: {str(e)}")
|
||||
self.is_monitoring = False
|
||||
|
||||
def _get_redis_collector(self):
|
||||
"""获取Redis信息收集器"""
|
||||
try:
|
||||
from core.redis_monitor import RedisInfoCollector
|
||||
|
||||
# 从配置中获取Redis配置
|
||||
redis_host = settings.REDIS_HOST
|
||||
redis_port = settings.REDIS_PORT
|
||||
redis_password = settings.REDIS_PASSWORD or None
|
||||
redis_db = settings.REDIS_DB
|
||||
|
||||
if redis_password == '':
|
||||
redis_password = None
|
||||
|
||||
return RedisInfoCollector(
|
||||
host=redis_host,
|
||||
port=redis_port,
|
||||
password=redis_password,
|
||||
db=redis_db
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning("RedisInfoCollector not available")
|
||||
return None
|
||||
|
||||
async def send_redis_overview(self):
|
||||
"""发送Redis概览信息"""
|
||||
try:
|
||||
collector = self._get_redis_collector()
|
||||
if collector is None:
|
||||
await self.send_error('Redis监控模块未安装')
|
||||
return
|
||||
|
||||
# 在线程池中执行同步方法
|
||||
loop = asyncio.get_event_loop()
|
||||
overview_data = await loop.run_in_executor(
|
||||
None,
|
||||
collector.get_all_info,
|
||||
'project_redis',
|
||||
'项目Redis'
|
||||
)
|
||||
|
||||
await self.send_message('redis_overview', 'Redis概览信息', overview_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取Redis概览失败: {str(e)}")
|
||||
await self.send_error(f'获取Redis概览失败: {str(e)}')
|
||||
|
||||
async def send_realtime_stats(self):
|
||||
"""发送Redis实时统计信息"""
|
||||
try:
|
||||
collector = self._get_redis_collector()
|
||||
if collector is None:
|
||||
await self.send_error('Redis监控模块未安装')
|
||||
return
|
||||
|
||||
# 在线程池中执行同步方法
|
||||
loop = asyncio.get_event_loop()
|
||||
realtime_data = await loop.run_in_executor(
|
||||
None,
|
||||
collector.get_realtime_stats,
|
||||
'project_redis'
|
||||
)
|
||||
|
||||
await self.send_message('redis_realtime', 'Redis实时统计', realtime_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取Redis实时统计失败: {str(e)}")
|
||||
await self.send_error(f'获取Redis实时统计失败: {str(e)}')
|
||||
|
||||
async def test_redis_connection(self):
|
||||
"""测试Redis连接"""
|
||||
try:
|
||||
collector = self._get_redis_collector()
|
||||
if collector is None:
|
||||
await self.send_error('Redis监控模块未安装')
|
||||
return
|
||||
|
||||
# 在线程池中执行同步方法
|
||||
loop = asyncio.get_event_loop()
|
||||
test_result = await loop.run_in_executor(None, collector.test_connection)
|
||||
|
||||
await self.send_message('connection_test', 'Redis连接测试结果', test_result)
|
||||
except Exception as e:
|
||||
logger.error(f"Redis连接测试失败: {str(e)}")
|
||||
await self.send_error(f'Redis连接测试失败: {str(e)}')
|
||||
@@ -0,0 +1,161 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
服务器监控 WebSocket 消费者
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer, manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ServerMonitorConsumer(TokenAuthWebSocketConsumer):
|
||||
"""服务器监控WebSocket消费者"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
super().__init__(websocket)
|
||||
self.monitor_task: Optional[asyncio.Task] = None
|
||||
self.is_monitoring = False
|
||||
self.monitor_interval = 2 # 固定2秒更新一次
|
||||
# 创建持久的收集器实例以保持缓存数据
|
||||
self.server_collector = None
|
||||
|
||||
def _get_server_collector(self):
|
||||
"""懒加载服务器信息收集器"""
|
||||
if self.server_collector is None:
|
||||
try:
|
||||
from core.server_monitor.server_info import ServerInfoCollector
|
||||
self.server_collector = ServerInfoCollector()
|
||||
except ImportError:
|
||||
logger.warning("ServerInfoCollector not available")
|
||||
return self.server_collector
|
||||
|
||||
async def connect(self):
|
||||
"""连接并开始监控"""
|
||||
await super().connect()
|
||||
if self.is_authenticated and self.user_id:
|
||||
# 加入服务器监控组
|
||||
await manager.group_add(
|
||||
"server_monitor",
|
||||
self.websocket
|
||||
)
|
||||
|
||||
async def disconnect(self, close_code: int = 1000):
|
||||
"""断开连接并停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self.user_id:
|
||||
await manager.group_discard(
|
||||
"server_monitor",
|
||||
self.websocket
|
||||
)
|
||||
await super().disconnect(close_code)
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""处理服务器监控消息"""
|
||||
message_type = data.get('type', 'unknown')
|
||||
|
||||
if message_type == 'start_monitor':
|
||||
await self.start_monitoring()
|
||||
elif message_type == 'stop_monitor':
|
||||
await self.stop_monitoring()
|
||||
elif message_type == 'get_overview':
|
||||
await self.send_server_overview()
|
||||
elif message_type == 'get_realtime':
|
||||
await self.send_realtime_stats()
|
||||
else:
|
||||
await self.send_error(f'未知的监控命令: {message_type}')
|
||||
|
||||
async def start_monitoring(self):
|
||||
"""开始监控"""
|
||||
if self.is_monitoring:
|
||||
await self.send_message('monitor_status', '监控已在运行')
|
||||
return
|
||||
|
||||
self.is_monitoring = True
|
||||
self.monitor_task = asyncio.create_task(self.monitor_loop())
|
||||
await self.send_message('monitor_started', f'开始监控，间隔{self.monitor_interval}秒')
|
||||
|
||||
async def stop_monitoring(self):
|
||||
"""停止监控"""
|
||||
self.is_monitoring = False
|
||||
if self.monitor_task:
|
||||
self.monitor_task.cancel()
|
||||
try:
|
||||
await self.monitor_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self.monitor_task = None
|
||||
await self.send_message('monitor_stopped', '监控已停止')
|
||||
|
||||
async def restart_monitoring(self):
|
||||
"""重启监控"""
|
||||
await self.stop_monitoring()
|
||||
await asyncio.sleep(0.1) # 短暂延迟
|
||||
await self.start_monitoring()
|
||||
|
||||
async def monitor_loop(self):
|
||||
"""监控循环"""
|
||||
try:
|
||||
while self.is_monitoring:
|
||||
try:
|
||||
await self.send_realtime_stats()
|
||||
except Exception as e:
|
||||
logger.error(f"发送实时数据失败: {str(e)}")
|
||||
# 发送错误消息但不停止监控循环
|
||||
try:
|
||||
await self.send_error(f'获取监控数据失败: {str(e)}')
|
||||
except:
|
||||
pass
|
||||
|
||||
# 等待下一次监控间隔
|
||||
await asyncio.sleep(self.monitor_interval)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("监控循环被取消")
|
||||
except Exception as e:
|
||||
logger.error(f"监控循环严重错误: {str(e)}")
|
||||
self.is_monitoring = False
|
||||
|
||||
async def send_server_overview(self):
|
||||
"""发送服务器概览信息"""
|
||||
try:
|
||||
collector = self._get_server_collector()
|
||||
if collector is None:
|
||||
await self.send_error('服务器监控模块未安装')
|
||||
return
|
||||
|
||||
# 在线程池中执行同步方法
|
||||
loop = asyncio.get_event_loop()
|
||||
overview_data = await loop.run_in_executor(None, collector.get_all_info)
|
||||
|
||||
await self.send_message('server_overview', '服务器概览信息', overview_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取服务器概览失败: {str(e)}")
|
||||
await self.send_error(f'获取服务器概览失败: {str(e)}')
|
||||
|
||||
async def send_realtime_stats(self):
|
||||
"""发送实时统计信息"""
|
||||
try:
|
||||
collector = self._get_server_collector()
|
||||
if collector is None:
|
||||
await self.send_error('服务器监控模块未安装')
|
||||
return
|
||||
|
||||
# 在线程池中执行同步方法
|
||||
loop = asyncio.get_event_loop()
|
||||
realtime_data = await loop.run_in_executor(None, collector.get_realtime_stats)
|
||||
|
||||
await self.send_message('realtime_stats', '实时统计信息', realtime_data)
|
||||
except Exception as e:
|
||||
logger.error(f"获取实时统计失败: {str(e)}")
|
||||
await self.send_error(f'获取实时统计失败: {str(e)}')
|
||||
@@ -0,0 +1,42 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
测试 WebSocket 消费者
|
||||
"""
|
||||
import platform
|
||||
from datetime import datetime
|
||||
from typing import Dict, Any
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from core.websocket.consumers.base import TokenAuthWebSocketConsumer
|
||||
|
||||
|
||||
class TestWebSocketConsumer(TokenAuthWebSocketConsumer):
|
||||
"""测试WebSocket消费者"""
|
||||
|
||||
def __init__(self, websocket: WebSocket):
|
||||
super().__init__(websocket)
|
||||
|
||||
async def handle_message(self, data: Dict[str, Any]):
|
||||
"""处理测试消息"""
|
||||
message_type = data.get('type', 'unknown')
|
||||
content = data.get('content', '')
|
||||
|
||||
if message_type == 'echo':
|
||||
await self.send_message('echo_response', f'回声: {content}')
|
||||
elif message_type == 'chat':
|
||||
await self.send_message('chat_response', f'收到聊天消息: {content}', {
|
||||
'user': f'user_{self.user_id}',
|
||||
'original_message': content
|
||||
})
|
||||
elif message_type == 'system_info':
|
||||
# 获取系统信息
|
||||
system_info = {
|
||||
'hostname': platform.node(),
|
||||
'system': platform.system(),
|
||||
'python_version': platform.python_version(),
|
||||
'timestamp': datetime.now().isoformat()
|
||||
}
|
||||
await self.send_message('system_info_response', '系统信息', system_info)
|
||||
else:
|
||||
await self.send_message('unknown_response', f'未知消息类型: {message_type}')
|
||||
@@ -0,0 +1,108 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WebSocket 路由
|
||||
定义 WebSocket 端点
|
||||
"""
|
||||
from fastapi import APIRouter, WebSocket
|
||||
|
||||
from core.websocket.consumers import (
|
||||
TestWebSocketConsumer,
|
||||
NotificationConsumer,
|
||||
ServerMonitorConsumer,
|
||||
RedisMonitorConsumer,
|
||||
DatabaseMonitorConsumer,
|
||||
)
|
||||
from core.chat.ws import ChatConsumer
|
||||
|
||||
router = APIRouter(redirect_slashes=False)
|
||||
|
||||
|
||||
@router.websocket("/ws/test")
|
||||
async def websocket_test(websocket: WebSocket):
|
||||
"""WebSocket测试连接"""
|
||||
consumer = TestWebSocketConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/notifications")
|
||||
async def websocket_notifications(websocket: WebSocket):
|
||||
"""通知推送连接"""
|
||||
consumer = NotificationConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/notification")
|
||||
async def websocket_notification(websocket: WebSocket):
|
||||
"""通知推送连接（兼容路径）"""
|
||||
consumer = NotificationConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/notification/")
|
||||
async def websocket_notification_slash(websocket: WebSocket):
|
||||
"""通知推送连接（带斜杠兼容）"""
|
||||
consumer = NotificationConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/notifications/")
|
||||
async def websocket_notifications_slash(websocket: WebSocket):
|
||||
"""通知推送连接（复数带斜杠兼容）"""
|
||||
consumer = NotificationConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/server-monitor")
|
||||
async def websocket_server_monitor(websocket: WebSocket):
|
||||
"""服务器监控连接"""
|
||||
consumer = ServerMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/server-monitor/")
|
||||
async def websocket_server_monitor_slash(websocket: WebSocket):
|
||||
"""服务器监控连接（带斜杠兼容）"""
|
||||
consumer = ServerMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/redis-monitor")
|
||||
async def websocket_redis_monitor(websocket: WebSocket):
|
||||
"""Redis监控连接"""
|
||||
consumer = RedisMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/redis-monitor/")
|
||||
async def websocket_redis_monitor_slash(websocket: WebSocket):
|
||||
"""Redis监控连接（带斜杠兼容）"""
|
||||
consumer = RedisMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/database-monitor")
|
||||
async def websocket_database_monitor(websocket: WebSocket):
|
||||
"""数据库监控连接"""
|
||||
consumer = DatabaseMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/database-monitor/")
|
||||
async def websocket_database_monitor_slash(websocket: WebSocket):
|
||||
"""数据库监控连接（带斜杠兼容）"""
|
||||
consumer = DatabaseMonitorConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/chat")
|
||||
async def websocket_chat(websocket: WebSocket):
|
||||
"""聊天连接"""
|
||||
consumer = ChatConsumer(websocket)
|
||||
await consumer.run()
|
||||
|
||||
|
||||
@router.websocket("/ws/chat/")
|
||||
async def websocket_chat_slash(websocket: WebSocket):
|
||||
"""聊天连接（带斜杠兼容）"""
|
||||
consumer = ChatConsumer(websocket)
|
||||
await consumer.run()
|
||||
Reference in New Issue
Block a user