fix: close late scheduler poll during drain

src/wf_scheduling/lifecycle.py

@@ -429,19 +429,25 @@ class SchedulerService:
         # unwind through the gate fence (in-memory accounting dropped,
         # durable marks kept), so recovery sees the exact crash shape.
         current = asyncio.current_task()
-        with self._lock:
+        await self._acquire_lock_async()
+        try:
             resume_pending = [
                 task
                 for task in self._resume_tasks.values()
                 if task is not current and not task.done()
             ]
+        finally:
+            self._lock.release()
         for task in resume_pending:
             task.cancel()
         if resume_pending:
             await asyncio.gather(*resume_pending, return_exceptions=True)
-        with self._lock:
+        await self._acquire_lock_async()
+        try:
             self._resume_tasks.clear()
             self._live_resumes.clear()
+        finally:
+            self._lock.release()
         resume_cancelled = len([t for t in resume_pending if t.cancelled()])
         pending = list(self._executions)
         if pending:
@@ -457,12 +463,15 @@ class SchedulerService:
             for task in pending:
                 self._executions.discard(task)
         cancelled = len([t for t in pending if t.cancelled()]) + resume_cancelled
-        with self._lock:
+        await self._acquire_lock_async()
+        try:
             report = DrainReport(
                 settled=self._settled,
                 abandoned=self._abandoned,
                 cancelled=cancelled,
             )
+        finally:
+            self._lock.release()
         self.ownership.release()
         self._started = False
         return report
@@ -475,6 +484,8 @@ class SchedulerService:
         """
         if not self._started:
             raise SchedulerStartupError("scheduler service is not started")
+        if self._stopping:
+            raise SchedulerStartupError("scheduler service is stopping")
         instant = now if now is not None else self.clock()
         try:
             return await asyncio.to_thread(self._tick, instant)
@@ -588,6 +599,11 @@ class SchedulerService:
 
     def _tick(self, now: datetime) -> dict[str, str]:
         with self._lock:
+            # poll_once can be submitted just as stop publishes the drain
+            # intent. Check again under the scheduler lock so a late worker
+            # cannot dispatch while stop joins the poll boundary.
+            if self._stopping:
+                raise SchedulerStartupError("scheduler service is stopping")
             self._refresh_sources()
             return self._scheduler.poll(now)
 

tests/scheduling/test_lifecycle.py

@@ -457,6 +457,23 @@ async def test_cancelled_resume_release_waits_for_lock_cleanup(
         await service.stop()
 
 
+async def test_poll_once_rejects_after_stop_begins(tmp_path: Path) -> None:
+    """A late manual tick cannot enter the scheduler during drain."""
+    intended = ts(2026, 9, 8, 12, 0)
+    service = _service(tmp_path, ScriptedRuntime("complete"))
+    try:
+        await service.start()
+        service.schedule_store.create_schedule(_sched_model("a", intended))
+        service._stopping = True
+
+        with pytest.raises(SchedulerStartupError, match="stopping"):
+            await service.poll_once(intended + timedelta(seconds=1))
+        assert service.run_store.list_runs() == []
+    finally:
+        service._stopping = False
+        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"))