From 4a2cf7d19c37ac7a06bc87bf51c87290048794d3 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:23:05 +0800 Subject: [PATCH] Revert background worker cancellation changes --- reme/components/job/background_job.py | 90 +++----------- tests/unit/test_job.py | 173 -------------------------- 2 files changed, 19 insertions(+), 244 deletions(-) diff --git a/reme/components/job/background_job.py b/reme/components/job/background_job.py index 5304c874..ec8406e6 100644 --- a/reme/components/job/background_job.py +++ b/reme/components/job/background_job.py @@ -5,7 +5,6 @@ import contextlib import random import threading import time -from concurrent.futures import Future from .base_job import BaseJob from ..component_registry import R @@ -24,10 +23,7 @@ class BackgroundJob(BaseJob): supervisor cannot trigger a restart. 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. + ``close_timeout`` seconds to exit gracefully; after that it is cancelled. If a single run survives at least ``attempt_reset_after`` seconds before crashing, the backoff attempt counter resets — so a long-stable job that eventually crashes restarts quickly rather than at the capped delay. @@ -54,89 +50,41 @@ class BackgroundJob(BaseJob): self.attempt_reset_after: float = attempt_reset_after self.use_thread_pool: bool = use_thread_pool self._stop_event: asyncio.Event | threading.Event | None = None - self._task: asyncio.Future | None = None - self._executor_future: Future | None = None - self._worker_lock = threading.Lock() - self._worker_task: asyncio.Task | None = None + self._task: asyncio.Task | None = None async def _start(self) -> None: await super()._start() if self.use_thread_pool and self.app_context.thread_pool: self._stop_event = threading.Event() - self._executor_future = self.app_context.thread_pool.submit(self._run_in_thread) - self._task = asyncio.wrap_future(self._executor_future) + self._task = asyncio.get_event_loop().run_in_executor( + self.app_context.thread_pool, + lambda: asyncio.run(self._run_with_supervisor()), + ) else: self._stop_event = asyncio.Event() self._task = asyncio.create_task(self._run_with_supervisor()) - def _run_in_thread(self) -> None: - """Keep executor completion tied to actual event-loop shutdown.""" - with contextlib.suppress(asyncio.CancelledError): - asyncio.run(self._run_thread_task()) - - async def _run_thread_task(self) -> None: - with self._worker_lock: - self._worker_task = asyncio.current_task() - try: - # close() may have set stop before the worker published its task. - if not self._stop_event.is_set(): - await self._run_with_supervisor() - finally: - # Keep cancellation from targeting a loop that has already closed. - with self._worker_lock: - self._worker_task = None - - def _cancel_task(self) -> None: - if self._executor_future is None: - self._task.cancel() - return - # Cancel queued work without waiting for another long-running job to free - # a worker. If already running, cancel the coroutine on its own loop. - self._executor_future.cancel() - with self._worker_lock: - if self._worker_task is not None: - self._worker_task.get_loop().call_soon_threadsafe(self._worker_task.cancel) - async def _close(self) -> None: if self._stop_event is not None: self._stop_event.set() - try: - await self._shutdown_task() - finally: - await super()._close() + await self._shutdown_task() + await super()._close() async def _shutdown_task(self) -> None: - """Wait for graceful exit, then request cancellation and await completion.""" + """Wait close_timeout for graceful exit, then force-cancel.""" if self._task is None: return - - 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() + 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._task.cancel() + with contextlib.suppress(BaseException): + await self._task + except Exception: + self.logger.exception(f"Background task '{self.name}' raised during close") 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 4fad942c..191a27d9 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -3,8 +3,6 @@ # pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,no-self-argument,unused-argument import asyncio -import threading -from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -339,177 +337,6 @@ def test_shutdown_task_cancels_on_timeout(): asyncio.run(run()) -def test_thread_pool_shutdown_waits_for_worker_cleanup(): - async def run(): - started = threading.Event() - events = [] - - class WaitingJob(BackgroundJob): - async def __call__(self, **kwargs): - started.set() - try: - # Bounded fallback ensures a regression cannot hang pytest. - await asyncio.sleep(2) - except asyncio.CancelledError: - events.append("cancelled") - raise - finally: - await asyncio.sleep(0.02) - events.append("cleaned") - - with ThreadPoolExecutor(max_workers=1) as pool: - job = WaitingJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) - await job.start() - assert await asyncio.to_thread(started.wait, 1) - executor_future = job._executor_future - await asyncio.wait_for(job.close(), 1) - assert events == ["cancelled", "cleaned"] - assert executor_future.done() - assert job._task is None - - asyncio.run(run()) - - -@pytest.mark.parametrize("queued", [True, False], ids=["queued", "starting"]) -def test_thread_pool_shutdown_before_coroutine_starts(queued): - async def run(): - entered = threading.Event() - release = threading.Event() - - class DelayedJob(BackgroundJob): - def _run_in_thread(self): - entered.set() - release.wait(2) - super()._run_in_thread() - - def _cancel_task(self): - super()._cancel_task() - release.set() - - with ThreadPoolExecutor(max_workers=1) as pool: - if queued: - pool.submit(release.wait, 2) - job = DelayedJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01) - body = AsyncMock() - job._run_with_supervisor = body - try: - await job.start() - if not queued: - assert await asyncio.to_thread(entered.wait, 1) - executor_future = job._executor_future - await asyncio.wait_for(job.close(), 1) - assert executor_future.done() - assert executor_future.cancelled() == queued - body.assert_not_awaited() - finally: - release.set() - - 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