diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py index 370bf9ee..bd7ced95 100644 --- a/src/wf_scheduling/lifecycle.py +++ b/src/wf_scheduling/lifecycle.py @@ -146,9 +146,12 @@ class RuntimeDispatcher: a settlement failure such as a torn write — goes through the service-provided ``abandon`` hook, which runs the same per-run reconciliation as startup recovery but scoped to the broken run id: - sibling executions are never touched. Cancellation (shutdown drain) is - re-raised unsettled: the durable executing mark is left for startup - recovery, which abandons it truthfully. + sibling executions are never touched. Cancellation before the runtime + returns is re-raised unsettled: the durable executing mark is left for + startup recovery, which abandons it truthfully. Once execution has + returned, settlement or scoped recovery is joined to completion even if + its task is cancelled, because a synchronous store transition cannot be + allowed to outlive the lifecycle owner. """ def __init__( @@ -195,14 +198,35 @@ class RuntimeDispatcher: except asyncio.CancelledError: raise except Exception as exc: - self._abandon(admission.id, exc) + await self._run_blocking_joined(self._abandon, admission.id, exc) return try: - self._settle(admission.id, state) + await self._run_blocking_joined(self._settle, admission.id, state) except asyncio.CancelledError: raise except Exception as exc: - self._abandon(admission.id, exc) + await self._run_blocking_joined(self._abandon, admission.id, exc) + + async def _run_blocking_joined( + self, + operation: Callable[..., Any], + *args: Any, + ) -> Any: + """Run a sync lifecycle operation off-loop and join it if cancelled. + + ``asyncio.to_thread`` alone lets its worker continue after the + awaiting task is cancelled. The outer execution remains in + ``SchedulerService._executions``; swallowing cancellation until this + worker finishes reuses that existing lifecycle tracking and prevents + persistence from racing shutdown ownership release. + """ + worker = asyncio.create_task(asyncio.to_thread(operation, *args)) + while True: + try: + return await asyncio.shield(worker) + except asyncio.CancelledError: + if worker.done(): + return await worker @dataclass(slots=True) @@ -534,11 +558,24 @@ class SchedulerService: """Register one execution task; dispatch always spawns exactly one.""" if self._loop is None: # pragma: no cover - start() sets this first raise SchedulerStartupError("scheduler service is not started") - # Submitted from the poll worker thread: wrap explicitly against - # the server loop, which has no current-loop binding out there. - tracked = asyncio.wrap_future( - asyncio.run_coroutine_threadsafe(coro, self._loop), loop=self._loop + # Submitted from the poll worker thread: create and track the actual + # loop task before returning. Tracking a wrap_future proxy would let + # cancellation mark the proxy done while the underlying coroutine + # (and any joined settlement worker) was still unwinding. + registration = asyncio.run_coroutine_threadsafe( + self._register_execution(coro), self._loop ) + try: + return registration.result() + except BaseException: + coro.close() + raise + + async def _register_execution( + self, coro: Coroutine[Any, Any, None] + ) -> asyncio.Task[Any]: + """Create one execution task on the service event loop.""" + tracked = asyncio.create_task(coro) self._executions.add(tracked) tracked.add_done_callback(self._executions.discard) return tracked diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py index 118d48e1..f1df75e6 100644 --- a/tests/scheduling/test_lifecycle.py +++ b/tests/scheduling/test_lifecycle.py @@ -11,6 +11,7 @@ dedicated integration tests. from __future__ import annotations import asyncio +import threading from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any, cast @@ -238,6 +239,95 @@ async def test_scheduled_failure_records_failed_history(tmp_path: Path) -> None: await service.stop() +async def test_settlement_does_not_block_the_event_loop( + tmp_path: Path, monkeypatch: Any +) -> None: + """Synchronous settlement persistence must not stall other coroutines.""" + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("complete")) + settlement_started = threading.Event() + settlement_finished = threading.Event() + release = threading.Event() + original = Scheduler.record_stopped_execution + + def _blocking_settle(self: Any, run_id: str, state: Any, now: datetime) -> None: + settlement_started.set() + release.wait(timeout=2) + original(self, run_id, state, now) + settlement_finished.set() + + monkeypatch.setattr(Scheduler, "record_stopped_execution", _blocking_settle) + release_timer = threading.Timer(1.0, release.set) + release_timer.start() + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + await asyncio.sleep(0.05) + assert settlement_started.is_set() + assert not settlement_finished.is_set() + assert service.live_executions == 1 + release.set() + await _wait_for(lambda: service.live_executions == 0) + finally: + release.set() + release_timer.cancel() + await service.stop() + + +async def test_shutdown_joins_cancelled_settlement_before_releasing_ownership( + tmp_path: Path, monkeypatch: Any +) -> None: + """Shutdown waits for an offloaded settlement after cancelling its task.""" + intended = ts(2026, 9, 8, 12, 0) + service = _service( + tmp_path, + ScriptedRuntime("complete"), + SchedulerServiceConfig( + poll_interval_s=0.01, + drain_grace_s=0, + auto_tick=False, + ), + ) + settlement_started = threading.Event() + release = threading.Event() + original = Scheduler.record_stopped_execution + + def _blocking_settle(self: Any, run_id: str, state: Any, now: datetime) -> None: + settlement_started.set() + release.wait(timeout=2) + original(self, run_id, state, now) + + monkeypatch.setattr(Scheduler, "record_stopped_execution", _blocking_settle) + release_timer = threading.Timer(1.0, release.set) + release_timer.start() + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + for _ in range(50): + if settlement_started.is_set(): + break + await asyncio.sleep(0.01) + else: + pytest.fail("settlement did not start") + + stop_task = asyncio.create_task(service.stop()) + await asyncio.sleep(0.05) + assert not stop_task.done() + release.set() + report = await stop_task + assert report.settled == 1 + assert report.cancelled == 0 + assert service.run_store.get_run( + _only_run_id(service.run_store) + ).status.value == ("completed") + finally: + release.set() + release_timer.cancel() + await service.stop() + + async def test_scheduled_interrupt_stays_resumable(tmp_path: Path) -> None: intended = ts(2026, 9, 8, 12, 0) service = _service(tmp_path, ScriptedRuntime("interrupt"))