feat: 添加异步任务管理器和本地存储管理器,重构统计模块
This commit is contained in:
150
src/manager/async_task_manager.py
Normal file
150
src/manager/async_task_manager.py
Normal 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()
|
||||
"""全局异步任务管理器实例"""
|
||||
67
src/manager/local_store_manager.py
Normal file
67
src/manager/local_store_manager.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from src.common.logger_manager import get_logger
|
||||
|
||||
LOCAL_STORE_FILE_PATH = "data/local_store.json"
|
||||
|
||||
logger = get_logger("local_storage")
|
||||
|
||||
|
||||
class LocalStoreManager:
|
||||
file_path: str
|
||||
"""本地存储路径"""
|
||||
|
||||
store: dict[str, str | list | dict | int | float | bool]
|
||||
"""本地存储数据"""
|
||||
|
||||
def __init__(self, local_store_path: str | None = None):
|
||||
self.file_path = local_store_path or LOCAL_STORE_FILE_PATH
|
||||
self.store = {}
|
||||
self.load_local_store()
|
||||
|
||||
def __getitem__(self, item: str) -> str | list | dict | int | float | bool | None:
|
||||
"""获取本地存储数据"""
|
||||
return self.store.get(item, None)
|
||||
|
||||
def __setitem__(self, key: str, value: str | list | dict | int | float | bool):
|
||||
"""设置本地存储数据"""
|
||||
self.store[key] = value
|
||||
self.save_local_store()
|
||||
|
||||
def __contains__(self, item: str) -> bool:
|
||||
"""检查本地存储数据是否存在"""
|
||||
return item in self.store
|
||||
|
||||
def load_local_store(self):
|
||||
"""加载本地存储数据"""
|
||||
if os.path.exists(self.file_path):
|
||||
# 存在本地存储文件,加载数据
|
||||
logger.info("正在阅读记事本......我在看,我真的在看!")
|
||||
logger.debug(f"加载本地存储数据: {self.file_path}")
|
||||
try:
|
||||
with open(self.file_path, "r", encoding="utf-8") as f:
|
||||
self.store = json.load(f)
|
||||
logger.success("全都记起来了!")
|
||||
except json.JSONDecodeError:
|
||||
logger.warning("啊咧?记事本被弄脏了,正在重建记事本......")
|
||||
self.store = {}
|
||||
with open(self.file_path, "w", encoding="utf-8") as f:
|
||||
json.dump({}, f, ensure_ascii=False, indent=4)
|
||||
logger.success("记事本重建成功!")
|
||||
else:
|
||||
# 不存在本地存储文件,创建新的目录和文件
|
||||
logger.warning("啊咧?记事本不存在,正在创建新的记事本......")
|
||||
os.makedirs(os.path.dirname(self.file_path), exist_ok=True)
|
||||
with open(self.file_path, "w", encoding="utf-8") as f:
|
||||
json.dump({}, f, ensure_ascii=False, indent=4)
|
||||
logger.success("记事本创建成功!")
|
||||
|
||||
def save_local_store(self):
|
||||
"""保存本地存储数据"""
|
||||
logger.debug(f"保存本地存储数据: {self.file_path}")
|
||||
with open(self.file_path, "w", encoding="utf-8") as f:
|
||||
json.dump(self.store, f, ensure_ascii=False, indent=4)
|
||||
|
||||
|
||||
local_storage = LocalStoreManager("data/local_store.json") # 全局单例化
|
||||
Reference in New Issue
Block a user