fix: join cancelled scheduler startup
This commit is contained in:
@@ -357,7 +357,14 @@ class SchedulerService:
|
||||
) from exc
|
||||
try:
|
||||
self._build_scheduler()
|
||||
await asyncio.to_thread(self._recover)
|
||||
await self._recover_joined()
|
||||
except asyncio.CancelledError:
|
||||
# Cancellation stops the awaiter, not a to_thread worker. Join
|
||||
# completion is guaranteed by _recover_joined before this path,
|
||||
# so ownership can be released without a recovery write racing a
|
||||
# later service owner.
|
||||
self.ownership.release()
|
||||
raise
|
||||
except Exception as exc:
|
||||
self.ownership.release()
|
||||
raise SchedulerStartupError(
|
||||
@@ -482,6 +489,37 @@ class SchedulerService:
|
||||
pass
|
||||
|
||||
# -- internals -----------------------------------------------------
|
||||
async def _acquire_lock_async(self) -> None:
|
||||
"""Acquire the service lock without blocking the event loop.
|
||||
|
||||
Polling and settlement run in worker threads and may be waiting for
|
||||
an event-loop registration callback while holding this lock. A normal
|
||||
blocking ``Lock.acquire`` from the event loop would form a circular
|
||||
wait, so retry non-blocking acquisition while yielding to callbacks.
|
||||
"""
|
||||
while not self._lock.acquire(blocking=False):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def _recover_joined(self) -> list[str]:
|
||||
"""Run startup recovery in a worker and join it across cancellation."""
|
||||
worker = asyncio.create_task(asyncio.to_thread(self._recover))
|
||||
cancelled = False
|
||||
while True:
|
||||
try:
|
||||
result = await asyncio.shield(worker)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
continue
|
||||
except BaseException:
|
||||
# If recovery failed after startup cancellation, cancellation
|
||||
# remains the caller-visible result, but the worker is done.
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
raise
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return result
|
||||
|
||||
def _build_scheduler(self) -> None:
|
||||
from wf_scheduling.poll import Scheduler
|
||||
|
||||
|
||||
@@ -531,6 +531,62 @@ async def test_failed_startup_releases_lock(tmp_path: Path) -> None:
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_cancelled_startup_joins_gated_recovery_before_releasing_ownership(
|
||||
tmp_path: Path, monkeypatch: Any
|
||||
) -> None:
|
||||
"""Cancelled startup waits for recovery before releasing its store lock."""
|
||||
service = _service(tmp_path, ScriptedRuntime("complete"))
|
||||
recovery_started = threading.Event()
|
||||
recovery_release = threading.Event()
|
||||
recovery_ready = threading.Event()
|
||||
recovery_continue = threading.Event()
|
||||
recovery_finished = threading.Event()
|
||||
original = SchedulerService._recover
|
||||
|
||||
def gated_recover(self: SchedulerService) -> list[str]:
|
||||
recovery_started.set()
|
||||
if not recovery_release.wait(5):
|
||||
raise AssertionError("recovery release was not signalled")
|
||||
recovery_ready.set()
|
||||
if not recovery_continue.wait(5):
|
||||
raise AssertionError("recovery continuation was not signalled")
|
||||
result = original(self)
|
||||
recovery_finished.set()
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(SchedulerService, "_recover", gated_recover)
|
||||
start_task = asyncio.create_task(service.start())
|
||||
try:
|
||||
assert await asyncio.to_thread(recovery_started.wait, 5)
|
||||
start_task.cancel()
|
||||
recovery_release.set()
|
||||
assert await asyncio.to_thread(recovery_ready.wait, 5)
|
||||
# The cancelled start is still joining the active recovery worker.
|
||||
assert not start_task.done()
|
||||
|
||||
retry = _service(tmp_path, ScriptedRuntime("complete"))
|
||||
with pytest.raises(SecondOwnerError):
|
||||
retry.ownership.acquire()
|
||||
|
||||
recovery_continue.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(start_task, 5)
|
||||
assert recovery_finished.is_set()
|
||||
assert service.running is False
|
||||
assert service.live_executions == 0
|
||||
assert service._tick_task is None
|
||||
|
||||
await retry.start()
|
||||
await retry.stop()
|
||||
finally:
|
||||
recovery_release.set()
|
||||
recovery_continue.set()
|
||||
if not start_task.done():
|
||||
start_task.cancel()
|
||||
await asyncio.gather(start_task, return_exceptions=True)
|
||||
await service.stop()
|
||||
|
||||
|
||||
async def test_corrupt_store_startup_releases_lock(tmp_path: Path) -> None:
|
||||
class BrokenRuns(FileRunStore):
|
||||
def list_admissions(self) -> list[Any]:
|
||||
|
||||
Reference in New Issue
Block a user