diff --git a/browser_pool_worker.py b/browser_pool_worker.py index 0d56fcf..12e632d 100755 --- a/browser_pool_worker.py +++ b/browser_pool_worker.py @@ -165,21 +165,48 @@ class BrowserWorker(threading.Thread): self.idle = False - if task is None: # None作为停止信号 - self.log("收到停止信号") - break - - # 按需创建或确保浏览器可用 - if not self._ensure_browser(): - self.log("浏览器不可用,任务失败") - task['callback'](None, "浏览器不可用") - self.failed_tasks += 1 - continue - - # 执行任务 - task_func = task.get('func') - task_args = task.get('args', ()) - task_kwargs = task.get('kwargs', {}) + 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 @@ -315,12 +342,13 @@ class BrowserWorkerPool: self.log("警告:线程池未初始化") return False - task = { - 'func': task_func, - 'args': args, - 'kwargs': kwargs, - 'callback': callback - } + task = { + 'func': task_func, + 'args': args, + 'kwargs': kwargs, + 'callback': callback, + 'retry_count': 0, + } try: self.task_queue.put(task, timeout=1) diff --git a/tests/test_browser_pool_worker.py b/tests/test_browser_pool_worker.py new file mode 100644 index 0000000..0212f2c --- /dev/null +++ b/tests/test_browser_pool_worker.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import queue + +from browser_pool_worker import BrowserWorker + + +class _AlwaysFailEnsureWorker(BrowserWorker): + def __init__(self, *, worker_id: int, task_queue: queue.Queue): + super().__init__(worker_id=worker_id, task_queue=task_queue, pre_warm=False) + self.ensure_calls = 0 + + def _ensure_browser(self) -> bool: # noqa: D401 - matching base naming + self.ensure_calls += 1 + if self.ensure_calls >= 2: + self.running = False + return False + + def _close_browser(self): + self.browser_instance = None + + +def test_requeue_task_when_browser_unavailable(): + task_queue: queue.Queue = queue.Queue() + callback_calls: list[tuple[object, object]] = [] + + def callback(result, error): + callback_calls.append((result, error)) + + task = { + "func": lambda *_args, **_kwargs: None, + "args": (), + "kwargs": {}, + "callback": callback, + "retry_count": 0, + } + + worker = _AlwaysFailEnsureWorker(worker_id=1, task_queue=task_queue) + worker.start() + task_queue.put(task) + worker.join(timeout=5) + + assert worker.is_alive() is False + assert worker.ensure_calls == 2 # 本地最多尝试2次创建浏览器 + assert callback_calls == [] # 第一次失败会重新入队,不应立即回调失败 + + requeued = task_queue.get_nowait() + assert requeued["retry_count"] == 1 + + +def test_fail_task_after_second_assignment(): + task_queue: queue.Queue = queue.Queue() + callback_calls: list[tuple[object, object]] = [] + + def callback(result, error): + callback_calls.append((result, error)) + + task = { + "func": lambda *_args, **_kwargs: None, + "args": (), + "kwargs": {}, + "callback": callback, + "retry_count": 1, # 已重新分配过1次 + } + + worker = _AlwaysFailEnsureWorker(worker_id=1, task_queue=task_queue) + worker.start() + task_queue.put(task) + worker.join(timeout=5) + + assert worker.is_alive() is False + assert callback_calls == [(None, "浏览器不可用")] + assert worker.total_tasks == 1 + assert worker.failed_tasks == 1 +