K2cr2O1
c851b49db9
重构 WS 类中的连接获取和释放逻辑，使用 connection 上下文管理器确保连接正确释放。 同时改进消息处理循环中的异常处理和连接管理。 refactor(ws_pool): 增强连接池的健壮性和管理能力 1. 添加连接上下文管理器支持 2. 改进连接获取和释放逻辑，增加连接计数和锁保护 3. 优化连接健康检查和清理机制 4. 增强错误处理和日志记录 fix(bot_status): 增加系统信息获取和渲染的错误处理 1. 为系统信息获取添加超时和错误处理 2. 为Redis数据获取添加异常捕获 3. 为图片渲染添加异常处理 4. 改进日志记录和用户反馈
329 lines
12 KiB
Python
329 lines
12 KiB
Python
"""
|
|
WebSocket 连接池模块
|
|
|
|
该模块实现了 WebSocket 连接池功能，用于管理多个 WebSocket 连接，
|
|
提高并发处理能力和连接复用效率。
|
|
"""
|
|
import asyncio
|
|
import websockets
|
|
from websockets.legacy.client import WebSocketClientProtocol
|
|
from typing import Optional, Dict, Any, cast, Union, AsyncGenerator
|
|
import uuid
|
|
from loguru import logger
|
|
import contextlib
|
|
|
|
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:
|
|
# 使用 wait_for 包装 ping
|
|
pong_waiter = await self.conn.ping()
|
|
await asyncio.wait_for(pong_waiter, timeout=timeout)
|
|
return True
|
|
except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed, Exception):
|
|
self.is_active = False
|
|
return False
|
|
|
|
async def close(self):
|
|
"""
|
|
关闭 WebSocket 连接
|
|
"""
|
|
if self.is_active:
|
|
self.is_active = False
|
|
try:
|
|
await self.conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
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._current_size = 0 # 当前管理的连接数（包括池中和借出的）
|
|
self._lock = asyncio.Lock() # 用于保护 _current_size 的修改
|
|
|
|
# 从全局配置读取参数
|
|
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)
|
|
async with self._lock:
|
|
self._current_size += 1
|
|
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}")
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def connection(self) -> AsyncGenerator[WSConnection, None]:
|
|
"""
|
|
获取连接的上下文管理器
|
|
"""
|
|
conn = await self.get_connection()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
await self.release_connection(conn)
|
|
|
|
async def get_connection(self) -> WSConnection:
|
|
"""
|
|
从连接池获取一个健康的连接，包含健康检查。
|
|
"""
|
|
if self._closed:
|
|
raise WebSocketError("连接池已关闭")
|
|
|
|
start_time = asyncio.get_event_loop().time()
|
|
timeout = 10 # 获取连接的总超时时间
|
|
|
|
while True:
|
|
if asyncio.get_event_loop().time() - start_time > timeout:
|
|
raise WebSocketError("获取连接超时")
|
|
|
|
try:
|
|
# 1. 尝试从池中获取
|
|
conn = self.pool.get_nowait()
|
|
|
|
# 健康检查
|
|
if await conn.ping():
|
|
logger.debug(f"连接 {conn.conn_id} 健康检查通过")
|
|
return conn
|
|
else:
|
|
logger.warning(f"连接 {conn.conn_id} 健康检查失败，丢弃")
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
# 继续循环，尝试获取下一个或创建新的
|
|
continue
|
|
|
|
except asyncio.QueueEmpty:
|
|
# 池为空，检查是否可以创建新连接
|
|
async with self._lock:
|
|
if self._current_size < self.pool_size:
|
|
# 有配额，创建新连接
|
|
self._current_size += 1 # 先占位
|
|
create_new = True
|
|
else:
|
|
create_new = False
|
|
|
|
if create_new:
|
|
try:
|
|
conn = await self._create_connection()
|
|
return conn
|
|
except Exception as e:
|
|
async with self._lock:
|
|
self._current_size -= 1 # 回滚占位
|
|
logger.error(f"创建新连接失败: {e}")
|
|
await asyncio.sleep(1) # 避免快速失败循环
|
|
continue
|
|
else:
|
|
# 没有配额，等待池中有可用连接
|
|
try:
|
|
conn = await asyncio.wait_for(self.pool.get(), timeout=1.0)
|
|
# 获取到了，进行健康检查（在下一次循环中处理，或者这里直接处理）
|
|
# 为了代码复用，我们把 conn 放回去（或者直接用），这里直接用
|
|
if await conn.ping():
|
|
return conn
|
|
else:
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
continue
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
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} 已失效，不返回连接池")
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
return
|
|
|
|
try:
|
|
# 尝试放回池中
|
|
self.pool.put_nowait(conn)
|
|
logger.debug(f"连接 {conn.conn_id} 已返回连接池")
|
|
except asyncio.QueueFull:
|
|
# 理论上不应该发生，除非 _current_size 逻辑有误
|
|
logger.warning(f"连接池已满，关闭多余连接 {conn.conn_id}")
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
except Exception as e:
|
|
logger.error(f"释放连接失败: {e}")
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
|
|
async def _cleanup_idle_connections(self):
|
|
"""
|
|
清理空闲连接任务
|
|
"""
|
|
while not self._closed:
|
|
await asyncio.sleep(60) # 每分钟检查一次
|
|
|
|
try:
|
|
# 我们不替换队列，而是取出检查再放回
|
|
# 这样比较安全，但可能会暂时清空池子
|
|
# 更好的做法是只检查队头的连接
|
|
|
|
# 获取当前队列大小
|
|
qsize = self.pool.qsize()
|
|
for _ in range(qsize):
|
|
try:
|
|
conn = self.pool.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
current_time = asyncio.get_event_loop().time()
|
|
if current_time - conn.last_used > self.max_idle_time:
|
|
logger.info(f"清理空闲连接 {conn.conn_id}")
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
else:
|
|
# 还没过期，放回去
|
|
try:
|
|
self.pool.put_nowait(conn)
|
|
except asyncio.QueueFull:
|
|
# 竞争条件下可能满了
|
|
await conn.close()
|
|
async with self._lock:
|
|
self._current_size -= 1
|
|
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():
|
|
try:
|
|
conn = self.pool.get_nowait()
|
|
await conn.close()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
logger.info("WebSocket 连接池已关闭")
|