404 lines
15 KiB
Python
404 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
WebSocket处理器模块
|
|
提供双向实时通信功能
|
|
"""
|
|
|
|
import json
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Dict, Set, Optional, Callable
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
@dataclass
|
|
class WebSocketClient:
|
|
"""WebSocket客户端"""
|
|
client_id: str
|
|
connection: any # WebSocket连接对象
|
|
topics: Set[str] = field(default_factory=set)
|
|
connected_at: float = field(default_factory=time.time)
|
|
last_activity: float = field(default_factory=time.time)
|
|
metadata: dict = field(default_factory=dict)
|
|
|
|
|
|
class WebSocketManager:
|
|
"""WebSocket连接管理器"""
|
|
|
|
# 预定义主题
|
|
TOPIC_MONITOR_CPU = 'monitor:cpu'
|
|
TOPIC_MONITOR_MEMORY = 'monitor:memory'
|
|
TOPIC_MONITOR_DISK = 'monitor:disk'
|
|
TOPIC_MONITOR_NETWORK = 'monitor:network'
|
|
TOPIC_MONITOR_ALL = 'monitor:*'
|
|
TOPIC_SYNC_PROGRESS = 'sync:progress'
|
|
TOPIC_SYNC_STATUS = 'sync:status'
|
|
TOPIC_SYNC_ALL = 'sync:*'
|
|
TOPIC_DOWNLOAD_PROGRESS = 'download:progress'
|
|
TOPIC_DOWNLOAD_STATUS = 'download:status'
|
|
TOPIC_DOWNLOAD_ALL = 'download:*'
|
|
TOPIC_SERVER_STATUS = 'server:status'
|
|
TOPIC_ALL = '*'
|
|
|
|
def __init__(self, config: dict = None):
|
|
self.config = config or {}
|
|
self.clients: Dict[str, WebSocketClient] = {}
|
|
self.lock = threading.Lock()
|
|
self._running = False
|
|
|
|
# 消息队列（用于批量发送）
|
|
self.message_queues: Dict[str, list] = {}
|
|
|
|
# 心跳间隔（秒）
|
|
self.heartbeat_interval = 30
|
|
|
|
# 统计
|
|
self.stats = {
|
|
'total_connections': 0,
|
|
'total_messages_sent': 0,
|
|
'total_messages_received': 0
|
|
}
|
|
|
|
def register_client(self, client_id: str, connection, metadata: dict = None) -> WebSocketClient:
|
|
"""注册新客户端"""
|
|
with self.lock:
|
|
client = WebSocketClient(
|
|
client_id=client_id,
|
|
connection=connection,
|
|
topics={self.TOPIC_ALL}, # 默认订阅所有
|
|
metadata=metadata or {}
|
|
)
|
|
self.clients[client_id] = client
|
|
self.stats['total_connections'] += 1
|
|
|
|
# 启动心跳
|
|
self._start_heartbeat(client_id)
|
|
|
|
return client
|
|
|
|
def unregister_client(self, client_id: str):
|
|
"""注销客户端"""
|
|
with self.lock:
|
|
client = self.clients.pop(client_id, None)
|
|
if client:
|
|
self.message_queues.pop(client_id, None)
|
|
|
|
def subscribe(self, client_id: str, *topics: str):
|
|
"""订阅主题"""
|
|
with self.lock:
|
|
client = self.clients.get(client_id)
|
|
if client:
|
|
client.topics.update(topics)
|
|
|
|
def unsubscribe(self, client_id: str, *topics: str):
|
|
"""取消订阅"""
|
|
with self.lock:
|
|
client = self.clients.get(client_id)
|
|
if client:
|
|
for topic in topics:
|
|
client.topics.discard(topic)
|
|
|
|
def is_subscribed(self, client_id: str, topic: str) -> bool:
|
|
"""检查客户端是否订阅了主题"""
|
|
with self.lock:
|
|
client = self.clients.get(client_id)
|
|
if not client:
|
|
return False
|
|
|
|
if self.TOPIC_ALL in client.topics:
|
|
return True
|
|
|
|
# 检查精确匹配或通配符匹配
|
|
if topic in client.topics:
|
|
return True
|
|
|
|
# 检查通配符匹配
|
|
topic_parts = topic.split(':')
|
|
for subscribed_topic in client.topics:
|
|
if subscribed_topic.endswith('*'):
|
|
prefix = subscribed_topic.rstrip('*').rstrip(':')
|
|
if topic.startswith(prefix):
|
|
return True
|
|
|
|
return False
|
|
|
|
def send_to_client(self, client_id: str, event_type: str, data: dict, callback: Callable = None):
|
|
"""发送消息到指定客户端"""
|
|
with self.lock:
|
|
client = self.clients.get(client_id)
|
|
if not client:
|
|
return False
|
|
|
|
message = self._format_message(event_type, data)
|
|
|
|
try:
|
|
if hasattr(client.connection, 'send'):
|
|
client.connection.send(message)
|
|
self.stats['total_messages_sent'] += 1
|
|
client.last_activity = time.time()
|
|
|
|
if callback:
|
|
callback(client_id, True)
|
|
return True
|
|
else:
|
|
# 放入消息队列
|
|
if client_id not in self.message_queues:
|
|
self.message_queues[client_id] = []
|
|
self.message_queues[client_id].append(message)
|
|
|
|
if callback:
|
|
callback(client_id, True)
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"WebSocket发送失败 {client_id}: {e}")
|
|
if callback:
|
|
callback(client_id, False)
|
|
return False
|
|
|
|
def broadcast(self, event_type: str, data: dict, topic: str = None):
|
|
"""广播消息到所有客户端"""
|
|
with self.lock:
|
|
sent_count = 0
|
|
failed_clients = []
|
|
|
|
for client_id, client in self.clients.items():
|
|
# 检查是否匹配主题
|
|
if topic and not self.is_subscribed(client_id, topic):
|
|
continue
|
|
|
|
message = self._format_message(event_type, data)
|
|
|
|
try:
|
|
if hasattr(client.connection, 'send'):
|
|
client.connection.send(message)
|
|
sent_count += 1
|
|
else:
|
|
if client_id not in self.message_queues:
|
|
self.message_queues[client_id] = []
|
|
self.message_queues[client_id].append(message)
|
|
sent_count += 1
|
|
|
|
client.last_activity = time.time()
|
|
|
|
except Exception as e:
|
|
print(f"广播到 {client_id} 失败: {e}")
|
|
failed_clients.append(client_id)
|
|
|
|
self.stats['total_messages_sent'] += sent_count
|
|
|
|
# 清理失败的客户端
|
|
for client_id in failed_clients:
|
|
self.unregister_client(client_id)
|
|
|
|
return sent_count
|
|
|
|
def broadcast_to_topic(self, topic: str, event_type: str, data: dict):
|
|
"""广播到订阅特定主题的客户端"""
|
|
with self.lock:
|
|
sent_count = 0
|
|
|
|
for client_id, client in self.clients.items():
|
|
if self.is_subscribed(client_id, topic):
|
|
message = self._format_message(event_type, data)
|
|
|
|
try:
|
|
if hasattr(client.connection, 'send'):
|
|
client.connection.send(message)
|
|
sent_count += 1
|
|
else:
|
|
if client_id not in self.message_queues:
|
|
self.message_queues[client_id] = []
|
|
self.message_queues[client_id].append(message)
|
|
sent_count += 1
|
|
|
|
except Exception as e:
|
|
print(f"发送到 {client_id} 失败: {e}")
|
|
|
|
return sent_count
|
|
|
|
def broadcast_monitor_update(self, stats: dict):
|
|
"""广播监控更新"""
|
|
# 提取关键指标
|
|
cpu_percent = stats.get('cpu', {}).get('percent', 0)
|
|
memory_percent = stats.get('memory', {}).get('percent', 0)
|
|
disk_percent = stats.get('disk', {}).get('percent', 0)
|
|
|
|
# 按阈值过滤
|
|
if cpu_percent > 0 or memory_percent > 0 or disk_percent > 0:
|
|
self.broadcast('monitor:stats', stats, 'monitor:*')
|
|
|
|
def broadcast_sync_progress(self, task_id: str, progress: dict):
|
|
"""广播同步进度"""
|
|
self.broadcast('sync:progress', progress, 'sync:*')
|
|
|
|
def broadcast_download_progress(self, download_id: str, progress: dict):
|
|
"""广播下载进度"""
|
|
self.broadcast('download:progress', progress, 'download:*')
|
|
|
|
def get_client_count(self) -> int:
|
|
"""获取客户端数量"""
|
|
with self.lock:
|
|
return len(self.clients)
|
|
|
|
def get_client_topics(self, client_id: str) -> Set[str]:
|
|
"""获取客户端订阅的主题"""
|
|
with self.lock:
|
|
client = self.clients.get(client_id)
|
|
return client.topics.copy() if client else set()
|
|
|
|
def get_all_topics(self) -> Set[str]:
|
|
"""获取所有被订阅的主题"""
|
|
with self.lock:
|
|
topics = set()
|
|
for client in self.clients.values():
|
|
topics.update(client.topics)
|
|
return topics
|
|
|
|
def get_stats(self) -> dict:
|
|
"""获取统计信息"""
|
|
with self.lock:
|
|
return {
|
|
**self.stats,
|
|
'connected_clients': len(self.clients),
|
|
'active_topics': len(self.get_all_topics()),
|
|
'queues_queued': sum(len(q) for q in self.message_queues.values())
|
|
}
|
|
|
|
def _format_message(self, event_type: str, data: dict) -> str:
|
|
"""格式化消息"""
|
|
return json.dumps({
|
|
'type': event_type,
|
|
'timestamp': time.time(),
|
|
'data': data
|
|
}, ensure_ascii=False)
|
|
|
|
def _start_heartbeat(self, client_id: str):
|
|
"""启动心跳"""
|
|
def heartbeat():
|
|
while client_id in self.clients:
|
|
try:
|
|
# 发送心跳
|
|
self.send_to_client(
|
|
client_id,
|
|
'ping',
|
|
{'timestamp': time.time()}
|
|
)
|
|
except Exception:
|
|
break
|
|
|
|
time.sleep(self.heartbeat_interval)
|
|
|
|
thread = threading.Thread(target=heartbeat, daemon=True)
|
|
thread.start()
|
|
|
|
def generate_client_id(self) -> str:
|
|
"""生成客户端ID"""
|
|
return f"ws_{uuid.uuid4().hex[:12]}"
|
|
|
|
def process_messages(self, client_id: str, messages: list):
|
|
"""处理客户端消息"""
|
|
for message in messages:
|
|
self._handle_message(client_id, message)
|
|
|
|
def _handle_message(self, client_id: str, message: str):
|
|
"""处理客户端消息"""
|
|
self.stats['total_messages_received'] += 1
|
|
|
|
try:
|
|
data = json.loads(message)
|
|
|
|
event_type = data.get('type')
|
|
payload = data.get('data', {})
|
|
|
|
if event_type == 'subscribe':
|
|
# 订阅主题
|
|
topics = payload.get('topics', [])
|
|
self.subscribe(client_id, *topics)
|
|
|
|
elif event_type == 'unsubscribe':
|
|
# 取消订阅
|
|
topics = payload.get('topics', [])
|
|
self.unsubscribe(client_id, *topics)
|
|
|
|
elif event_type == 'ping':
|
|
# 心跳响应
|
|
self.send_to_client(client_id, 'pong', {'timestamp': time.time()})
|
|
|
|
elif event_type == 'status':
|
|
# 请求状态 - 返回服务器和客户端状态
|
|
response_data = {'status': 'ok'}
|
|
|
|
if payload.get('monitor'):
|
|
# 获取实时监控数据
|
|
try:
|
|
if handler.monitor:
|
|
response_data['monitor'] = handler.monitor.get_realtime_stats()
|
|
else:
|
|
# 回退到直接使用 psutil
|
|
import psutil
|
|
response_data['monitor'] = {
|
|
'timestamp': datetime.now().isoformat(),
|
|
'cpu': {
|
|
'percent': psutil.cpu_percent(interval=0.1),
|
|
'count': psutil.cpu_count()
|
|
},
|
|
'memory': psutil.virtual_memory()._asdict(),
|
|
'disk': psutil.disk_usage(handler.config.get('base_dir', './downloads'))._asdict()
|
|
}
|
|
except Exception as e:
|
|
response_data['monitor'] = {'error': str(e)}
|
|
|
|
if payload.get('ws'):
|
|
# 获取 WebSocket 统计
|
|
response_data['ws'] = self.get_stats()
|
|
|
|
if payload.get('client'):
|
|
# 获取当前客户端信息
|
|
client = self.clients.get(client_id)
|
|
if client:
|
|
response_data['client'] = {
|
|
'client_id': client.client_id,
|
|
'connected_at': client.connected_at,
|
|
'last_activity': client.last_activity,
|
|
'topics': list(client.topics),
|
|
'metadata': client.metadata
|
|
}
|
|
|
|
self.send_to_client(client_id, 'status:response', response_data)
|
|
|
|
elif event_type == 'sync':
|
|
# 同步相关操作
|
|
sync_action = payload.get('action', 'status')
|
|
|
|
if sync_action == 'status':
|
|
# 获取同步状态
|
|
response_data = {'action': 'status'}
|
|
try:
|
|
from core.sync_scheduler import SyncScheduler
|
|
scheduler = SyncScheduler()
|
|
response_data['sync_status'] = scheduler.get_status() if hasattr(scheduler, 'get_status') else {'message': 'sync scheduler running'}
|
|
except Exception as e:
|
|
response_data['sync_status'] = {'error': str(e)}
|
|
|
|
self.send_to_client(client_id, 'sync:response', response_data)
|
|
|
|
elif sync_action == 'list':
|
|
# 获取同步任务列表
|
|
response_data = {'action': 'list'}
|
|
self.send_to_client(client_id, 'sync:response', response_data)
|
|
|
|
elif sync_action == 'trigger':
|
|
# 触发手动同步
|
|
response_data = {'action': 'trigger', 'status': 'pending'}
|
|
self.send_to_client(client_id, 'sync:response', response_data)
|
|
|
|
except json.JSONDecodeError:
|
|
pass
|
|
except Exception as e:
|
|
print(f"处理WebSocket消息失败: {e}")
|