243 lines
10 KiB
Python
243 lines
10 KiB
Python
import asyncio
|
||
from typing import List, Callable, Dict, Coroutine, Tuple
|
||
from utils.logger import session_logger
|
||
|
||
|
||
class InterruptHandler:
|
||
def __init__(self):
|
||
self.cleanup_handlers: Dict[
|
||
tuple, List[Tuple[int, Callable[..., Coroutine]]]
|
||
] = {}
|
||
self.interrupt_states: Dict[tuple, bool] = {}
|
||
self.DEFAULT_PRIORITY = 100
|
||
self.PRIORITY = {
|
||
"NETWORK_CONNECTION": 10, # 网络连接最优先关闭
|
||
"QUEUE_CLEANUP": 20, # 队列清理的优先级
|
||
"TASK_CANCELLATION": 30, # 任务取消的优先级
|
||
"DEFAULT": 100, # 默认优先级
|
||
}
|
||
self.session_locks: Dict[tuple, asyncio.Lock] = (
|
||
{}
|
||
) # 添加锁字典，用于保护中断状态更新
|
||
self.interrupt_processing: Dict[tuple, bool] = {}
|
||
self.pending_interrupts: Dict[tuple, int] = {}
|
||
|
||
async def get_session_lock(self, session_key: tuple) -> asyncio.Lock:
|
||
if session_key not in self.session_locks:
|
||
self.session_locks[session_key] = asyncio.Lock()
|
||
return self.session_locks[session_key]
|
||
|
||
async def register_session(self, session_key: tuple) -> None:
|
||
lock = await self.get_session_lock(session_key)
|
||
async with lock:
|
||
if session_key not in self.cleanup_handlers:
|
||
self.cleanup_handlers[session_key] = []
|
||
self.interrupt_states[session_key] = False
|
||
self.interrupt_processing[session_key] = False # 初始化处理状态
|
||
self.pending_interrupts[session_key] = 0 # 初始化挂起的中断请求计数
|
||
device_id, session_id = session_key
|
||
session_logger.info(device_id, session_id, "已注册中断处理")
|
||
|
||
async def register_cleanup_handler(
|
||
self,
|
||
session_key: tuple,
|
||
handler: Callable[..., Coroutine],
|
||
priority: int = None,
|
||
*args,
|
||
**kwargs,
|
||
) -> None:
|
||
if session_key not in self.cleanup_handlers:
|
||
await self.register_session(session_key)
|
||
if priority is None:
|
||
priority = self.DEFAULT_PRIORITY
|
||
|
||
async def wrapped_handler():
|
||
try:
|
||
await handler(*args, **kwargs)
|
||
except Exception as e:
|
||
device_id, session_id = session_key
|
||
session_logger.error(
|
||
device_id, session_id, f"执行中断清理处理器出错: {e}"
|
||
)
|
||
|
||
self.cleanup_handlers[session_key].append((priority, wrapped_handler))
|
||
|
||
async def handle_interrupt(self, session_key: tuple) -> None:
|
||
device_id, session_id = session_key
|
||
await self.set_interrupt_state(session_key, True) # 使用异步方法设置状态
|
||
lock = await self.get_session_lock(session_key)
|
||
async with lock:
|
||
if self.interrupt_processing.get(session_key, False):
|
||
self.pending_interrupts[session_key] = (
|
||
self.pending_interrupts.get(session_key, 0) + 1
|
||
)
|
||
session_logger.info(
|
||
device_id,
|
||
session_id,
|
||
f"收到连续中断请求，当前有 {self.pending_interrupts[session_key]} 个中断正在等待",
|
||
)
|
||
return
|
||
self.interrupt_processing[session_key] = True
|
||
handlers_count = len(self.cleanup_handlers.get(session_key, []))
|
||
try:
|
||
session_logger.info(
|
||
device_id, session_id, f"处理中断, 开始执行 {handlers_count} 个清理任务"
|
||
)
|
||
if session_key not in self.cleanup_handlers:
|
||
session_logger.warning(
|
||
device_id, session_id, "未找到该会话的中断处理器"
|
||
)
|
||
return
|
||
handlers = self.cleanup_handlers[session_key]
|
||
if not handlers:
|
||
session_logger.info(device_id, session_id, "该会话没有注册的中断处理器")
|
||
return
|
||
|
||
sorted_handlers = sorted(handlers, key=lambda x: x[0])
|
||
priority_groups = {}
|
||
for priority, handler in sorted_handlers:
|
||
if priority not in priority_groups:
|
||
priority_groups[priority] = []
|
||
priority_groups[priority].append(handler)
|
||
|
||
for priority in sorted(priority_groups.keys()):
|
||
group_handlers = priority_groups[priority]
|
||
group_size = len(group_handlers)
|
||
session_logger.info(
|
||
device_id,
|
||
session_id,
|
||
f"执行优先级 {priority} 的 {group_size} 个清理任务"
|
||
)
|
||
|
||
timeout = min(1.0 + 0.5 * group_size, 3.0)
|
||
from services.task_manager import task_manager
|
||
tasks = []
|
||
for handler in group_handlers:
|
||
task = await task_manager.create_task(
|
||
handler(),
|
||
device_id=device_id,
|
||
session_key=session_key,
|
||
task_type="interrupt_cleanup"
|
||
)
|
||
tasks.append(task)
|
||
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.gather(*tasks, return_exceptions=True),
|
||
timeout=timeout
|
||
)
|
||
except asyncio.TimeoutError:
|
||
session_logger.warning(
|
||
device_id,
|
||
session_id,
|
||
f"优先级 {priority} 的清理任务执行超时({timeout:.1f}秒)"
|
||
)
|
||
|
||
self.cleanup_handlers[session_key] = []
|
||
session_logger.info(device_id, session_id, "所有中断清理任务已完成")
|
||
finally:
|
||
async with lock:
|
||
self.interrupt_processing[session_key] = False
|
||
pending_count = self.pending_interrupts.get(session_key, 0)
|
||
if pending_count > 0:
|
||
self.pending_interrupts[session_key] = 0
|
||
session_logger.info(
|
||
device_id,
|
||
session_id,
|
||
f"有 {pending_count} 个连续中断请求，立即处理",
|
||
)
|
||
await task_manager.create_task(
|
||
self.handle_pending_interrupts(session_key, pending_count),
|
||
device_id=device_id,
|
||
session_key=session_key,
|
||
task_type="pending_interrupt"
|
||
)
|
||
|
||
async def handle_pending_interrupts(self, session_key: tuple, count: int) -> None:
|
||
device_id, session_id = session_key
|
||
try:
|
||
session_logger.info(device_id, session_id, f"处理 {count} 个挂起的中断请求")
|
||
await self.handle_interrupt(session_key)
|
||
except Exception as e:
|
||
session_logger.error(device_id, session_id, f"处理挂起中断请求时出错: {e}")
|
||
|
||
async def remove_session(self, session_key: tuple) -> None:
|
||
lock = await self.get_session_lock(session_key)
|
||
async with lock:
|
||
if session_key in self.cleanup_handlers:
|
||
del self.cleanup_handlers[session_key]
|
||
if session_key in self.interrupt_states:
|
||
del self.interrupt_states[session_key]
|
||
if session_key in self.interrupt_processing:
|
||
del self.interrupt_processing[session_key]
|
||
if session_key in self.pending_interrupts:
|
||
del self.pending_interrupts[session_key]
|
||
if session_key in self.session_locks:
|
||
del self.session_locks[session_key]
|
||
device_id, session_id = session_key
|
||
session_logger.info(device_id, session_id, "已移除会话的中断处理器和状态")
|
||
|
||
def is_interrupted(self, session_key: tuple) -> bool:
|
||
return self.interrupt_states.get(session_key, False)
|
||
|
||
def is_processing_interrupt(self, session_key: tuple) -> bool:
|
||
return self.interrupt_processing.get(session_key, False)
|
||
|
||
async def set_interrupt_state(
|
||
self, session_key: tuple, state: bool
|
||
) -> None:
|
||
if session_key not in self.session_locks:
|
||
await self.register_session(session_key)
|
||
lock = await self.get_session_lock(session_key)
|
||
async with lock:
|
||
old_state = self.interrupt_states.get(session_key, False)
|
||
if old_state != state:
|
||
self.interrupt_states[session_key] = state
|
||
device_id, session_id = session_key
|
||
status = "中断" if state else "正常"
|
||
session_logger.info(device_id, session_id, f"会话状态已设置为{status}")
|
||
else:
|
||
self.interrupt_states[session_key] = state
|
||
|
||
async def notify_client_interrupt_processed(
|
||
self, session_key: tuple, websocket
|
||
) -> bool:
|
||
if not websocket:
|
||
return False
|
||
device_id, session_id = session_key
|
||
max_retries = 3
|
||
retry_count = 0
|
||
while retry_count < max_retries:
|
||
try:
|
||
if websocket.client_state.name == "CONNECTED" and not getattr(
|
||
websocket, "_closed", False
|
||
):
|
||
await websocket.send_text("INTERRUPT_PROCESSED")
|
||
session_logger.info(
|
||
device_id, session_id, "已通知客户端中断处理完成"
|
||
)
|
||
return True
|
||
else:
|
||
session_logger.warning(
|
||
device_id, session_id, "WebSocket连接已关闭，无法通知客户端"
|
||
)
|
||
return False
|
||
except Exception as e:
|
||
retry_count += 1
|
||
if retry_count < max_retries:
|
||
session_logger.warning(
|
||
device_id,
|
||
session_id,
|
||
f"通知客户端失败，尝试重试 ({retry_count}/{max_retries}): {e}",
|
||
)
|
||
await asyncio.sleep(0.2)
|
||
else:
|
||
session_logger.error(
|
||
device_id, session_id, f"通知客户端失败，已达最大重试次数: {e}"
|
||
)
|
||
return False
|
||
return False
|
||
|
||
|
||
interrupt_handler = InterruptHandler()
|