196 lines
7.4 KiB
Python
196 lines
7.4 KiB
Python
import functools
|
||
import inspect
|
||
import logging
|
||
import weakref
|
||
import sys
|
||
from typing import Callable, Optional, Dict, List, Any
|
||
|
||
LOG = logging.getLogger("JobDecorator")
|
||
|
||
# 任务注册表类，管理所有任务和实例
|
||
class TaskRegistry:
|
||
_instance = None
|
||
|
||
@classmethod
|
||
def get_instance(cls):
|
||
if cls._instance is None:
|
||
cls._instance = TaskRegistry()
|
||
return cls._instance
|
||
|
||
def __init__(self):
|
||
if TaskRegistry._instance is not None:
|
||
raise RuntimeError("请使用 TaskRegistry.get_instance() 获取实例")
|
||
|
||
self.tasks = [] # 存储所有被装饰的任务
|
||
self.instances = [] # 存储所有实例的弱引用
|
||
|
||
def register_task(self, task_info):
|
||
"""注册任务信息"""
|
||
self.tasks.append(task_info)
|
||
LOG.info(f"已注册任务: {task_info['name']}, cron: {task_info['cron']}")
|
||
|
||
def register_instance(self, instance):
|
||
"""注册实例到实例列表"""
|
||
# 检查实例是否已注册
|
||
for ref in self.instances:
|
||
if ref() is instance:
|
||
return
|
||
|
||
# 添加实例的弱引用
|
||
self.instances.append(weakref.ref(instance))
|
||
LOG.info(f"已注册实例: {instance.__class__.__name__}")
|
||
|
||
# 清理已失效的弱引用
|
||
self.instances = [ref for ref in self.instances if ref() is not None]
|
||
|
||
# 尝试为这个实例注册任务
|
||
self.register_tasks_for_instance(instance)
|
||
|
||
def register_tasks_for_instance(self, instance):
|
||
"""为指定实例注册任务"""
|
||
class_name = instance.__class__.__name__
|
||
registered_count = 0
|
||
|
||
for task in self.tasks:
|
||
# 检查任务是否属于这个类
|
||
if task["class_name"] == class_name and hasattr(instance, task["method_name"]):
|
||
method = getattr(instance, task["method_name"])
|
||
|
||
# 使用schedule库注册任务
|
||
self._register_with_schedule(instance, method, task)
|
||
registered_count += 1
|
||
|
||
LOG.info(f"为 {class_name} 注册了 {registered_count} 个定时任务")
|
||
return registered_count
|
||
|
||
def _register_with_schedule(self, instance, method, task):
|
||
"""使用schedule库注册任务"""
|
||
import schedule
|
||
|
||
cron = task["cron"]
|
||
name = task["name"]
|
||
LOG.info(f"使用schedule注册任务: {name}, cron: {cron}")
|
||
|
||
# 解析cron表达式
|
||
parts = cron.split()
|
||
if len(parts) < 5:
|
||
LOG.error(f"无效的cron表达式: {cron}")
|
||
return
|
||
|
||
minute, hour, day, month, weekday = parts[:5]
|
||
|
||
# 处理简单的情况
|
||
if minute == "*/1" and hour == "*" and day == "*" and month == "*" and weekday == "*":
|
||
# 每分钟执行
|
||
LOG.info(f"设置每分钟执行任务: {name}")
|
||
schedule.every(1).minutes.do(method)
|
||
elif minute.isdigit() and hour.isdigit() and day == "*" and month == "*":
|
||
# 每天固定时间执行
|
||
time_str = f"{hour.zfill(2)}:{minute.zfill(2)}"
|
||
LOG.info(f"设置每天 {time_str} 执行任务: {name}")
|
||
schedule.every().day.at(time_str).do(method)
|
||
else:
|
||
LOG.warning(f"不支持的cron表达式: {cron}，将使用onEveryTime方法")
|
||
# 尝试使用onEveryTime方法
|
||
if hasattr(instance, 'onEveryTime'):
|
||
if minute.isdigit() and hour.isdigit():
|
||
time_str = f"{hour.zfill(2)}:{minute.zfill(2)}"
|
||
LOG.info(f"使用onEveryTime注册任务: {name}, 时间: {time_str}")
|
||
instance.onEveryTime(time_str, method)
|
||
|
||
def scan_for_instances(self):
|
||
"""扫描已加载的模块，查找所有已创建的类实例"""
|
||
LOG.info("扫描已加载的模块，查找类实例...")
|
||
|
||
# 获取所有已加载的模块
|
||
for module_name, module in sys.modules.items():
|
||
# 跳过内置模块和标准库模块
|
||
if module_name.startswith('_') or not hasattr(module, '__file__') or module.__file__ is None:
|
||
continue
|
||
|
||
# 查找模块中的所有全局变量
|
||
for var_name in dir(module):
|
||
try:
|
||
var = getattr(module, var_name)
|
||
# 检查是否是类实例
|
||
if isinstance(var, object) and not isinstance(var, type) and hasattr(var, '__class__'):
|
||
# 检查类是否有被装饰的方法
|
||
class_name = var.__class__.__name__
|
||
has_decorated_method = False
|
||
|
||
for task in self.tasks:
|
||
if task["class_name"] == class_name and hasattr(var, task["method_name"]):
|
||
has_decorated_method = True
|
||
break
|
||
|
||
if has_decorated_method:
|
||
LOG.info(f"在模块 {module_name} 中找到实例: {class_name}")
|
||
self.register_instance(var)
|
||
except Exception as e:
|
||
# 忽略获取属性时的错误
|
||
pass
|
||
|
||
def initialize(self):
|
||
"""初始化任务系统，扫描所有已创建的实例并注册任务"""
|
||
LOG.info("初始化任务系统")
|
||
|
||
# 扫描已加载的模块，查找所有已创建的类实例
|
||
self.scan_for_instances()
|
||
|
||
# 清理已失效的弱引用
|
||
self.instances = [ref for ref in self.instances if ref() is not None]
|
||
|
||
# 为所有实例注册任务
|
||
for ref in self.instances:
|
||
instance = ref()
|
||
if instance:
|
||
self.register_tasks_for_instance(instance)
|
||
|
||
# 输出已注册的任务信息
|
||
LOG.info(f"已注册 {len(self.tasks)} 个任务")
|
||
for task in self.tasks:
|
||
LOG.info(f"任务: {task['name']}, 方法: {task['method_name']}, cron: {task['cron']}")
|
||
|
||
|
||
# 获取任务注册表实例
|
||
registry = TaskRegistry.get_instance()
|
||
|
||
|
||
def scheduled_job(cron: str, name: Optional[str] = None, enabled: bool = True):
|
||
"""
|
||
定时任务装饰器，用于标记需要定时执行的方法
|
||
|
||
:param cron: cron表达式，例如 "0 0 * * *" 表示每天0点执行一次
|
||
:param name: 任务名称，默认使用方法名
|
||
:param enabled: 是否启用该任务，默认为True
|
||
"""
|
||
|
||
def decorator(func):
|
||
task_name = name or func.__name__
|
||
|
||
# 记录任务信息
|
||
task_info = {
|
||
"func": func,
|
||
"cron": cron,
|
||
"name": task_name,
|
||
"enabled": enabled,
|
||
"method_name": func.__name__,
|
||
"class_name": func.__qualname__.split('.')[0] if '.' in func.__qualname__ else None
|
||
}
|
||
registry.register_task(task_info)
|
||
|
||
@functools.wraps(func)
|
||
def wrapper(*args, **kwargs):
|
||
# 如果是实例方法的第一次调用，注册实例
|
||
if args and not isinstance(args[0], type) and hasattr(args[0], '__class__'):
|
||
registry.register_instance(args[0])
|
||
return func(*args, **kwargs)
|
||
|
||
return wrapper
|
||
|
||
return decorator
|
||
|
||
|
||
def initialize_job_system():
|
||
"""初始化任务系统，扫描所有已创建的实例并注册任务"""
|
||
registry.initialize() |