From fc953262503613b424daa4aa2f3fd27db2a3696a Mon Sep 17 00:00:00 2001
From: lda <lda@ldlda.com>
Date: Wed, 9 Sep 2026 23:13:49 +0700
Subject: [PATCH] fix: join cancelled scheduler startup

---
 src/wf_scheduling/lifecycle.py     | 40 ++++++++++++++++++++-
 tests/scheduling/test_lifecycle.py | 56 ++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py
index bd7ced95..af73fccb 100644
--- a/src/wf_scheduling/lifecycle.py
+++ b/src/wf_scheduling/lifecycle.py
@@ -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
 
diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py
index 47924fad..67a05955 100644
--- a/tests/scheduling/test_lifecycle.py
+++ b/tests/scheduling/test_lifecycle.py
@@ -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]: