#!/usr/bin/env python3 # -*- coding: utf-8 -*- """截图线程池管理 - 工作线程池模式（并发执行截图任务）""" import os import threading import queue import time from typing import Callable, Optional, Dict, Any # 安全修复: 将魔法数字提取为可配置常量 BROWSER_IDLE_TIMEOUT = int(os.environ.get("BROWSER_IDLE_TIMEOUT", "300")) # 空闲超时(秒)，默认5分钟 TASK_QUEUE_TIMEOUT = int(os.environ.get("TASK_QUEUE_TIMEOUT", "10")) # 队列获取超时(秒) TASK_QUEUE_MAXSIZE = int(os.environ.get("BROWSER_TASK_QUEUE_MAXSIZE", "200")) # 队列最大长度(0表示无限制) BROWSER_MAX_USE_COUNT = int(os.environ.get("BROWSER_MAX_USE_COUNT", "0")) # 每个执行环境最大复用次数(0表示不限制) # 新增：自适应资源配置 ADAPTIVE_CONFIG = os.environ.get("BROWSER_ADAPTIVE_CONFIG", "1").strip().lower() in ("1", "true", "yes", "on") LOAD_HISTORY_SIZE = 50 # 负载历史记录大小 class AdaptiveResourceManager: """自适应资源管理器""" def __init__(self): self._load_history = [] self._current_load = 0 self._last_adjustment = 0 self._adjustment_cooldown = 30 # 调整冷却时间30秒 def record_task_interval(self, interval: float): """记录任务间隔，更新负载历史""" if len(self._load_history) >= LOAD_HISTORY_SIZE: self._load_history.pop(0) self._load_history.append(interval) # 计算当前负载 if len(self._load_history) >= 2: recent_intervals = self._load_history[-10:] # 最近10个任务 avg_interval = sum(recent_intervals) / len(recent_intervals) # 负载越高，间隔越短 self._current_load = 1.0 / max(avg_interval, 0.1) def should_adjust_timeout(self) -> bool: """判断是否应该调整超时配置""" if not ADAPTIVE_CONFIG: return False current_time = time.time() if current_time - self._last_adjustment < self._adjustment_cooldown: return False return len(self._load_history) >= 10 # 至少需要10个数据点 def calculate_optimal_idle_timeout(self) -> int: """基于历史负载计算最优空闲超时""" if not self._load_history: return BROWSER_IDLE_TIMEOUT # 计算最近任务间隔的平均值 recent_intervals = self._load_history[-20:] # 最近20个任务 if len(recent_intervals) < 2: return BROWSER_IDLE_TIMEOUT avg_interval = sum(recent_intervals) / len(recent_intervals) # 根据负载动态调整超时 # 高负载时缩短超时，低负载时延长超时 if self._current_load > 2.0: # 高负载 optimal_timeout = min(avg_interval * 1.5, 600) # 最多10分钟 elif self._current_load < 0.5: # 低负载 optimal_timeout = min(avg_interval * 3.0, 1800) # 最多30分钟 else: # 正常负载 optimal_timeout = min(avg_interval * 2.0, 900) # 最多15分钟 return max(int(optimal_timeout), 60) # 最少1分钟 def get_optimal_queue_timeout(self) -> int: """获取最优队列超时""" if not self._load_history: return TASK_QUEUE_TIMEOUT # 根据任务频率调整队列超时 if self._current_load > 2.0: # 高负载时减少等待 return max(TASK_QUEUE_TIMEOUT // 2, 3) elif self._current_load < 0.5: # 低负载时可以增加等待 return min(TASK_QUEUE_TIMEOUT * 2, 30) else: return TASK_QUEUE_TIMEOUT def record_adjustment(self): """记录一次调整操作""" self._last_adjustment = time.time() class BrowserWorker(threading.Thread): """截图工作线程 - 每个worker维护自己的执行环境""" def __init__( self, worker_id: int, task_queue: queue.Queue, log_callback: Optional[Callable] = None, pre_warm: bool = False, ): super().__init__(daemon=True) self.worker_id = worker_id self.task_queue = task_queue self.log_callback = log_callback self.browser_instance = None self.running = True self.idle = True self.total_tasks = 0 self.failed_tasks = 0 self.pre_warm = pre_warm self.last_activity_ts = 0.0 self.task_start_time = 0.0 # 初始化自适应资源管理器 if ADAPTIVE_CONFIG: self._adaptive_mgr = AdaptiveResourceManager() else: self._adaptive_mgr = None def log(self, message: str): """日志输出""" if self.log_callback: self.log_callback(f"[Worker-{self.worker_id}] {message}") else: print(f"[截图池][Worker-{self.worker_id}] {message}") def _create_browser(self): """创建截图执行环境（逻辑占位，无需真实浏览器）""" created_at = time.time() self.browser_instance = { "created_at": created_at, "use_count": 0, "worker_id": self.worker_id, } self.last_activity_ts = created_at self.log("截图执行环境就绪") return True def _close_browser(self): """关闭截图执行环境""" if self.browser_instance: self.log(f"执行环境已释放（共处理{self.browser_instance.get('use_count', 0)}个任务）") self.browser_instance = None def _check_browser_health(self) -> bool: """检查执行环境是否就绪""" return bool(self.browser_instance) def _ensure_browser(self) -> bool: """确保执行环境可用""" if self._check_browser_health(): return True self.log("执行环境不可用，尝试重新创建...") self._close_browser() return self._create_browser() def run(self): """工作线程主循环 - 按需启动执行环境模式""" if self.pre_warm: self.log("Worker启动（预热模式，启动即准备执行环境）") else: self.log("Worker启动（按需模式，等待任务时不占用资源）") if self.pre_warm and not self.browser_instance: self._create_browser() self.pre_warm = False while self.running: try: # 允许运行中触发预热（例如池在初始化后调用 warmup） if self.pre_warm and not self.browser_instance: self._create_browser() self.pre_warm = False # 从队列获取任务（带超时，以便能响应停止信号和空闲检查） self.idle = True # 使用自适应队列超时 queue_timeout = ( self._adaptive_mgr.get_optimal_queue_timeout() if self._adaptive_mgr else TASK_QUEUE_TIMEOUT ) try: task = self.task_queue.get(timeout=queue_timeout) except queue.Empty: # 检查是否需要释放空闲的执行环境 if self.browser_instance and self.last_activity_ts > 0: idle_time = time.time() - self.last_activity_ts # 使用自适应空闲超时 optimal_timeout = ( self._adaptive_mgr.calculate_optimal_idle_timeout() if self._adaptive_mgr else BROWSER_IDLE_TIMEOUT ) if idle_time > optimal_timeout: self.log(f"空闲{int(idle_time)}秒（优化超时:{optimal_timeout}秒），释放执行环境") self._close_browser() continue self.idle = False if task is None: # None作为停止信号 self.log("收到停止信号") break # 按需创建或确保执行环境可用 browser_ready = False for attempt in range(2): if self._ensure_browser(): browser_ready = True break if attempt < 1: self.log("执行环境创建失败，重试...") time.sleep(0.5) if not browser_ready: retry_count = int(task.get("retry_count", 0) or 0) if isinstance(task, dict) else 0 if retry_count < 1 and isinstance(task, dict): task["retry_count"] = retry_count + 1 try: self.task_queue.put(task, timeout=1) self.log("执行环境不可用，任务重新入队") except queue.Full: self.log("任务队列已满，无法重新入队，任务失败") callback = task.get("callback") if callable(callback): callback(None, "执行环境不可用") self.total_tasks += 1 self.failed_tasks += 1 continue self.log("执行环境不可用，任务失败") callback = task.get("callback") if isinstance(task, dict) else None if callable(callback): callback(None, "执行环境不可用") self.total_tasks += 1 self.failed_tasks += 1 continue # 执行任务 task_func = task.get("func") task_args = task.get("args", ()) task_kwargs = task.get("kwargs", {}) callback = task.get("callback") self.total_tasks += 1 # 确保browser_instance存在后再访问 if self.browser_instance is None: self.log("执行环境不可用，任务失败") if callable(callback): callback(None, "执行环境不可用") self.failed_tasks += 1 continue self.browser_instance["use_count"] += 1 self.log(f"开始执行任务（第{self.browser_instance['use_count']}次执行）") # 记录任务开始时间 task_start_time = time.time() try: # 将执行环境实例传递给任务函数 result = task_func(self.browser_instance, *task_args, **task_kwargs) callback(result, None) self.log(f"任务执行成功") # 记录任务完成并更新负载历史 task_end_time = time.time() task_interval = task_end_time - task_start_time if self._adaptive_mgr: self._adaptive_mgr.record_task_interval(task_interval) self.last_activity_ts = time.time() except Exception as e: self.log(f"任务执行失败: {e}") callback(None, str(e)) self.failed_tasks += 1 self.last_activity_ts = time.time() # 任务失败后，检查执行环境健康 if not self._check_browser_health(): self.log("任务失败导致执行环境异常，将在下次任务前重建") self._close_browser() # 定期重启执行环境，释放可能累积的资源 if self.browser_instance and BROWSER_MAX_USE_COUNT > 0: if self.browser_instance.get("use_count", 0) >= BROWSER_MAX_USE_COUNT: self.log(f"执行环境已复用{self.browser_instance['use_count']}次，重启释放资源") self._close_browser() except Exception as e: self.log(f"Worker出错: {e}") time.sleep(1) # 清理资源 self._close_browser() self.log(f"Worker停止（总任务:{self.total_tasks}, 失败:{self.failed_tasks}）") def stop(self): """停止worker""" self.running = False class BrowserWorkerPool: """截图工作线程池""" def __init__(self, pool_size: int = 3, log_callback: Optional[Callable] = None): self.pool_size = pool_size self.log_callback = log_callback maxsize = TASK_QUEUE_MAXSIZE if TASK_QUEUE_MAXSIZE > 0 else 0 self.task_queue = queue.Queue(maxsize=maxsize) self.workers = [] self.initialized = False self.lock = threading.Lock() def log(self, message: str): """日志输出""" if self.log_callback: self.log_callback(message) else: print(f"[截图池] {message}") def initialize(self): """初始化工作线程池（按需模式，默认预热1个执行环境）""" with self.lock: if self.initialized: return self.log(f"正在初始化截图线程池（{self.pool_size}个worker，按需启动执行环境）...") for i in range(self.pool_size): worker = BrowserWorker( worker_id=i + 1, task_queue=self.task_queue, log_callback=self.log_callback, pre_warm=(i < 1), ) worker.start() self.workers.append(worker) self.initialized = True self.log(f"[OK] 截图线程池初始化完成（{self.pool_size}个worker就绪，执行环境将在有任务时按需启动）") # 初始化完成后，默认预热1个执行环境，降低容器重启后前几批任务的冷启动开销 self.warmup(1) def warmup(self, count: int = 1) -> int: """预热截图线程池 - 预创建指定数量的执行环境""" if count <= 0: return 0 if not self.initialized: self.log("警告：线程池未初始化，无法预热") return 0 with self.lock: target_workers = list(self.workers[: min(count, len(self.workers))]) self.log(f"预热截图线程池（预创建{len(target_workers)}个执行环境）...") for worker in target_workers: if not worker.browser_instance: worker.pre_warm = True # 等待预热完成（最多等待20秒，避免阻塞过久） deadline = time.time() + 20 while time.time() < deadline: warmed = sum(1 for w in target_workers if w.browser_instance) if warmed >= len(target_workers): break time.sleep(0.1) warmed = sum(1 for w in target_workers if w.browser_instance) self.log(f"[OK] 截图线程池预热完成（{warmed}个执行环境就绪）") return warmed def submit_task(self, task_func: Callable, callback: Callable, *args, **kwargs) -> bool: """ 提交任务到队列 Args: task_func: 任务函数，签名为 func(browser_instance, *args, **kwargs) callback: 回调函数，签名为 callback(result, error) *args, **kwargs: 传递给task_func的参数 Returns: 是否成功提交 """ if not self.initialized: self.log("警告：线程池未初始化") return False task = { "func": task_func, "args": args, "kwargs": kwargs, "callback": callback, "retry_count": 0, } try: self.task_queue.put(task, timeout=1) return True except queue.Full: self.log(f"警告：任务队列已满（maxsize={self.task_queue.maxsize}），拒绝提交任务") return False def get_stats(self) -> Dict[str, Any]: """获取线程池统计信息""" workers = list(self.workers or []) idle_count = sum(1 for w in workers if getattr(w, "idle", False)) total_tasks = sum(int(getattr(w, "total_tasks", 0) or 0) for w in workers) failed_tasks = sum(int(getattr(w, "failed_tasks", 0) or 0) for w in workers) worker_details = [] for w in workers: browser_instance = getattr(w, "browser_instance", None) browser_use_count = 0 browser_created_at = None if isinstance(browser_instance, dict): browser_use_count = int(browser_instance.get("use_count", 0) or 0) browser_created_at = browser_instance.get("created_at") worker_details.append( { "worker_id": getattr(w, "worker_id", None), "idle": bool(getattr(w, "idle", False)), "has_browser": bool(browser_instance), "total_tasks": int(getattr(w, "total_tasks", 0) or 0), "failed_tasks": int(getattr(w, "failed_tasks", 0) or 0), "browser_use_count": browser_use_count, "browser_created_at": browser_created_at, "last_active_ts": float(getattr(w, "last_activity_ts", 0) or 0), "thread_alive": bool(getattr(w, "is_alive", lambda: False)()), } ) return { "pool_size": self.pool_size, "idle_workers": idle_count, "busy_workers": max(0, len(workers) - idle_count), "queue_size": self.task_queue.qsize(), "total_tasks": total_tasks, "failed_tasks": failed_tasks, "success_rate": f"{(total_tasks - failed_tasks) / total_tasks * 100:.1f}%" if total_tasks > 0 else "N/A", "workers": worker_details, "timestamp": time.time(), } def wait_for_completion(self, timeout: Optional[float] = None): """等待所有任务完成""" start_time = time.time() while not self.task_queue.empty(): if timeout and (time.time() - start_time) > timeout: self.log("等待超时") return False time.sleep(0.5) # 再等待一下确保正在执行的任务完成 time.sleep(2) return True def shutdown(self): """关闭线程池""" self.log("正在关闭工作线程池...") # 发送停止信号 for _ in self.workers: self.task_queue.put(None) # 等待所有worker停止 for worker in self.workers: worker.join(timeout=10) self.workers.clear() self.initialized = False self.log("[OK] 工作线程池已关闭") # 全局实例 _global_pool: Optional[BrowserWorkerPool] = None _pool_lock = threading.Lock() def get_browser_worker_pool(pool_size: int = 3, log_callback: Optional[Callable] = None) -> BrowserWorkerPool: """获取全局截图工作线程池（单例）""" global _global_pool with _pool_lock: if _global_pool is None: _global_pool = BrowserWorkerPool(pool_size=pool_size, log_callback=log_callback) _global_pool.initialize() return _global_pool def init_browser_worker_pool(pool_size: int = 3, log_callback: Optional[Callable] = None): """初始化全局截图工作线程池""" get_browser_worker_pool(pool_size=pool_size, log_callback=log_callback) def _shutdown_pool_when_idle(pool: BrowserWorkerPool) -> None: try: pool.wait_for_completion(timeout=60) except Exception: pass try: pool.shutdown() except Exception: pass def resize_browser_worker_pool(pool_size: int, log_callback: Optional[Callable] = None) -> bool: """调整截图线程池并发（新任务走新池，旧池空闲后自动关闭）""" global _global_pool try: target_size = max(1, int(pool_size)) except Exception: target_size = 1 with _pool_lock: old_pool = _global_pool if old_pool and int(getattr(old_pool, "pool_size", 0) or 0) == target_size: return False effective_log_callback = log_callback or (getattr(old_pool, "log_callback", None) if old_pool else None) _global_pool = BrowserWorkerPool(pool_size=target_size, log_callback=effective_log_callback) _global_pool.initialize() if old_pool: threading.Thread(target=_shutdown_pool_when_idle, args=(old_pool,), daemon=True).start() return True def shutdown_browser_worker_pool(): """关闭全局截图工作线程池""" global _global_pool with _pool_lock: if _global_pool: _global_pool.shutdown() _global_pool = None if __name__ == "__main__": # 测试代码 print("测试截图工作线程池...") def test_task(browser_instance, url: str, task_id: int): """测试任务：访问URL""" print(f"[Task-{task_id}] 开始访问: {url}") time.sleep(2) # 模拟截图耗时 return {"task_id": task_id, "url": url, "status": "success"} def test_callback(result, error): """测试回调""" if error: print(f"任务失败: {error}") else: print(f"任务成功: {result}") # 创建线程池（2个worker） pool = BrowserWorkerPool(pool_size=2) pool.initialize() # 提交4个任务 for i in range(4): pool.submit_task(test_task, test_callback, f"https://example.com/{i}", i + 1) print("\n任务已提交，等待完成...") pool.wait_for_completion() print("\n统计信息:", pool.get_stats()) # 关闭线程池 pool.shutdown() print("\n测试完成！")