diff --git a/src/wf_scheduling/lifecycle.py b/src/wf_scheduling/lifecycle.py index 214689f0..b81699f1 100644 --- a/src/wf_scheduling/lifecycle.py +++ b/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) diff --git a/tests/scheduling/test_lifecycle.py b/tests/scheduling/test_lifecycle.py index 7e6d62b0..86c86750 100644 --- a/tests/scheduling/test_lifecycle.py +++ b/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"))