From 7241e1053e91ee71793e398eb1f2cbe5f2d80728 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Thu, 15 Oct 2026 17:12:36 +0800 Subject: [PATCH] Finish background worker cleanup before propagating cancellation --- reme/components/job/background_job.py | 43 ++++++++--- tests/unit/test_job.py | 103 ++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 12 deletions(-) diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index d4a99c99..5304c874 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -26,6 +26,7 @@ class BackgroundJob(BaseJob): On close, the stop_event is set and the task is given up to ``close_timeout`` seconds to exit gracefully; after that cancellation is requested on its event loop and close waits for the actual worker to exit. + Cancellation of close is propagated only after worker cleanup finishes. Job bodies must cooperate with cancellation, including in thread-pool mode. If a single run survives at least ``attempt_reset_after`` seconds before crashing, the backoff attempt counter resets — so a long-stable job that @@ -99,25 +100,43 @@ class BackgroundJob(BaseJob): async def _close(self) -> None: if self._stop_event is not None: self._stop_event.set() - await self._shutdown_task() - await super()._close() + try: + await self._shutdown_task() + finally: + await super()._close() async def _shutdown_task(self) -> None: """Wait for graceful exit, then request cancellation and await completion.""" if self._task is None: return - try: - # shield prevents wait_for's cancellation from propagating to the task itself, - # so a timeout here truly times out instead of cancelling silently. - await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) - except asyncio.TimeoutError: - self._cancel_task() - with contextlib.suppress(BaseException): - await asyncio.shield(self._task) - except Exception: - self.logger.exception(f"Background task '{self.name}' raised during close") + + async def wait_for_exit() -> None: + try: + try: + # A timeout must not cancel the executor's completion future. + await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) + except asyncio.TimeoutError: + self._cancel_task() + await self._task + except asyncio.CancelledError: + pass # The worker or queued executor future was cancelled. + except Exception: + self.logger.exception(f"Background task '{self.name}' raised during close") + + # Shield the whole shutdown so repeated caller cancellations cannot + # interrupt worker cleanup or cancel the executor's completion future. + shutdown = asyncio.create_task(wait_for_exit()) + cancelled = None + while not shutdown.done(): + try: + await asyncio.shield(shutdown) + except asyncio.CancelledError as exc: + cancelled = exc + shutdown.result() self._task = None self._executor_future = None + if cancelled is not None: + raise cancelled def _backoff_delay(self, attempt: int) -> float: """Exponential backoff with ±50% jitter, capped at backoff_cap.""" diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 1d68adbf..4fad942c 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -407,6 +407,109 @@ def test_thread_pool_shutdown_before_coroutine_starts(queued): asyncio.run(run()) +@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) +@pytest.mark.parametrize("phase", ["graceful", "forced"]) +def test_close_finishes_worker_cleanup_before_propagating_repeated_cancellation(use_thread_pool, phase): + async def run(): + started, cleaning, release, finished = (threading.Event() for _ in range(4)) + + class WaitingJob(BackgroundJob): + async def __call__(self, **kwargs): + started.set() + try: + await asyncio.sleep(1) + finally: + cleaning.set() + deadline = asyncio.get_running_loop().time() + 1 + while not release.is_set() and asyncio.get_running_loop().time() < deadline: + await asyncio.sleep(0.001) + finished.set() + + with ThreadPoolExecutor(max_workers=1) as pool: + job = WaitingJob( + app_context=SimpleNamespace(thread_pool=pool), + use_thread_pool=use_thread_pool, + close_timeout=0.1 if phase == "graceful" else 0.01, + ) + await job.start() + worker = job._task + executor_future = job._executor_future + closing = None + try: + assert await asyncio.to_thread(started.wait, 1) + closing = asyncio.create_task(job.close()) + if phase == "forced": + assert await asyncio.to_thread(cleaning.wait, 1) + else: + await asyncio.sleep(0) + for _ in range(3): + closing.cancel() + await asyncio.sleep(0) + assert not closing.done() + assert not worker.done() + assert job.is_started + assert await asyncio.to_thread(cleaning.wait, 1) + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(closing, 1) + assert finished.is_set() + assert worker.done() + assert executor_future is None or executor_future.done() + assert job._task is None + assert job._executor_future is None + assert not job.is_started + finally: + release.set() + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + await asyncio.gather(worker, return_exceptions=True) + await job.close() + + asyncio.run(run()) + + +@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"]) +def test_close_logs_worker_cleanup_failure_after_timeout(use_thread_pool): + async def run(): + started = threading.Event() + errors = [] + + class FailingCleanupJob(BackgroundJob): + async def __call__(self, **kwargs): + started.set() + try: + await asyncio.sleep(1) + finally: + raise RuntimeError("worker cleanup failed") + + def capture_error(_message): + import sys + + errors.append(sys.exception()) + + with ThreadPoolExecutor(max_workers=1) as pool: + job = FailingCleanupJob( + app_context=SimpleNamespace(thread_pool=pool), + use_thread_pool=use_thread_pool, + supervisor=False, + close_timeout=0.01, + ) + job.logger = MagicMock() + job.logger.exception.side_effect = capture_error + await job.start() + try: + assert await asyncio.to_thread(started.wait, 1) + finally: + await job.close() + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert str(errors[0]) == "worker cleanup failed" + assert job._task is None + assert job._executor_future is None + + asyncio.run(run()) + + def test_cron_uses_configured_timezone(monkeypatch): from zoneinfo import ZoneInfo