From 5d364f005a97d986add881d8240aef29798d8d7f Mon Sep 17 00:00:00 2001 From: lda Date: Wed, 9 Sep 2026 23:26:18 +0700 Subject: [PATCH] fix: finish cancelled resume cleanup --- src/wf_scheduling/lifecycle.py | 17 ++++++++-- src/wf_scheduling/resume_gate.py | 45 ++++++++++++++++---------- tests/scheduling/test_lifecycle.py | 52 ++++++++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 19 deletions(-) diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py index caeda0b0..214689f0 100644 --- a/src/wf_scheduling/lifecycle.py +++ b/src/wf_scheduling/lifecycle.py @@ -491,16 +491,29 @@ class SchedulerService: pass # -- internals ----------------------------------------------------- - async def _acquire_lock_async(self) -> None: + async def _acquire_lock_async(self, *, continue_after_cancel: bool = False) -> bool: """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. + Cleanup callers can request cancellation to be deferred until after + they have acquired and released the lock; otherwise a cancelled wait + exits immediately. The return value reports deferred cancellation. """ + cancelled = False while not self._lock.acquire(blocking=False): - await asyncio.sleep(0) + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + if not continue_after_cancel: + raise + # The synchronous operation guarded by this lock owns its + # cleanup boundary. Do not let cancellation strand capacity; + # propagate it after the caller has finished the operation. + cancelled = True + return cancelled async def _recover_joined(self) -> list[str]: """Run startup recovery in a worker and join it across cancellation.""" diff --git a/src/wf_scheduling/resume_gate.py b/src/wf_scheduling/resume_gate.py index 0180b25e..1a89004f 100644 --- a/src/wf_scheduling/resume_gate.py +++ b/src/wf_scheduling/resume_gate.py @@ -129,13 +129,15 @@ class SchedulerResumeGate: recovery reconcile by attempt identity. """ service = self._service - await service._acquire_lock_async() + cancelled = await service._acquire_lock_async(continue_after_cancel=True) try: service.run_store.clear_executing(run_id) service._live_resumes.discard(run_id) service._resume_tasks.pop(run_id, None) finally: service._lock.release() + if cancelled: + raise asyncio.CancelledError async def fence(self, run_id: str) -> None: """Forget a shutdown-cancelled execution; keep durable marks. @@ -148,12 +150,14 @@ class SchedulerResumeGate: never silently resumable, never replayed. """ service = self._service - await service._acquire_lock_async() + cancelled = await service._acquire_lock_async(continue_after_cancel=True) try: service._live_resumes.discard(run_id) service._resume_tasks.pop(run_id, None) finally: service._lock.release() + if cancelled: + raise asyncio.CancelledError async def reconcile_cancelled(self, run_id: str) -> None: """Reconcile caller cancellation while the scheduler stays live. @@ -166,7 +170,7 @@ class SchedulerResumeGate: shutdown; a shutdown that wins the race keeps the durable fence. """ service = self._service - await service._acquire_lock_async() + cancelled = await service._acquire_lock_async(continue_after_cancel=True) try: if service._started and not service._stopping and service._scheduler: try: @@ -192,6 +196,8 @@ class SchedulerResumeGate: service._resume_tasks.pop(run_id, None) finally: service._lock.release() + if cancelled: + raise asyncio.CancelledError async def note_resumed_result( self, @@ -212,21 +218,26 @@ class SchedulerResumeGate: the note must never break a completed resume. """ service = self._service - await service._acquire_lock_async() + cancelled = await service._acquire_lock_async(continue_after_cancel=True) try: if not service._started: - return False - scheduler = service._scheduler - if scheduler is None: - return False - try: - return scheduler.record_resumed_stopped_result( - run_id, - status_value=status_value, - checkpoint_id=checkpoint_id, - now=service.clock(), - ) - except SecondOwnerError, OSError, ValueError: - return False + noted = False + else: + scheduler = service._scheduler + if scheduler is None: + noted = False + else: + try: + noted = scheduler.record_resumed_stopped_result( + run_id, + status_value=status_value, + checkpoint_id=checkpoint_id, + now=service.clock(), + ) + except SecondOwnerError, OSError, ValueError: + noted = False finally: service._lock.release() + if cancelled: + raise asyncio.CancelledError + return noted diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py index 4bb7a1af..7e6d62b0 100644 --- a/tests/scheduling/test_lifecycle.py +++ b/tests/scheduling/test_lifecycle.py @@ -36,6 +36,7 @@ from wf_scheduling.lifecycle import ( from wf_scheduling.models import Schedule from wf_scheduling.ownership import SchedulerOwnership, SecondOwnerError from wf_scheduling.poll import Scheduler +from wf_scheduling.resume_gate import SchedulerResumeGate from wf_scheduling.store import FileScheduleStore @@ -405,6 +406,57 @@ asyncio.run(main()) assert "registration-stop-complete" in output +async def test_cancelled_resume_release_waits_for_lock_cleanup( + tmp_path: Path, +) -> None: + """Cancellation cannot skip scheduled-resume slot cleanup.""" + intended = ts(2026, 9, 8, 12, 0) + service = _service(tmp_path, ScriptedRuntime("interrupt")) + lock_held = threading.Event() + release_lock = threading.Event() + holder: threading.Thread | None = None + release_task: asyncio.Task[Any] | None = None + try: + await service.start() + service.schedule_store.create_schedule(_sched_model("a", intended)) + await service.poll_once(intended + timedelta(seconds=1)) + run_id = _only_run_id(service.run_store) + await _wait_for(lambda: service.live_executions == 0) + + gate = SchedulerResumeGate(service) + assert await gate.acquire(run_id, owner_task=asyncio.current_task()) + + def hold_service_lock() -> None: + service._lock.acquire() + lock_held.set() + release_lock.wait(5) + service._lock.release() + + holder = threading.Thread(target=hold_service_lock, daemon=True) + holder.start() + assert await asyncio.to_thread(lock_held.wait, 5) + + release_task = asyncio.create_task(gate.release(run_id)) + await asyncio.sleep(0.05) + assert not release_task.done() + release_task.cancel() + release_lock.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(release_task, 5) + + assert not service.run_store.is_executing(run_id) + assert run_id not in service._live_resumes + assert run_id not in service._resume_tasks + finally: + release_lock.set() + if release_task is not None and not release_task.done(): + release_task.cancel() + await asyncio.gather(release_task, return_exceptions=True) + if holder is not None: + holder.join(timeout=5) + 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"))