fix: join cancelled scheduler startup

This commit is contained in:
lda
2026-09-09 23:13:49 +07:00 Verified
parent 8d0c61737f
commit fc95326250
2 changed files with 95 additions and 1 deletions
+39 -1
View File
@@ -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
+56
View File
@@ -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]: