feat: 添加异步任务管理器和本地存储管理器,重构统计模块

This commit is contained in:
Oct-autumn
2025-05-07 18:20:26 +08:00
parent b2b43c140f
commit 46d15b1fe7
5 changed files with 725 additions and 332 deletions

View File

@@ -0,0 +1,150 @@
from abc import abstractmethod
import asyncio
from asyncio import Task, Event, Lock
from typing import Callable, Dict
from src.common.logger_manager import get_logger
logger = get_logger("async_task_manager")
class AsyncTask:
"""异步任务基类"""
def __init__(self, task_name: str | None = None, wait_before_start: int = 0, run_interval: int = 0):
self.task_name: str = task_name or self.__class__.__name__
"""任务名称"""
self.wait_before_start: int = wait_before_start
"""运行任务前是否进行等待单位设为0则不等待"""
self.run_interval: int = run_interval
"""多次运行的时间间隔单位设为0则仅运行一次"""
@abstractmethod
async def run(self):
"""
任务的执行过程
"""
pass
async def start_task(self, abort_flag: asyncio.Event):
if self.wait_before_start > 0:
# 等待指定时间后开始任务
await asyncio.sleep(self.wait_before_start)
while not abort_flag.is_set():
await self.run()
if self.run_interval > 0:
await asyncio.sleep(self.run_interval)
else:
break
class AsyncTaskManager:
"""异步任务管理器"""
def __init__(self):
self.tasks: Dict[str, Task] = {}
"""任务列表"""
self.abort_flag: Event = Event()
"""是否中止任务标志"""
self._lock: Lock = Lock()
"""异步锁当可能出现await时需要加锁"""
def _remove_task_call_back(self, task: Task):
"""
call_back: 任务完成后移除任务
"""
task_name = task.get_name()
if task_name in self.tasks:
# 任务完成后移除任务
del self.tasks[task_name]
logger.debug(f"已移除任务 '{task_name}'")
else:
logger.warning(f"尝试移除不存在的任务 '{task_name}'")
@staticmethod
def _default_finish_call_back(task: Task):
"""
call_back: 默认的任务完成回调函数
"""
try:
task.result()
logger.debug(f"任务 '{task.get_name()}' 完成")
except asyncio.CancelledError:
logger.debug(f"任务 '{task.get_name()}' 被取消")
except Exception as e:
logger.error(f"任务 '{task.get_name()}' 执行时发生异常: {e}", exc_info=True)
async def add_task(self, task: AsyncTask, call_back: Callable[[asyncio.Task], None] | None = None):
"""
添加任务
"""
if not issubclass(task.__class__, AsyncTask):
raise TypeError(f"task '{task.__class__.__name__}' 必须是继承 AsyncTask 的子类")
with self._lock: # 由于可能需要await等待任务完成所以需要加异步锁
if task.task_name in self.tasks:
logger.warning(f"已存在名称为 '{task.task_name}' 的任务,正在尝试取消并替换")
self.tasks[task.task_name].cancel() # 取消已存在的任务
await self.tasks[task.task_name] # 等待任务完成
logger.info(f"成功结束任务 '{task.task_name}'")
# 创建新任务
task_inst = asyncio.create_task(task.start_task(self.abort_flag))
task_inst.set_name(task.task_name)
task_inst.add_done_callback(self._remove_task_call_back) # 添加完成回调函数-完成任务后自动移除任务
task_inst.add_done_callback(
call_back or self._default_finish_call_back
) # 添加完成回调函数-用户自定义或默认的FallBack
self.tasks[task.task_name] = task_inst # 将任务添加到任务列表
logger.info(f"已启动任务 '{task.task_name}'")
def get_tasks_status(self) -> Dict[str, Dict[str, str]]:
"""
获取所有任务的状态
"""
tasks_status = {}
for task_name, task in self.tasks.items():
tasks_status[task_name] = {
"status": "running" if not task.done() else "done",
}
return tasks_status
async def stop_and_wait_all_tasks(self):
"""
终止所有任务并等待它们完成该方法会阻塞其它尝试add_task()的操作)
"""
with self._lock: # 由于可能需要await等待任务完成所以需要加异步锁
# 设置中止标志
self.abort_flag.set()
# 取消所有任务
for name, inst in self.tasks.items():
try:
inst.cancel()
except asyncio.CancelledError:
logger.info(f"已取消任务 '{name}'")
# 等待所有任务完成
for task_name, task_inst in self.tasks.items():
if not task_inst.done():
try:
await task_inst
except asyncio.CancelledError: # 此处再次捕获取消异常防止stop_all_tasks()时延迟抛出异常
logger.info(f"任务 {task_name} 已取消")
except Exception as e:
logger.error(f"任务 {task_name} 执行时发生异常: {e}", ext_info=True)
# 清空任务列表
self.tasks.clear()
self.abort_flag.clear()
logger.info("所有异步任务已停止")
async_task_manager = AsyncTaskManager()
"""全局异步任务管理器实例"""