Cancel and drain in-flight ctx.run tasks before raising 409 on cancellation

python/restate/server_context.py

@@ -447,6 +447,26 @@ def cancel(self):
         for task in to_cancel:
             task.cancel()
 
+    async def cancel_and_drain(self):
+        """Cancel all tasks in the list and wait until they have finished unwinding."""
+        to_cancel = list(self.tasks)
+        self.tasks.clear()
+        for task in to_cancel:
+            task.cancel()
+        if not to_cancel:
+            return
+        # Shield the children: cancelling the caller must not re-cancel tasks that are still
+        # running their cleanup. The caller's cancellation is re-raised once they have finished.
+        gather_fut = asyncio.gather(*to_cancel, return_exceptions=True)
+        caller_cancelled = False
+        while not gather_fut.done():
+            try:
+                await asyncio.shield(gather_fut)
+            except asyncio.CancelledError:
+                caller_cancelled = True
+        if caller_cancelled:
+            raise asyncio.CancelledError()
+
 
 restate_context_is_replaying = contextvars.ContextVar("restate_context_is_replaying", default=False)
 
@@ -694,6 +714,9 @@ async def create_poll_or_cancel_coroutine(self, unresolved_future: UnresolvedFut
                 # One of the handles completed
                 return
             if isinstance(do_progress_response, DoProgressCancelSignalReceived):
+                # In-flight ctx.run closures must finish unwinding before the handler's
+                # cancellation handling (e.g. compensation) starts running.
+                await self.tasks.cancel_and_drain()
                 raise TerminalError("cancelled", 409)
             if isinstance(do_progress_response, DoProgressExecuteRun):
                 fn = self.run_coros_to_execute[do_progress_response.handle]

tests/disconnect_hotloop.py

@@ -161,6 +161,105 @@ async def mock_receive() -> ASGIReceiveEvent:
         await receive_channel.close()
 
 
+async def test_cancel_signal_drains_in_flight_run_before_raising():
+    """
+    When the cancel signal arrives while a ctx.run closure is still executing, the closure must
+    observe CancelledError and finish its cleanup before the handler sees TerminalError(409).
+    """
+    from restate.exceptions import TerminalError
+    from restate.server_context import ServerInvocationContext, Tasks
+    from restate.vm import (
+        DoProgressCancelSignalReceived,
+        DoProgressExecuteRun,
+        DoProgressWaitExternalProgress,
+        SingleUnresolvedFuture,
+    )
+
+    run_started = asyncio.Event()
+    events: list[str] = []
+
+    async def run_action():
+        run_started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            events.append("run cancelled")
+            raise
+        finally:
+            # Cleanup that itself awaits, so it only completes if the canceller waits for it.
+            await asyncio.sleep(0.01)
+            events.append("run cleanup done")
+
+    class FakeReceive:
+        async def __call__(self) -> ASGIReceiveEvent:
+            # Hand control to the run task until it is actually executing the closure.
+            await run_started.wait()
+            return cast(ASGIReceiveEvent, {"type": "restate.run_completed", "data": None})
+
+        async def enqueue_restate_event(self, what):
+            pass
+
+    vm = MagicMock()
+    vm.take_output.return_value = None
+    vm.do_progress.side_effect = [
+        DoProgressExecuteRun(handle=1),
+        DoProgressWaitExternalProgress(),
+        DoProgressCancelSignalReceived(),
+    ]
+
+    ctx = ServerInvocationContext.__new__(ServerInvocationContext)
+    ctx.vm = vm
+    ctx.receive = FakeReceive()  # type: ignore[assignment]
+    ctx.run_coros_to_execute = {1: run_action}
+    ctx.tasks = Tasks()
+
+    with pytest.raises(TerminalError) as exc_info:
+        await asyncio.wait_for(ctx.create_poll_or_cancel_coroutine(SingleUnresolvedFuture(1)), timeout=2.0)
+    events.append("handler saw TerminalError")
+
+    assert exc_info.value.status_code == 409
+    assert events == ["run cancelled", "run cleanup done", "handler saw TerminalError"]
+    assert not ctx.tasks.tasks
+
+
+async def test_cancel_and_drain_finishes_cleanup_when_caller_is_cancelled():
+    """Cancelling the drainer must not interrupt a child that is still awaiting in its finally block."""
+    from restate.server_context import Tasks
+
+    in_cleanup = asyncio.Event()
+    release_cleanup = asyncio.Event()
+    events: list[str] = []
+
+    async def child():
+        try:
+            await asyncio.Event().wait()
+        finally:
+            in_cleanup.set()
+            await release_cleanup.wait()
+            events.append("child cleanup done")
+
+    tasks = Tasks()
+    child_task = asyncio.create_task(child())
+    tasks.add(child_task)
+    await asyncio.sleep(0)
+
+    drainer = asyncio.create_task(tasks.cancel_and_drain())
+    await asyncio.wait_for(in_cleanup.wait(), timeout=1.0)
+
+    drainer.cancel()
+    await asyncio.sleep(0.01)
+    assert not drainer.done()
+    assert not child_task.done()
+
+    release_cleanup.set()
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(drainer, timeout=1.0)
+    events.append("drainer cancelled")
+
+    assert events == ["child cleanup done", "drainer cancelled"]
+    assert child_task.cancelled()
+
+
 # ---- Shutdown / SIGTERM tests ----