From 6a87711f893ce94e8b5c93408e4c54256fa2c923 Mon Sep 17 00:00:00 2001 From: lda Date: Wed, 9 Sep 2026 23:14:17 +0700 Subject: [PATCH] fix: avoid scheduler registration deadlock --- src/wf_api/runs.py | 10 ++-- src/wf_scheduling/lifecycle.py | 12 +++-- src/wf_scheduling/resume_gate.py | 35 ++++++++++---- tests/scheduling/test_lifecycle.py | 76 ++++++++++++++++++++++++++++++ 4 files changed, 113 insertions(+), 20 deletions(-) diff --git a/src/wf_api/runs.py b/src/wf_api/runs.py index 06e35b94..3d9936ee 100644 --- a/src/wf_api/runs.py +++ b/src/wf_api/runs.py @@ -207,7 +207,7 @@ class WorkflowRunApi: # the gate. gate = self.resume_slot_gate if gate is not None: - slot = gate.acquire(run_id, owner_task=asyncio.current_task()) + slot = await gate.acquire(run_id, owner_task=asyncio.current_task()) else: slot = None slot_held = slot is not None @@ -226,7 +226,7 @@ class WorkflowRunApi: # still live is reconciled for this run immediately, so its # ambiguous failure frees capacity without becoming retryable. if slot_held and gate is not None: - gate.reconcile_cancelled(run_id) + await gate.reconcile_cancelled(run_id) raise except BaseException: # Pre-persist errors, torn persists, and validation failures: @@ -235,11 +235,11 @@ class WorkflowRunApi: # attempt / stopped result already captures the ambiguity for # recovery to reconcile by attempt identity. if slot_held and gate is not None: - gate.release(run_id) + await gate.release(run_id) raise else: if slot_held and gate is not None: - gate.release(run_id) + await gate.release(run_id) return result async def _resume_scheduled_or_manual( @@ -352,7 +352,7 @@ class WorkflowRunApi: # quietly — restart recovery reconciles those instead. gate = self.resume_slot_gate if gate is not None: - gate.note_resumed_result( + await gate.note_resumed_result( run_id, status_value=run.status.value, checkpoint_id=next_record.latest_checkpoint_id, diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py index af73fccb..caeda0b0 100644 --- a/src/wf_scheduling/lifecycle.py +++ b/src/wf_scheduling/lifecycle.py @@ -397,11 +397,13 @@ class SchedulerService: """ if not self._started: return DrainReport() - with self._lock: - # Under the service lock, atomically with gate acquisition: - # every resume that passed the stopping check is already - # tracked, and every later one is rejected. - self._stopping = True + # Publish the drain intent before waiting for a poll that may be + # blocked in _submit. Event-loop gate callers cannot pass this check + # after stop begins, while the lock acquisition below still joins the + # already-running poll before the drain snapshots live work. + self._stopping = True + await self._acquire_lock_async() + self._lock.release() if self._tick_task is not None: self._tick_task.cancel() try: diff --git a/src/wf_scheduling/resume_gate.py b/src/wf_scheduling/resume_gate.py index c91f157e..0180b25e 100644 --- a/src/wf_scheduling/resume_gate.py +++ b/src/wf_scheduling/resume_gate.py @@ -49,7 +49,7 @@ class SchedulerResumeGate: def __init__(self, service: Any) -> None: self._service = service - def acquire( + async def acquire( self, run_id: str, *, owner_task: asyncio.Task[Any] | None = None ) -> Any | None: """Take the shared execution slot for a scheduled resume. @@ -74,7 +74,8 @@ class SchedulerResumeGate: rejected instead of running unfenced. """ service = self._service - with service._lock: + await service._acquire_lock_async() + try: if not service._started: return None try: @@ -114,8 +115,10 @@ class SchedulerResumeGate: if owner_task is not None: service._resume_tasks[run_id] = owner_task return admission + finally: + service._lock.release() - def release(self, run_id: str) -> None: + async def release(self, run_id: str) -> None: """Free the slot held by :meth:`acquire` after the stopped persist. Correct for every non-cancellation outcome: on success the @@ -126,12 +129,15 @@ class SchedulerResumeGate: recovery reconcile by attempt identity. """ service = self._service - with service._lock: + await service._acquire_lock_async() + 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() - def fence(self, run_id: str) -> None: + async def fence(self, run_id: str) -> None: """Forget a shutdown-cancelled execution; keep durable marks. The shutdown drain cancels a resume that outlives the grace @@ -142,11 +148,14 @@ class SchedulerResumeGate: never silently resumable, never replayed. """ service = self._service - with service._lock: + await service._acquire_lock_async() + try: service._live_resumes.discard(run_id) service._resume_tasks.pop(run_id, None) + finally: + service._lock.release() - def reconcile_cancelled(self, run_id: str) -> None: + async def reconcile_cancelled(self, run_id: str) -> None: """Reconcile caller cancellation while the scheduler stays live. Shutdown cancellation must retain the executing marker and ACTIVE @@ -157,7 +166,8 @@ class SchedulerResumeGate: shutdown; a shutdown that wins the race keeps the durable fence. """ service = self._service - with service._lock: + await service._acquire_lock_async() + try: if service._started and not service._stopping and service._scheduler: try: from wf_scheduling import recovery @@ -180,8 +190,10 @@ class SchedulerResumeGate: del service._errors[: max(0, len(service._errors) - 20)] service._live_resumes.discard(run_id) service._resume_tasks.pop(run_id, None) + finally: + service._lock.release() - def note_resumed_result( + async def note_resumed_result( self, run_id: str, *, @@ -200,7 +212,8 @@ class SchedulerResumeGate: the note must never break a completed resume. """ service = self._service - with service._lock: + await service._acquire_lock_async() + try: if not service._started: return False scheduler = service._scheduler @@ -215,3 +228,5 @@ class SchedulerResumeGate: ) except SecondOwnerError, OSError, ValueError: return False + finally: + service._lock.release() diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py index 67a05955..4bb7a1af 100644 --- a/tests/scheduling/test_lifecycle.py +++ b/tests/scheduling/test_lifecycle.py @@ -11,6 +11,8 @@ dedicated integration tests. from __future__ import annotations import asyncio +import subprocess +import sys import threading from datetime import UTC, datetime, timedelta from pathlib import Path @@ -329,6 +331,80 @@ async def test_shutdown_joins_cancelled_settlement_before_releasing_ownership( await service.stop() +def test_submit_and_stop_do_not_deadlock_on_registration(tmp_path: Path) -> None: + """The poll worker and event loop must not wait on each other.""" + child = """ +import asyncio +import sys +import threading +import time +from pathlib import Path + +from tests.scheduling.test_lifecycle import ( + SchedulerServiceConfig, + ScriptedRuntime, + _sched_model, + _service, + ts, +) +import wf_scheduling.lifecycle as lifecycle + + +async def main() -> None: + root = Path(sys.argv[1]) + service = _service( + root, + ScriptedRuntime("complete"), + config=SchedulerServiceConfig(poll_interval_s=0.01, auto_tick=False), + ) + await service.start() + service.schedule_store.create_schedule( + _sched_model("a", ts(2026, 9, 8, 12, 0)) + ) + registration_entered = threading.Event() + release_registration = threading.Event() + original = asyncio.run_coroutine_threadsafe + + def blocked_registration(coro, loop): + registration_entered.set() + if not release_registration.wait(5): + raise AssertionError("registration release was not signalled") + return original(coro, loop) + + lifecycle.asyncio.run_coroutine_threadsafe = blocked_registration + poll_task = asyncio.create_task( + service.poll_once(ts(2026, 9, 8, 12, 1)) + ) + if not await asyncio.to_thread(registration_entered.wait, 5): + raise AssertionError("poll did not reach the registration gate") + stop_task = asyncio.create_task(service.stop()) + threading.Thread( + target=lambda: (time.sleep(0.1), release_registration.set()), + daemon=True, + ).start() + await asyncio.wait_for(asyncio.gather(poll_task, stop_task), 4) + print("registration-stop-complete", flush=True) + + +asyncio.run(main()) +""" + process = subprocess.Popen( + [sys.executable, "-c", child, str(tmp_path)], + cwd=Path(__file__).resolve().parents[2], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + output, _ = process.communicate(timeout=7) + except subprocess.TimeoutExpired: + process.kill() + output, _ = process.communicate(timeout=5) + pytest.fail(f"registration/stop deadlocked; child output: {output}") + assert process.returncode == 0, output + assert "registration-stop-complete" in output + + async def test_scheduled_interrupt_stays_resumable(tmp_path: Path) -> None: intended = ts(2026, 9, 8, 12, 0) service = _service(tmp_path, ScriptedRuntime("interrupt"))