fix: avoid scheduler registration deadlock

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,

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:

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()

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"))