"""
工作线程核心管理模块
负责管理工作线程的生命周期，包括创建、扩展和回收。
"""
import logging
import os
import queue
import threading
import time
import uuid
from typing import Any, Callable, Dict, List, Optional
from ..types import TaskTuple
from .adaptive import AdaptiveWorkerManager
[文档]
class WorkerManager:
"""
工作线程管理器
负责管理工作线程的生命周期，包括创建、扩展和回收。
支持动态线程池，根据任务队列大小和CPU使用率自动调整线程数量。
使用自适应策略，在负载高时扩容，空闲时缩容。
线程安全说明：
- 所有对worker_threads列表的操作都在threads_lock保护下进行
- 线程退出时会从列表中安全移除，避免竞态条件
- 自适应扩缩容操作在threads_lock保护下进行
- 使用退出事件确认机制确保线程完全退出
"""
# 配置常量
QUEUE_GET_TIMEOUT = 1 # 队列获取超时时间（秒）
QUEUE_PUT_TIMEOUT = 1 # 队列放入超时时间（秒）
# 默认自适应配置
DEFAULT_CPU_THRESHOLD = 0.8 # 默认CPU阈值
DEFAULT_QUEUE_THRESHOLD_HIGH = 100 # 默认扩容队列阈值
DEFAULT_QUEUE_THRESHOLD_LOW = 10 # 默认缩容队列阈值
DEFAULT_SCALE_UP_COOLDOWN = 5.0 # 默认扩容冷却期
DEFAULT_SCALE_DOWN_COOLDOWN = 30.0 # 默认缩容冷却期
# 实例属性类型声明
_adaptive_manager: Optional["AdaptiveWorkerManager"]
[文档]
def __init__(
self,
logger: logging.Logger,
task_queue: "queue.Queue[TaskTuple]",
worker_threads: List[threading.Thread],
threads_lock: threading.Lock,
running_event: threading.Event,
min_workers: int,
max_workers: int,
idle_timeout: int,
task_timeout: Optional[float],
execute_task_func: Callable[[TaskTuple], None],
adaptive_worker_enabled: bool = True,
cpu_threshold: Optional[float] = None,
queue_threshold_high: Optional[int] = None,
queue_threshold_low: Optional[int] = None,
scale_up_cooldown: Optional[float] = None,
scale_down_cooldown: Optional[float] = None,
use_cpu_monitoring: bool = True,
):
"""
初始化工作线程管理器
Args:
logger: 日志记录器
task_queue: 任务队列
worker_threads: 工作线程列表
threads_lock: 线程锁
running_event: 运行事件
min_workers: 最小工作线程数
max_workers: 最大工作线程数
idle_timeout: 空闲超时时间（秒）
task_timeout: 任务超时时间（秒）
execute_task_func: 任务执行函数
adaptive_worker_enabled: 是否启用自适应扩缩容
cpu_threshold: CPU使用率阈值（可选）
queue_threshold_high: 扩容队列积压阈值（可选）
queue_threshold_low: 缩容队列空闲阈值（可选）
scale_up_cooldown: 扩容冷却期（秒）（可选）
scale_down_cooldown: 缩容冷却期（秒）（可选）
use_cpu_monitoring: 是否启用CPU监控
"""
self.logger = logger
self.task_queue = task_queue
self.worker_threads = worker_threads
self.threads_lock = threads_lock
self._running_event = running_event
self.min_workers = min_workers
self.max_workers = max_workers
self.idle_timeout = idle_timeout
self.task_timeout = task_timeout
self._execute_task = execute_task_func
self.adaptive_worker_enabled = adaptive_worker_enabled
# 线程退出事件跟踪（用于确认线程已完全退出）
self._thread_exit_events: Dict[str, threading.Event] = {}
# 从环境变量读取配置（如果未提供）
self.cpu_threshold = self._load_config_float(
"WORKER_CPU_THRESHOLD", cpu_threshold, self.DEFAULT_CPU_THRESHOLD
)
self.queue_threshold_high = self._load_config_int(
"WORKER_QUEUE_THRESHOLD_HIGH", queue_threshold_high, self.DEFAULT_QUEUE_THRESHOLD_HIGH
)
self.queue_threshold_low = self._load_config_int(
"WORKER_QUEUE_THRESHOLD_LOW", queue_threshold_low, self.DEFAULT_QUEUE_THRESHOLD_LOW
)
self.scale_up_cooldown = self._load_config_float(
"WORKER_SCALE_UP_COOLDOWN", scale_up_cooldown, self.DEFAULT_SCALE_UP_COOLDOWN
)
self.scale_down_cooldown = self._load_config_float(
"WORKER_SCALE_DOWN_COOLDOWN", scale_down_cooldown, self.DEFAULT_SCALE_DOWN_COOLDOWN
)
# 初始化自适应工作线程管理器
if adaptive_worker_enabled:
self._adaptive_manager = AdaptiveWorkerManager(
min_workers=min_workers,
max_workers=max_workers,
cpu_threshold=self.cpu_threshold,
queue_threshold_high=self.queue_threshold_high,
queue_threshold_low=self.queue_threshold_low,
scale_up_cooldown=self.scale_up_cooldown,
scale_down_cooldown=self.scale_down_cooldown,
use_cpu_monitoring=use_cpu_monitoring,
)
else:
self._adaptive_manager = None
# 空闲时间跟踪
self._idle_start_time: Optional[float] = None
def _load_config_float(self, env_key: str, value: Optional[float], default: float) -> float:
"""加载浮点配置"""
if value is not None:
return value
env_value = os.getenv(env_key)
if env_value:
try:
return float(env_value)
except ValueError:
self.logger.warning(f"无效的 {env_key}: {env_value}，使用默认值 {default}")
return default
def _load_config_int(self, env_key: str, value: Optional[int], default: int) -> int:
"""加载整数配置"""
if value is not None:
return value
env_value = os.getenv(env_key)
if env_value:
try:
return int(env_value)
except ValueError:
self.logger.warning(f"无效的 {env_key}: {env_value}，使用默认值 {default}")
return default
[文档]
def start_initial_workers(self) -> None:
"""
启动初始工作线程
根据 min_workers 配置启动最小数量的工作线程。
这些线程会持续运行，不会被空闲超时机制回收。
"""
for _ in range(self.min_workers):
thread = threading.Thread(
target=self._worker_loop,
name=f"TaskWorker-{uuid.uuid4()}",
daemon=True,
)
thread.start()
self.worker_threads.append(thread)
self.logger.info(f"启动初始工作线程数: {len(self.worker_threads)}")
[文档]
def scale_up_workers_if_needed(self) -> None:
"""
根据队列大小和CPU使用率动态扩展工作线程
当队列中的任务数量超过当前线程数时，自动创建新线程。
线程数量不会超过 max_workers 限制。
使用自适应策略，考虑CPU使用率和队列积压情况。
"""
with self.threads_lock:
current_thread_count = len(self.worker_threads)
queue_size = self.task_queue.qsize()
if current_thread_count >= self.max_workers:
return
if self._adaptive_manager is not None:
# 使用自适应策略判断
cpu_usage = self._adaptive_manager.get_cpu_usage()
if self._adaptive_manager.should_scale_up(
current_thread_count, queue_size, cpu_usage
):
self._create_and_start_worker()
else:
# 原始策略：队列积压超过当前线程数时扩容
if queue_size > current_thread_count:
self._create_and_start_worker()
def _create_and_start_worker(self) -> None:
"""
创建并启动新的工作线程
注意：此方法应在 threads_lock 保护下调用，或确保调用者已持有锁。
"""
thread = threading.Thread(
target=self._worker_loop,
name=f"TaskWorker-{uuid.uuid4()}",
daemon=True,
)
# 为线程创建退出事件
exit_event = threading.Event()
self._thread_exit_events[thread.name] = exit_event
thread.start()
self.worker_threads.append(thread)
self.logger.info("启动新工作线程，当前线程数: %d", len(self.worker_threads))
[文档]
def record_task_time(self, task_time: float) -> None:
"""
记录任务执行时间（用于自适应管理）
Args:
task_time: 任务执行时间（秒）
"""
if self._adaptive_manager is not None:
self._adaptive_manager.record_task_time(task_time)
[文档]
def get_idle_time(self) -> Optional[float]:
"""
获取当前队列空闲时间
Returns:
float: 队列空闲时间（秒），如果当前不空闲则返回None
"""
if self._idle_start_time is None:
return None
return time.time() - self._idle_start_time
[文档]
def should_scale_down(self) -> bool:
"""
判断是否应该缩减线程（使用自适应策略）
Returns:
bool: 如果应该缩容则返回True
"""
if self._adaptive_manager is None:
return False
with self.threads_lock:
current_workers = len(self.worker_threads)
queue_size = self.task_queue.qsize()
idle_time = self.get_idle_time()
return self._adaptive_manager.should_scale_down(current_workers, queue_size, idle_time)
[文档]
def get_adaptive_stats(self) -> Optional[Dict[str, Any]]:
"""
获取自适应管理器的统计信息
Returns:
Dict[str, Any]: 统计信息字典，如果自适应管理未启用则返回None
"""
if self._adaptive_manager is not None:
return self._adaptive_manager.get_stats()
return None
def _check_idle_timeout(self, thread_name: str, idle_start: Optional[float]) -> bool:
"""
检查空闲超时，如果超时则尝试退出线程
Args:
thread_name: 线程名称
idle_start: 空闲开始时间
Returns:
bool: 如果应该退出线程则返回True，否则返回False
"""
now = time.time()
if idle_start is None:
return False
idle_duration = now - idle_start
# 使用自适应策略判断是否应该缩容
if self._adaptive_manager is not None:
with self.threads_lock:
current_workers = len(self.worker_threads)
queue_size = self.task_queue.qsize()
if self._adaptive_manager.should_scale_down(
current_workers, queue_size, idle_duration
):
# 执行缩容
current_thread = threading.current_thread()
if len(self.worker_threads) > self.min_workers:
if current_thread in self.worker_threads:
self.worker_threads.remove(current_thread)
self.logger.info(f"空闲线程退出（自适应）: {thread_name}")
return True
else:
self.logger.debug(f"线程 {thread_name} 已达到最小线程数限制，不退出")
return True
return False
# 原始策略：使用 idle_timeout
if idle_duration >= self.idle_timeout:
# 检查是否可以退出（保持最小线程数）
# 在锁内检查并移除，避免竞态条件
with self.threads_lock:
current_thread = threading.current_thread()
# 再次检查线程数，确保在锁内的一致性
if len(self.worker_threads) > self.min_workers:
# 检查当前线程是否仍在列表中
if current_thread in self.worker_threads:
self.worker_threads.remove(current_thread)
self.logger.info(f"空闲线程退出: {thread_name}")
return True
else:
self.logger.debug(f"线程 {thread_name} 已达到最小线程数限制，不退出")
return False
def _worker_loop(self) -> None:
"""
工作线程主循环
从任务队列中获取任务并执行。当空闲时间超过 idle_timeout 且
当前线程数大于 min_workers 时，线程会自动退出以节省资源。
同时记录任务执行时间用于自适应管理。
使用退出事件确认机制，确保线程完全退出后才从列表中移除。
"""
thread_name = threading.current_thread().name
self.logger.debug("工作线程启动: %s", thread_name)
idle_start: Optional[float] = None
try:
while self._running_event.is_set():
try:
task = self.task_queue.get(timeout=self.QUEUE_GET_TIMEOUT)
# 退出信号
if task is None:
self.task_queue.task_done()
break
idle_start = None
task_start_time = time.time()
# 执行任务
self._execute_task(task)
self.task_queue.task_done()
# 记录任务执行时间（用于自适应管理）
task_end_time = time.time()
self.record_task_time(task_end_time - task_start_time)
except queue.Empty:
now = time.time()
if idle_start is None:
idle_start = now
# 更新空闲开始时间（用于自适应管理）
if self._adaptive_manager is not None:
self._idle_start_time = idle_start
elif self._check_idle_timeout(thread_name, idle_start):
break
continue
except KeyboardInterrupt:
# 键盘中断，正常退出
self.logger.info(f"工作线程收到中断信号: {thread_name}")
break
except SystemExit:
# 系统退出，重新抛出，不捕获
raise
except queue.Full:
# 队列满异常，不应该在工作线程中出现，记录警告
self.logger.warning(f"工作线程意外遇到队列满异常: {thread_name}")
except TimeoutError as e:
# 超时异常，记录警告但继续运行
self.logger.warning(f"工作线程任务执行超时: {e}")
except (IOError, OSError, ConnectionError) as e:
# I/O 相关异常，可能是临时性问题，记录警告并继续
self.logger.warning(f"工作线程遇到 I/O 相关异常 [{type(e).__name__}]: {e}")
except (MemoryError, ResourceWarning) as e:
# 资源相关异常，记录错误但尝试继续运行
self.logger.error(f"工作线程遇到资源异常 [{type(e).__name__}]: {e}", exc_info=True)
except (TimeoutError, RuntimeError) as e:
# 超时和运行时错误，记录错误但继续运行
self.logger.error(
f"工作线程遇到运行时异常 [{type(e).__name__}]: {e}", exc_info=True
)
except Exception as e:
# 记录未预期的异常，但不中断线程运行
# 这确保了单个任务的异常不会影响整个线程池的运行
error_type = type(e).__name__
# 注意：这里捕获所有异常是为了保持线程池运行
# 如果是编程错误（TypeError、AttributeError等），应该修复而不是掩盖
if error_type in (
"TypeError",
"AttributeError",
"ValueError",
"KeyError",
"IndexError",
):
# 编程错误应该向上传播，以便开发时发现
self.logger.critical(
f"工作线程遇到编程错误 [{error_type}]: {e}，建议修复代码", exc_info=True
)
# 重新抛出编程错误，以便快速失败
raise
else:
self.logger.error(f"工作线程执行未预期异常 [{error_type}]: {e}", exc_info=True)
finally:
# 设置线程退出事件，通知管理器线程已退出
if thread_name in self._thread_exit_events:
self._thread_exit_events[thread_name].set()
# 清理退出事件引用
del self._thread_exit_events[thread_name]
self.logger.debug("工作线程退出: %s", thread_name)
[文档]
def send_shutdown_signals(self) -> None:
"""
向所有工作线程发送退出信号
通过向任务队列中放入 None 值来通知工作线程退出。
如果队列已满，会尝试等待并重试。
"""
with self.threads_lock:
thread_count = len(self.worker_threads)
# 尝试发送退出信号,如果队列满则等待
signals_sent = 0
for _ in range(thread_count):
try:
# type: ignore[arg-type] # None是特殊的退出信号
self.task_queue.put_nowait(None) # type: ignore[arg-type]
signals_sent += 1
except queue.Full:
# 队列已满,尝试等待并重试
try:
# type: ignore[arg-type] # None是特殊的退出信号
self.task_queue.put(None, timeout=self.QUEUE_PUT_TIMEOUT) # type: ignore[arg-type]
signals_sent += 1
except queue.Full:
self.logger.warning(
f"无法发送退出信号给所有线程，已发送: {signals_sent}/{thread_count}"
)
break
if signals_sent < thread_count:
self.logger.warning(
f"只发送了 {signals_sent}/{thread_count} 个退出信号，"
f"部分线程可能需要等待超时退出"
)
else:
self.logger.info(f"成功发送 {signals_sent} 个退出信号给工作线程")
[文档]
def wait_for_threads_exit(self, join_timeout: int) -> None:
"""
等待所有工作线程退出
在锁内创建线程副本以避免竞态条件，然后逐个等待线程退出。
使用退出事件确认机制，确保线程完全退出。
如果线程在超时时间内未退出，会记录警告但继续执行。
Args:
join_timeout: 线程join超时时间（秒）
"""
# 在锁内创建线程副本，避免竞态条件
with self.threads_lock:
threads_to_wait = list(self.worker_threads)
for thread in threads_to_wait:
if thread.is_alive():
try:
# 等待线程退出事件设置（表示线程已完成清理）
exit_event = self._thread_exit_events.get(thread.name)
if exit_event:
exit_event.wait(timeout=join_timeout)
# 执行线程 join
thread.join(timeout=join_timeout)
if thread.is_alive():
self.logger.warning(f"线程 {thread.name} 在超时后仍未退出")
else:
self.logger.debug(f"线程 {thread.name} 已退出")
except RuntimeError as e:
# RuntimeError可能发生在join已死线程时
self.logger.warning(f"线程 {thread.name} join 失败 [RuntimeError]: {e}")
except Exception as e:
# 其他未预期的异常
error_type = type(e).__name__
self.logger.warning(f"线程 {thread.name} join 失败 [{error_type}]: {e}")
# 清理所有退出事件引用
self._thread_exit_events.clear()