From 86a52adda58ebf050d31b2ca740ef69b482cf68c Mon Sep 17 00:00:00 2001 From: lda Date: Wed, 9 Sep 2026 18:33:57 +0700 Subject: [PATCH] fix: reconcile live scheduled resume cancellation --- src/wf_api/runs.py | 12 ++- src/wf_scheduling/resume_gate.py | 37 ++++++++- tests/wf_server/test_scheduler_integration.py | 81 +++++++++++++++++++ 3 files changed, 122 insertions(+), 8 deletions(-) diff --git a/src/wf_api/runs.py b/src/wf_api/runs.py index e28bc809..b815c9eb 100644 --- a/src/wf_api/runs.py +++ b/src/wf_api/runs.py @@ -221,14 +221,12 @@ class WorkflowRunApi: store=store, ) except asyncio.CancelledError: - # Shutdown drain cancelled the execution mid-flight (or the - # caller went away): fence, don't release. The durable - # executing mark and the ACTIVE attempt stay exactly as a - # crash mid-resume would leave them, so restart recovery - # fails the run closed instead of presenting a half-resumed - # run as safe to retry. Cancellation is never swallowed. + # Shutdown cancellation keeps the durable crash shape for + # restart recovery. A caller cancellation while the service is + # 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.fence(run_id) + gate.reconcile_cancelled(run_id) raise except BaseException: # Pre-persist errors, torn persists, and validation failures: diff --git a/src/wf_scheduling/resume_gate.py b/src/wf_scheduling/resume_gate.py index 19753d63..c91f157e 100644 --- a/src/wf_scheduling/resume_gate.py +++ b/src/wf_scheduling/resume_gate.py @@ -132,7 +132,7 @@ class SchedulerResumeGate: service._resume_tasks.pop(run_id, None) def fence(self, run_id: str) -> None: - """Forget a cancelled execution; keep its durable marks for recovery. + """Forget a shutdown-cancelled execution; keep durable marks. The shutdown drain cancels a resume that outlives the grace deadline and joins it before releasing ownership, so no old @@ -146,6 +146,41 @@ class SchedulerResumeGate: service._live_resumes.discard(run_id) service._resume_tasks.pop(run_id, None) + 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 + attempt until startup recovery decides the ambiguous outcome. A + caller cancellation during ordinary operation has no future owner to + perform that recovery, so reconcile only this run immediately. The + service lock makes the stopping check and scoped recovery atomic with + shutdown; a shutdown that wins the race keeps the durable fence. + """ + service = self._service + with service._lock: + if service._started and not service._stopping and service._scheduler: + try: + from wf_scheduling import recovery + + recovery.recover( + schedule_store=service.schedule_store, + run_store=service.run_store, + now=service.clock(), + ownership=service.ownership, + only_run_id=run_id, + ) + except Exception as exc: + # Keep the durable marker when reconciliation cannot be + # proven. The in-memory task is still gone, but capacity + # remains conservatively occupied and the error is visible + # through the service diagnostics. + service._errors.append( + f"{run_id}: cancellation recovery failed: {exc}" + ) + del service._errors[: max(0, len(service._errors) - 20)] + service._live_resumes.discard(run_id) + service._resume_tasks.pop(run_id, None) + def note_resumed_result( self, run_id: str, diff --git a/tests/wf_server/test_scheduler_integration.py b/tests/wf_server/test_scheduler_integration.py index 4055ded7..1a1fc7c5 100644 --- a/tests/wf_server/test_scheduler_integration.py +++ b/tests/wf_server/test_scheduler_integration.py @@ -1408,6 +1408,87 @@ async def test_shutdown_timeout_fences_inflight_resume_and_restart_keeps_decisio await service.stop() +async def test_caller_cancellation_reconciles_and_frees_capacity( + tmp_path: Path, +) -> None: + """Caller cancellation fails one ambiguous resume without fencing capacity.""" + _gate_open.clear() + root = tmp_path / "store" + server = build_local_static_workflow_server(root, extra_sources=_gate_sources()) + await _seed( + server, + "ask_gate", + "ask_gate.default", + _interrupt_then_gate_plan("ask_gate"), + ["submitted"], + ) + await _seed(server, "const", "const.default", _constant_plan("const"), ["ok"]) + store = FileScheduleStore(root) + store.create_schedule(_one_shot("asker", "ask_gate.default", datetime.now(UTC))) + service = _scheduler(server, capacity=1, auto_tick=False) + try: + await service.start() + async with asyncio.timeout(20): + while not _run_ids(root): + await service.poll_once(datetime.now(UTC)) + await asyncio.sleep(0.02) + ask_id = _run_ids(root)[0] + await _wait_for( + lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted" + ) + + resume_task = asyncio.create_task( + server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + ) + await _wait_for(lambda: ask_id in service._live_resumes) + await _wait_for(lambda: FileRunStore(root).is_executing(ask_id)) + await _wait_for( + lambda: ( + (attempt := FileRunStore(root).get_resume_attempt(ask_id)) is not None + and attempt.state == "ACTIVE" + ) + ) + resume_task.cancel() + with pytest.raises(asyncio.CancelledError): + await resume_task + + assert service.running + assert not service._live_resumes + cancelled = FileRunStore(root).get_run(ask_id) + assert cancelled.status.value == "failed" + assert not FileRunStore(root).is_executing(ask_id) + attempt = FileRunStore(root).get_resume_attempt(ask_id) + assert attempt is not None and attempt.state == "ACTIVE" + assert len(_entries(root, "asker", "failed")) == 1 + + with pytest.raises(ValueError, match="ambiguous active resume attempt"): + await server.api.resume_run( + run_id=ask_id, resume_payload={}, resume_outcome="submitted" + ) + assert _entries(root, "asker", "completed") == [] + + store.create_schedule(_one_shot("after", "const.default", datetime.now(UTC))) + async with asyncio.timeout(20): + while True: + await service.poll_once(datetime.now(UTC)) + if len(_run_ids(root)) == 2: + after_id = [rid for rid in _run_ids(root) if rid != ask_id][0] + if FileRunStore(root).get_run(after_id).status.value == "completed": + break + await asyncio.sleep(0.02) + assert await _kinds(root, "asker") == [ + "admitted", + "interrupted", + "failed", + ] + assert await _kinds(root, "after") == ["admitted", "completed"] + finally: + _gate_open.set() + await service.stop() + + async def test_scheduled_resume_requested_during_drain_rejected( tmp_path: Path, ) -> None: