154 lines
6.0 KiB
Python
154 lines
6.0 KiB
Python
import asyncio
|
|
import weakref
|
|
from typing import Dict, Set, Optional
|
|
from utils.logger import session_logger
|
|
|
|
|
|
class TaskManager:
|
|
"""管理异步任务，防止内存泄露"""
|
|
|
|
def __init__(self):
|
|
self.device_tasks: Dict[str, Set[asyncio.Task]] = {}
|
|
self.session_tasks: Dict[tuple, Set[asyncio.Task]] = {}
|
|
self.cleanup_tasks: Set[asyncio.Task] = set()
|
|
self.lock = asyncio.Lock()
|
|
|
|
async def create_task(self, coro, device_id: Optional[str] = None, session_key: Optional[tuple] = None, task_type: str = "general"):
|
|
"""创建并跟踪异步任务"""
|
|
task = asyncio.create_task(coro)
|
|
|
|
async with self.lock:
|
|
# 按设备ID跟踪
|
|
if device_id:
|
|
if device_id not in self.device_tasks:
|
|
self.device_tasks[device_id] = set()
|
|
self.device_tasks[device_id].add(task)
|
|
|
|
# 按会话跟踪
|
|
if session_key:
|
|
if session_key not in self.session_tasks:
|
|
self.session_tasks[session_key] = set()
|
|
self.session_tasks[session_key].add(task)
|
|
|
|
# 跟踪清理任务
|
|
if task_type == "cleanup":
|
|
self.cleanup_tasks.add(task)
|
|
|
|
# 任务完成后自动清理
|
|
task.add_done_callback(lambda t: asyncio.create_task(self._cleanup_completed_task(t, device_id, session_key, task_type)))
|
|
|
|
session_logger.info(
|
|
device_id or "system",
|
|
session_key[1] if session_key else "task_manager",
|
|
f"创建任务: {task_type}, 总任务数: {len(self.cleanup_tasks) if task_type == 'cleanup' else 'tracked'}"
|
|
)
|
|
|
|
return task
|
|
|
|
async def _cleanup_completed_task(self, task: asyncio.Task, device_id: Optional[str], session_key: Optional[tuple], task_type: str):
|
|
"""清理已完成的任务"""
|
|
async with self.lock:
|
|
if device_id and device_id in self.device_tasks:
|
|
self.device_tasks[device_id].discard(task)
|
|
if not self.device_tasks[device_id]:
|
|
del self.device_tasks[device_id]
|
|
|
|
if session_key and session_key in self.session_tasks:
|
|
self.session_tasks[session_key].discard(task)
|
|
if not self.session_tasks[session_key]:
|
|
del self.session_tasks[session_key]
|
|
|
|
if task_type == "cleanup":
|
|
self.cleanup_tasks.discard(task)
|
|
|
|
async def cancel_device_tasks(self, device_id: str):
|
|
"""取消设备的所有任务"""
|
|
async with self.lock:
|
|
if device_id in self.device_tasks:
|
|
tasks = list(self.device_tasks[device_id])
|
|
session_logger.info(device_id, "task_manager", f"取消设备任务数量: {len(tasks)}")
|
|
|
|
for task in tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
# 等待任务取消完成
|
|
if tasks:
|
|
try:
|
|
await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=5.0)
|
|
except asyncio.TimeoutError:
|
|
session_logger.warning(device_id, "task_manager", "部分任务取消超时")
|
|
|
|
del self.device_tasks[device_id]
|
|
|
|
async def cancel_session_tasks(self, session_key: tuple):
|
|
"""取消会话的所有任务"""
|
|
async with self.lock:
|
|
if session_key in self.session_tasks:
|
|
tasks = list(self.session_tasks[session_key])
|
|
device_id, session_id = session_key
|
|
session_logger.info(device_id, session_id, f"取消会话任务数量: {len(tasks)}")
|
|
|
|
for task in tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
# 等待任务取消完成
|
|
if tasks:
|
|
try:
|
|
await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=3.0)
|
|
except asyncio.TimeoutError:
|
|
session_logger.warning(device_id, session_id, "部分会话任务取消超时")
|
|
|
|
del self.session_tasks[session_key]
|
|
|
|
async def get_task_stats(self) -> dict:
|
|
"""获取任务统计信息"""
|
|
async with self.lock:
|
|
device_count = sum(len(tasks) for tasks in self.device_tasks.values())
|
|
session_count = sum(len(tasks) for tasks in self.session_tasks.values())
|
|
|
|
return {
|
|
"device_tasks": device_count,
|
|
"session_tasks": session_count,
|
|
"cleanup_tasks": len(self.cleanup_tasks),
|
|
"total_tracked_devices": len(self.device_tasks),
|
|
"total_tracked_sessions": len(self.session_tasks)
|
|
}
|
|
|
|
async def cleanup_all_tasks(self):
|
|
"""清理所有任务"""
|
|
async with self.lock:
|
|
all_tasks = []
|
|
|
|
# 收集所有任务
|
|
for tasks in self.device_tasks.values():
|
|
all_tasks.extend(tasks)
|
|
for tasks in self.session_tasks.values():
|
|
all_tasks.extend(tasks)
|
|
all_tasks.extend(self.cleanup_tasks)
|
|
|
|
session_logger.info("system", "task_manager", f"开始清理所有任务，总数: {len(all_tasks)}")
|
|
|
|
# 取消所有任务
|
|
for task in all_tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
|
|
# 等待任务取消完成
|
|
if all_tasks:
|
|
try:
|
|
await asyncio.wait_for(asyncio.gather(*all_tasks, return_exceptions=True), timeout=10.0)
|
|
except asyncio.TimeoutError:
|
|
session_logger.warning("system", "task_manager", "部分任务清理超时")
|
|
|
|
# 清空所有跟踪
|
|
self.device_tasks.clear()
|
|
self.session_tasks.clear()
|
|
self.cleanup_tasks.clear()
|
|
|
|
session_logger.info("system", "task_manager", "任务管理器清理完成")
|
|
|
|
|
|
# 全局任务管理器实例
|
|
task_manager = TaskManager() |