"""
异步任务管理器
纯Python实现的异步任务管理器，支持线程池和动态伸缩。
适用于需要异步执行任务的场景，如后台任务处理、批量数据处理等。
"""
import logging
import os
import queue
import threading
import time
import uuid
from typing import Any, Callable, Dict, List, Optional
from .cleanup import CleanupThreadManager
from .config import ConfigLoader
from .task_status import TaskStatusManager
from .types import TaskStatusDict, TaskTuple
from .worker import TaskExecutor, WorkerManager
[文档]
class TaskQueueFullError(Exception):
"""任务队列已满异常"""
pass
[文档]
class TaskManager:
"""纯Python实现的异步任务管理器（线程池 + 动态伸缩）"""
# 默认配置常量
DEFAULT_QUEUE_SIZE = 1000
DEFAULT_MIN_WORKERS = 1
DEFAULT_IDLE_TIMEOUT = 60 # 秒
DEFAULT_TASK_STATUS_TTL = 3600 # 1小时
DEFAULT_MAX_TASK_STATUS_COUNT = 10000
DEFAULT_CLEANUP_INTERVAL = 300 # 5分钟
DEFAULT_THREAD_JOIN_TIMEOUT = 2 # 秒
DEFAULT_TASK_TIMEOUT = None # 默认无超时限制
# 工作线程配置常量
DEFAULT_MIN_MAX_WORKERS = 4 # 最大工作线程数的最小值
CPU_MULTIPLIER = 4 # CPU核心数的倍数，用于计算最大工作线程数
_instance_lock = threading.Lock()
_instances: Dict[str, "TaskManager"] = {}
# 实例属性类型声明
_instance_key: str
_timers: List["threading.Timer"]
[文档]
def __new__(cls, instance_key: str = "default") -> "TaskManager":
"""
单例模式实现
使用双重检查锁定模式确保线程安全。
支持通过 instance_key 创建多个不同的单例实例。
使用场景：
- 默认情况下，使用 "default" 作为 instance_key，所有调用返回同一个实例
- 如果需要多个独立的任务管理器实例（例如：不同业务模块使用不同的管理器），
可以使用不同的 instance_key，每个 key 对应一个独立的单例实例
示例：
# 获取默认实例
manager1 = TaskManager()
manager2 = TaskManager() # manager1 和 manager2 是同一个实例
# 获取不同业务模块的独立实例
order_manager = TaskManager(instance_key="order")
payment_manager = TaskManager(instance_key="payment") # 独立的实例
Args:
instance_key: 实例键名，默认为 "default"。不同的 key 对应不同的单例实例。
Returns:
TaskManager: 任务管理器实例（单例）
"""
if instance_key not in cls._instances:
with cls._instance_lock:
# 双重检查，避免多个线程同时创建实例
if instance_key not in cls._instances:
# 先创建实例对象
instance = super().__new__(cls)
# 存储 instance_key，用于后续验证
instance._instance_key = instance_key
# 确保实例已创建后再初始化
# 注意：此时实例属性尚未初始化，_init_task_manager会初始化所有属性
instance._init_task_manager()
cls._instances[instance_key] = instance
return cls._instances[instance_key]
[文档]
def __init__(self, instance_key: str = "default") -> None:
"""
初始化方法（单例模式下此方法不会重复执行）
注意：由于使用单例模式，实际的初始化在 __new__ 中通过
_init_task_manager 完成。此方法存在是为了符合Python对象
创建规范，主要用于验证 instance_key 的一致性。
单例模式实现说明：
- Python 的对象创建流程：__new__ -> __init__
- 在单例模式中，__new__ 负责创建或返回已有实例
- 如果实例已存在，Python 仍会调用 __init__，但此时实例已经初始化完成
- 此方法会验证 instance_key 是否匹配，避免混淆
Args:
instance_key: 实例键名，应该与创建实例时使用的 key 一致
"""
# 验证 instance_key 是否匹配（仅在实例已存在时）
#
# 注意：此检查存在 TOCTOU（Time-of-Check to Time-of-Use）窗口，
# 但在当前实现中是安全的：
# 1. _instance_key 在 __new__ 中设置，初始化后不会修改
# 2. 线程安全由 _instance_lock 保证
# 3. 此检查主要用于发现编程错误，不是关键安全检查
if hasattr(self, "_instance_key") and self._instance_key != instance_key:
# 安全访问 logger（可能在单例模式下尚未初始化）
logger = getattr(self, "logger", logging.getLogger(__name__))
logger.warning(
f"警告：尝试使用 instance_key='{instance_key}' 访问实例，"
f"但该实例的实际 instance_key 为 '{self._instance_key}'。"
f"这可能是编程错误，建议检查代码。"
)
def _init_task_manager(self) -> None:
"""初始化任务管理器"""
self.logger = logging.getLogger(__name__)
# 按步骤初始化各个组件
self._init_basic_structures()
config_loader = ConfigLoader(self.logger)
self._init_worker_config()
self._load_configurations(config_loader)
self._init_flags_and_tracking()
self._init_managers()
self._log_initialization_status()
self._start_initial_threads()
def _init_basic_structures(self) -> None:
"""初始化基础数据结构"""
self.task_queue: "queue.Queue[TaskTuple]" = queue.Queue(
maxsize=self.DEFAULT_QUEUE_SIZE
)
self.worker_threads: List[threading.Thread] = []
self.threads_lock = threading.Lock()
def _init_worker_config(self) -> None:
"""初始化工作线程配置"""
self.min_workers = self.DEFAULT_MIN_WORKERS
cpu_count = os.cpu_count() or 1
self.max_workers = max(self.DEFAULT_MIN_MAX_WORKERS, cpu_count * self.CPU_MULTIPLIER)
self.idle_timeout = self.DEFAULT_IDLE_TIMEOUT
def _load_configurations(self, config_loader: ConfigLoader) -> None:
"""加载并验证配置"""
self.task_status_ttl = config_loader.load_int_config(
"TASK_STATUS_TTL",
self.DEFAULT_TASK_STATUS_TTL,
"TASK_STATUS_TTL",
min_value=1,
max_value=config_loader.MAX_TTL,
)
self.max_task_status_count = config_loader.load_int_config(
"MAX_TASK_STATUS_COUNT",
self.DEFAULT_MAX_TASK_STATUS_COUNT,
"MAX_TASK_STATUS_COUNT",
min_value=1,
max_value=config_loader.MAX_TASK_STATUS_COUNT,
)
self.cleanup_interval = config_loader.load_int_config(
"TASK_CLEANUP_INTERVAL",
self.DEFAULT_CLEANUP_INTERVAL,
"TASK_CLEANUP_INTERVAL",
min_value=1,
max_value=config_loader.MAX_CLEANUP_INTERVAL,
)
self.task_timeout = config_loader.load_timeout_config(self.DEFAULT_TASK_TIMEOUT)
def _init_flags_and_tracking(self) -> None:
"""初始化运行标志和跟踪变量"""
self._running_event = threading.Event()
self._running_event.set()
self._timers: List[threading.Timer] = []
self._scale_scheduled = False
self._scale_lock = threading.Lock()
self._scale_check_interval = 0.1
self._scale_scheduled_time = 0.0
def _init_managers(self) -> None:
"""初始化各个管理器"""
self.status_manager = TaskStatusManager(
self.logger,
self.task_status_ttl,
self.max_task_status_count,
)
self.task_executor = TaskExecutor(
self.logger,
lambda: self.task_timeout,
self.status_manager.update_task_status,
)
self.worker_manager = WorkerManager(
self.logger,
self.task_queue,
self.worker_threads,
self.threads_lock,
self._running_event,
self.min_workers,
self.max_workers,
self.idle_timeout,
self.task_timeout,
self.task_executor.execute_task,
)
self.cleanup_manager = CleanupThreadManager(
self.logger,
self._running_event,
self.cleanup_interval,
self.status_manager.cleanup_old_task_status,
)
def _log_initialization_status(self) -> None:
"""记录初始化状态"""
self.logger.info("任务管理器初始化完成")
self.logger.debug(
f"任务管理器配置 - TTL: {self.task_status_ttl}s, "
f"最大状态数: {self.max_task_status_count}, "
f"清理间隔: {self.cleanup_interval}s, "
f"任务超时: {self.task_timeout or '无限制'}s"
)
def _start_initial_threads(self) -> None:
"""启动初始线程"""
self.worker_manager.start_initial_workers()
self.cleanup_manager.start()
def _schedule_scale_check(self) -> None:
"""
调度延迟的扩缩容检查
使用Timer实现延迟检查，避免每次提交任务都触发扩缩容检查，
减少锁竞争，提高高并发场景下的性能。
包含安全机制：如果调度标志长时间未重置（可能由于Timer异常），
会自动重置以防止后续检查被阻塞。
"""
with self._scale_lock:
# 检查是否需要重置过期的调度标志（安全机制）
current_time = time.time()
if self._scale_scheduled:
# 如果调度时间超过 30 秒未重置，强制重置（可能Timer异常）
if current_time - self._scale_scheduled_time > 30:
self.logger.warning(
f"检测到调度标志超时（{(current_time - self._scale_scheduled_time):.1f}秒），"
f"强制重置以防止扩缩容检查阻塞"
)
self._scale_scheduled = False
# 如果已经有一个检查在等待中，跳过
if self._scale_scheduled:
return
self._scale_scheduled = True
self._scale_scheduled_time = current_time
# 在后台线程中执行扩缩容检查
def delayed_scale_check() -> None:
"""
延迟执行的扩缩容检查函数
此函数作为Timer回调执行，负责检查是否需要扩展工作线程。
使用Timer延迟执行可以避免每次提交任务都触发扩缩容检查，
减少锁竞争，提高高并发场景下的性能。
内部异常会被捕获并记录日志，不会传播到外部。
无论执行是否成功，都会在finally块中重置调度标志。
"""
try:
self.worker_manager.scale_up_workers_if_needed()
except (IOError, OSError, ConnectionError) as e:
# I/O 相关异常，记录警告
self.logger.warning(f"扩缩容检查遇到 I/O 异常 [{type(e).__name__}]: {e}")
except (AttributeError, TypeError, ValueError) as e:
# 编程错误，记录严重错误并重新抛出
self.logger.critical(
f"扩缩容检查遇到编程错误 [{type(e).__name__}]: {e}，建议修复代码", exc_info=True
)
raise
except Exception as e:
# 其他未预期的异常
self.logger.error(f"扩缩容检查失败 [{type(e).__name__}]: {e}", exc_info=True)
finally:
with self._scale_lock:
self._scale_scheduled = False
scale_timer = threading.Timer(self._scale_check_interval, delayed_scale_check)
scale_timer.daemon = True
scale_timer.start()
# 跟踪 Timer 以便后续清理
with self._scale_lock:
self._timers.append(scale_timer)
[文档]
def submit_task(
self,
func: Callable[..., Any],
*args: Any,
block: bool = False,
timeout: Optional[float] = None,
**kwargs: Any,
) -> str:
"""
提交任务到任务队列
此方法是线程安全的，可以在多个线程中并发调用。
任务会被添加到队列中，由工作线程异步执行。
Args:
func: 要执行的任务函数
*args: 任务函数的 positional 参数
block: 如果队列已满，是否阻塞等待。默认为 False。
timeout: 阻塞等待的超时时间（秒）。仅在 block=True 时有效。
如果为 None，则无限等待。默认为 None。
**kwargs: 任务函数的关键字参数
Returns:
str: 任务ID（UUID格式的字符串）
Raises:
TaskQueueFullError: 当队列已满且 block=False 时抛出
Note:
- 任务函数应该是线程安全的
- 如果 block=False 且队列已满，会抛出 TaskQueueFullError 异常
- 如果 block=True，会阻塞等待直到队列有空间或超时
- 提交后任务状态为"pending"，执行时变为"running"，完成后变为"completed"或"failed"
Example:
>>> def my_task(name: str, value: int) -> str:
... return f"Task {name} completed with value {value}"
>>>
>>> manager = TaskManager()
>>> # 非阻塞模式提交任务
>>> task_id = manager.submit_task(my_task, "task1", value=100)
>>>
>>> # 阻塞模式提交任务（最多等待10秒）
>>> task_id = manager.submit_task(
... my_task, "task2", value=200, block=True, timeout=10.0
... )
"""
task_id = str(uuid.uuid4())
try:
if block:
self.task_queue.put((task_id, func, args, kwargs), timeout=timeout)
else:
self.task_queue.put_nowait((task_id, func, args, kwargs))
except queue.Full:
error_msg = f"任务队列已满，无法提交任务 {task_id}"
self.logger.error(error_msg)
raise TaskQueueFullError(error_msg)
# 更新任务状态为pending
self.status_manager.update_task_status(task_id, "pending", submit_time=time.time())
# 调度延迟的扩缩容检查，避免每次提交都触发检查
self._schedule_scale_check()
self.logger.debug("已提交任务 %s 到队列", task_id)
return task_id
[文档]
def get_task_status(self, task_id: str) -> Optional[TaskStatusDict]:
"""
获取任务状态
此方法是线程安全的，可以在多个线程中并发调用。
Args:
task_id: 任务ID（由 submit_task 返回的字符串）
Returns:
Optional[TaskStatusDict]: 任务状态字典，包含以下字段：
- status: 任务状态（"pending"、"running"、"completed"、"failed"）
- submit_time: 提交时间（Unix时间戳，可选）
- start_time: 开始执行时间（Unix时间戳，可选）
- end_time: 结束时间（Unix时间戳，可选）
- result: 任务执行结果（仅当status为"completed"时存在）
- error: 错误信息（仅当status为"failed"时存在）
如果任务不存在或已被清理，则返回None
Example:
>>> task_id = manager.submit_task(my_task, "task1")
>>> status = manager.get_task_status(task_id)
>>> if status:
... print(f"任务状态: {status['status']}")
... if status['status'] == 'completed':
... print(f"结果: {status.get('result')}")
"""
# 刷新待处理的批量更新，确保返回最新状态
if hasattr(self.task_executor, "flush_pending_updates"):
self.task_executor.flush_pending_updates()
if (
hasattr(self.status_manager, "_batch_updater")
and self.status_manager._batch_updater is not None
):
self.status_manager._batch_updater.force_flush()
return self.status_manager.get_task_status(task_id)
[文档]
def clear_task_status(self, task_id: Optional[str] = None) -> None:
"""
清除指定任务状态或所有任务状态
此方法是线程安全的，可以在多个线程中并发调用。
Args:
task_id: 要清除的任务ID。如果为None，则清除所有任务状态。
Note:
- 清除任务状态不会影响正在执行的任务
- 清除后，get_task_status将返回None
"""
self.status_manager.clear_task_status(task_id)
[文档]
def shutdown(self) -> None:
"""
关闭任务管理器
优雅关闭流程：
1. 清除运行标志，停止接受新任务
2. 取消所有活动的 Timer
3. 发送退出信号给所有工作线程
4. 等待所有工作线程退出
5. 等待清理线程退出
6. 清理所有资源（线程列表、任务状态等）
注意：如果多次调用，只有第一次调用会生效。
"""
if not self._running_event.is_set():
self.logger.info("任务管理器已经关闭，跳过重复关闭")
return
self._running_event.clear()
self.logger.info("正在关闭任务管理器...")
# 取消所有活动的 Timer
with self._scale_lock:
for timer in self._timers:
if timer.is_alive():
timer.cancel()
self._timers.clear()
# 发送退出信号给所有工作线程
self.worker_manager.send_shutdown_signals()
# 等待所有工作线程退出
self.worker_manager.wait_for_threads_exit(self.DEFAULT_THREAD_JOIN_TIMEOUT)
# 等待清理线程退出
self.cleanup_manager.join(self.DEFAULT_THREAD_JOIN_TIMEOUT)
# 清理资源（在所有线程退出后再清理）
with self.threads_lock:
self.worker_threads.clear()
self.status_manager.clear_task_status()
self.logger.info("任务管理器已关闭")
[文档]
@classmethod
def destroy_instance(cls, instance_key: str = "default") -> bool:
"""
销毁指定的单例实例
清理并移除指定的单例实例，释放相关资源。
适用于需要完全释放 TaskManager 资源的场景。
Args:
instance_key: 要销毁的实例键名，默认为 "default"
Returns:
bool: 如果实例存在并被销毁则返回 True，否则返回 False
Example:
>>> # 销毁默认实例
>>> TaskManager.destroy_instance()
>>>
>>> # 销毁特定实例
>>> TaskManager.destroy_instance("order")
Note:
- 销毁后，再次创建相同 instance_key 的实例会重新初始化
- 如果实例正在运行中，会先调用 shutdown() 清理资源
- 此方法是线程安全的
"""
with cls._instance_lock:
if instance_key in cls._instances:
instance = cls._instances[instance_key]
# 如果实例仍在运行，先关闭
if instance._running_event.is_set():
instance.shutdown()
# 从实例字典中移除
del cls._instances[instance_key]
return True
return False