K2cr2O1
d458413e4b
- 新增 `/status` 指令，展示机器人运行状态和系统指标 - 实现Redis Lua脚本支持原子化计数器操作 - 添加消息收发统计功能 - 完善文档，包括插件开发和性能优化指南 - 重构WebSocket连接池，增加健康检查机制 - 移除旧版编译脚本，优化项目结构
247 lines
8.2 KiB
Python
247 lines
8.2 KiB
Python
"""
|
||
WebSocket 连接池模块
|
||
|
||
该模块实现了 WebSocket 连接池功能，用于管理多个 WebSocket 连接，
|
||
提高并发处理能力和连接复用效率。
|
||
"""
|
||
import asyncio
|
||
import websockets
|
||
from websockets.legacy.client import WebSocketClientProtocol
|
||
from typing import Optional, Dict, Any, cast, Union
|
||
import uuid
|
||
from loguru import logger
|
||
|
||
from .config_loader import global_config
|
||
from .utils.exceptions import WebSocketError, WebSocketConnectionError
|
||
|
||
|
||
class WSConnection:
|
||
"""
|
||
WebSocket 连接包装类
|
||
|
||
封装单个 WebSocket 连接的状态和操作
|
||
"""
|
||
def __init__(self, conn: WebSocketClientProtocol, conn_id: str):
|
||
self.conn = conn
|
||
self.conn_id = conn_id
|
||
self.last_used = asyncio.get_event_loop().time()
|
||
self.is_active = True
|
||
self._pending_requests: Dict[str, asyncio.Future] = {}
|
||
|
||
async def send(self, data: Union[Dict[Any, Any], bytes]):
|
||
"""
|
||
发送数据到 WebSocket 连接
|
||
"""
|
||
if not self.is_active:
|
||
raise WebSocketError(f"连接 {self.conn_id} 已关闭")
|
||
|
||
try:
|
||
await self.conn.send(data)
|
||
self.last_used = asyncio.get_event_loop().time()
|
||
except Exception as e:
|
||
self.is_active = False
|
||
raise WebSocketError(f"发送数据失败: {e}")
|
||
|
||
async def recv(self):
|
||
"""
|
||
从 WebSocket 连接接收数据
|
||
"""
|
||
if not self.is_active:
|
||
raise WebSocketError(f"连接 {self.conn_id} 已关闭")
|
||
|
||
try:
|
||
data = await self.conn.recv()
|
||
self.last_used = asyncio.get_event_loop().time()
|
||
return data
|
||
except Exception as e:
|
||
self.is_active = False
|
||
raise WebSocketError(f"接收数据失败: {e}")
|
||
|
||
async def ping(self, timeout: int = 5) -> bool:
|
||
"""
|
||
对 WebSocket 连接执行 ping-pong 健康检查
|
||
"""
|
||
if not self.is_active:
|
||
return False
|
||
try:
|
||
await asyncio.wait_for(self.conn.ping(), timeout=timeout)
|
||
return True
|
||
except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
|
||
self.is_active = False
|
||
return False
|
||
|
||
async def close(self):
|
||
"""
|
||
关闭 WebSocket 连接
|
||
"""
|
||
if self.is_active:
|
||
self.is_active = False
|
||
await self.conn.close()
|
||
|
||
|
||
class WSConnectionPool:
|
||
"""
|
||
WebSocket 连接池
|
||
|
||
管理多个 WebSocket 连接，提供连接的获取、释放和回收功能
|
||
"""
|
||
def __init__(self, pool_size: int = 3, max_idle_time: int = 300):
|
||
"""
|
||
初始化连接池
|
||
|
||
:param pool_size: 连接池大小
|
||
:param max_idle_time: 连接最大空闲时间（秒）
|
||
"""
|
||
self.pool_size = pool_size
|
||
self.max_idle_time = max_idle_time
|
||
self.pool: asyncio.Queue[WSConnection] = asyncio.Queue(maxsize=pool_size)
|
||
self._closed = False
|
||
self._cleanup_task: Optional[asyncio.Task] = None
|
||
|
||
# 从全局配置读取参数
|
||
self.url = global_config.napcat_ws.uri
|
||
self.token = global_config.napcat_ws.token
|
||
self.reconnect_interval = global_config.napcat_ws.reconnect_interval
|
||
|
||
logger.info(f"WebSocket 连接池初始化完成，大小: {pool_size}")
|
||
|
||
async def initialize(self):
|
||
"""
|
||
初始化连接池，创建初始连接
|
||
"""
|
||
if self._closed:
|
||
raise WebSocketError("连接池已关闭")
|
||
|
||
# 启动连接清理任务
|
||
self._cleanup_task = asyncio.create_task(self._cleanup_idle_connections())
|
||
|
||
# 创建初始连接
|
||
for _ in range(self.pool_size):
|
||
try:
|
||
conn = await self._create_connection()
|
||
await self.pool.put(conn)
|
||
logger.info(f"WebSocket 连接 {conn.conn_id} 已创建并加入连接池")
|
||
except Exception as e:
|
||
logger.error(f"创建初始连接失败: {e}")
|
||
|
||
async def _create_connection(self) -> WSConnection:
|
||
"""
|
||
创建新的 WebSocket 连接
|
||
"""
|
||
headers = {"Authorization": f"Bearer {self.token}"} if self.token else {}
|
||
|
||
try:
|
||
conn_id = str(uuid.uuid4())
|
||
websocket_raw = await websockets.connect(
|
||
self.url, additional_headers=headers
|
||
)
|
||
websocket = cast(WebSocketClientProtocol, websocket_raw)
|
||
|
||
conn = WSConnection(websocket, conn_id)
|
||
logger.info(f"WebSocket 连接 {conn_id} 已建立")
|
||
return conn
|
||
except Exception as e:
|
||
raise WebSocketConnectionError(f"创建 WebSocket 连接失败: {e}")
|
||
|
||
async def get_connection(self) -> WSConnection:
|
||
"""
|
||
从连接池获取一个健康的连接，包含健康检查。
|
||
"""
|
||
if self._closed:
|
||
raise WebSocketError("连接池已关闭")
|
||
|
||
try:
|
||
# 尝试从连接池获取连接
|
||
conn = await asyncio.wait_for(self.pool.get(), timeout=5)
|
||
|
||
# 健康检查
|
||
if await conn.ping():
|
||
logger.debug(f"连接 {conn.conn_id} 健康检查通过")
|
||
return conn
|
||
else:
|
||
logger.warning(f"连接 {conn.conn_id} 健康检查失败，丢弃并获取新连接")
|
||
await conn.close()
|
||
return await self.get_connection() # 递归获取下一个
|
||
|
||
except asyncio.TimeoutError:
|
||
# 连接池为空，创建新连接
|
||
logger.warning("连接池在5秒内无可用连接，创建新连接")
|
||
return await self._create_connection()
|
||
except Exception as e:
|
||
raise WebSocketError(f"获取连接时发生未知错误: {e}")
|
||
|
||
async def release_connection(self, conn: WSConnection):
|
||
"""
|
||
释放连接回连接池
|
||
"""
|
||
if self._closed:
|
||
await conn.close()
|
||
return
|
||
|
||
if not conn.is_active:
|
||
logger.warning(f"连接 {conn.conn_id} 已失效，不返回连接池")
|
||
return
|
||
|
||
try:
|
||
if self.pool.full():
|
||
# 连接池已满，关闭该连接
|
||
await conn.close()
|
||
logger.info(f"连接池已满，关闭连接 {conn.conn_id}")
|
||
else:
|
||
await self.pool.put(conn)
|
||
logger.debug(f"连接 {conn.conn_id} 已返回连接池")
|
||
except Exception as e:
|
||
logger.error(f"释放连接失败: {e}")
|
||
await conn.close()
|
||
|
||
async def _cleanup_idle_connections(self):
|
||
"""
|
||
清理空闲连接任务
|
||
"""
|
||
while not self._closed:
|
||
await asyncio.sleep(60) # 每分钟检查一次
|
||
|
||
try:
|
||
# 检查连接池中的连接
|
||
new_pool = asyncio.Queue(maxsize=self.pool_size)
|
||
current_time = asyncio.get_event_loop().time()
|
||
|
||
while not self.pool.empty():
|
||
conn = await self.pool.get()
|
||
|
||
if current_time - conn.last_used > self.max_idle_time:
|
||
# 连接空闲时间过长，关闭
|
||
await conn.close()
|
||
logger.info(f"清理空闲连接 {conn.conn_id}")
|
||
else:
|
||
# 放回新队列
|
||
await new_pool.put(conn)
|
||
|
||
# 替换原连接池
|
||
self.pool = new_pool
|
||
except Exception as e:
|
||
logger.error(f"清理空闲连接失败: {e}")
|
||
|
||
async def close(self):
|
||
"""
|
||
关闭连接池
|
||
"""
|
||
if self._closed:
|
||
return
|
||
|
||
self._closed = True
|
||
|
||
# 停止清理任务
|
||
if self._cleanup_task:
|
||
self._cleanup_task.cancel()
|
||
try:
|
||
await self._cleanup_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
# 关闭所有连接
|
||
while not self.pool.empty():
|
||
conn = await self.pool.get()
|
||
await conn.close()
|
||
|
||
logger.info("WebSocket 连接池已关闭") |