Add explicit cancellation (preview)

python/restate/_preview.py

@@ -12,7 +12,10 @@
 
 import typing
 
+from .context import Context, RestateDurableFuture
+from .exceptions import TerminalError
 from .handler import handler_from_callable
+from .server_context import ServerInvocationContext
 from .service import Service
 from .object import VirtualObject
 from .workflow import Workflow
@@ -53,3 +56,44 @@ async def greet(ctx: Context, name: str) -> str:
     else:
         handler_from_callable(target).explicit_recovery = True  # type: ignore
     return target
+
+
+def cancellation(ctx: Context) -> RestateDurableFuture[None]:
+    """
+    Return a future that completes when Restate cancels the invocation of `ctx`.
+
+    Await it together with the handler's other work, for example with `restate.select`::
+
+        @greeter.handler(explicit_cancellation=True)
+        async def greet(ctx: Context, name: str) -> str:
+            work = ctx.service_call(slow_greeting, arg=name)
+            done, _ = await restate.select(work=work, cancelled=cancellation(ctx))
+            if done == "cancelled":
+                await cancel_previous_calls(ctx)
+                raise TerminalError("cancelled", status_code=409)
+            ...
+
+    The handler must run with ``explicit_cancellation=True`` (on the handler, or on its service): then
+    Restate cancelling the invocation neither fails the handler's awaits nor cancels its calls.
+
+    NOTE: This API is in preview and might change without notice.
+    """
+    return _server_context(ctx, "cancellation").cancellation()
+
+
+async def cancel_previous_calls(ctx: Context) -> typing.List[str]:
+    """
+    Cancel the calls that this attempt of the handler made and that have not completed yet.
+
+    This is what Restate does on its own for a handler without explicit cancellation. Returns the invocation
+    ids of the cancelled calls. The handler must run with ``explicit_cancellation=True``.
+
+    NOTE: This API is in preview and might change without notice.
+    """
+    return await _server_context(ctx, "cancel_previous_calls").cancel_previous_calls()
+
+
+def _server_context(ctx: Context, operation: str) -> ServerInvocationContext:
+    if not isinstance(ctx, ServerInvocationContext):
+        raise TerminalError(f"{operation}() takes the context a handler is invoked with")
+    return ctx

python/restate/handler.py

@@ -256,6 +256,8 @@ class Handler(Generic[I, O]):
         ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress.
         invocation_retry_policy: Optional retry policy configuration applied to this handler.
         explicit_recovery: If true, the handler runs with explicit recovery, see restate._preview.explicit_recovery.
+        explicit_cancellation: If true, the handler runs with explicit cancellation (preview):
+            see restate._preview.cancellation.
     """
 
     service_tag: ServiceTag
@@ -277,6 +279,7 @@ class Handler(Generic[I, O]):
     invocation_retry_policy: Optional[InvocationRetryPolicy] = None
     context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None
     explicit_recovery: bool = False
+    explicit_cancellation: bool = False
 
 
 # disable too many arguments warning
@@ -301,6 +304,7 @@ def make_handler(
     ingress_private: Optional[bool] = None,
     invocation_retry_policy: Optional[InvocationRetryPolicy] = None,
     context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None,
+    explicit_cancellation: bool = False,
 ) -> Handler[I, O]:
     """
     Factory function to create a handler.
@@ -325,6 +329,7 @@ def make_handler(
         ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress.
         invocation_retry_policy: Retry policy used by Restate when invoking this handler.
             NOTE: only supported with restate-server >= 1.5.
+        explicit_cancellation: If true, the handler runs with explicit cancellation (preview).
     """
     # try to deduce the handler name
     handler_name = name
@@ -360,6 +365,7 @@ def make_handler(
         ingress_private=ingress_private,
         invocation_retry_policy=invocation_retry_policy,
         context_managers=context_managers,
+        explicit_cancellation=explicit_cancellation,
     )
 
     vars(wrapped)[RESTATE_UNIQUE_HANDLER_SYMBOL] = handler

python/restate/object.py

@@ -75,6 +75,8 @@ class VirtualObject:
             NOTE: You can set this field only if you register this service against restate-server >= 1.4,
             otherwise the service discovery will fail.
         invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this virtual object.
+        explicit_cancellation (bool, optional): (Preview) Explicit cancellation for all handlers of this virtual object:
+            see the handler option of the same name. NOTE: in preview, it might change without notice.
     """
 
     handlers: typing.Dict[str, Handler[typing.Any, typing.Any]]
@@ -95,6 +97,7 @@ def __init__(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ):
         self.service_tag = ServiceTag("object", name, description, metadata)
         self.handlers = {}
@@ -109,6 +112,7 @@ def __init__(
         self.invocation_retry_policy = invocation_retry_policy
         self.context_managers = invocation_context_managers
         self.explicit_recovery = False
+        self.explicit_cancellation = explicit_cancellation
 
     @property
     def name(self):
@@ -138,6 +142,7 @@ def handler(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> typing.Callable[[T], T]:
         """
         Decorator for defining a handler function.
@@ -183,6 +188,9 @@ def handler(
                 NOTE: You can set this field only if you register this service against restate-server >= 1.4,
                 otherwise the service discovery will fail.
             invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this handler.
+            explicit_cancellation: (Preview) If true, Restate cancelling an invocation neither fails the handler's
+                awaits nor cancels its calls: the handler awaits `restate._preview.cancellation(ctx)` and decides
+                what to stop. NOTE: in preview, it might change without notice.
 
         Returns:
             Callable: The decorated function.
@@ -228,6 +236,7 @@ def wrapped(*args, **kwargs):
                 ingress_private,
                 invocation_retry_policy,
                 combined_context_managers,
+                explicit_cancellation=explicit_cancellation,
             )
             if handler.name in self.handlers:
                 raise ValueError(f"Handler {handler.name} already exists in {self.name}")

python/restate/server.py

@@ -360,6 +360,7 @@ async def app(scope: Scope, receive: Receive, send: Send):
                         request_headers,
                         disable_payload_checks=disable_payload_checks,
                         explicit_recovery=handler.explicit_recovery or service.explicit_recovery,
+                        explicit_cancellation=handler.explicit_cancellation or service.explicit_cancellation,
                     ),
                     handler,
                     dict(request_headers),

python/restate/server_context.py

@@ -59,7 +59,7 @@
 from restate.serde import BytesSerde, DefaultSerde, Serde
 from restate.server_types import ReceiveChannel, Send
 from restate.types import extract_core_type
-from restate.vm import Failure, Invocation, NotReady, VMWrapper, RunRetryConfig, Suspended  # pylint: disable=line-too-long
+from restate.vm import CANCEL_HANDLE, Failure, Invocation, NotReady, VMWrapper, RunRetryConfig, Suspended  # pylint: disable=line-too-long
 from restate.vm import (
     DoProgressAnyCompleted,
     DoProgressCancelSignalReceived,
@@ -534,6 +534,10 @@ def __init__(
         self.request_finished_event = asyncio.Event()
         self.tasks = Tasks()
         self.extension_data: Dict[str, Any] = {}
+        # Explicit cancellation (preview): this attempt's calls that may be in flight, and the
+        # cancellation's future (the cancel notification can be taken only once).
+        self._calls_in_flight: List[ServerCallDurableFuture[Any]] = []
+        self._cancellation: Optional[ServerDurableFuture[None]] = None
 
     def _encode(self, buffer: bytes) -> bytes:
         """Apply the journal value codec encode step, if a codec is configured."""
@@ -763,7 +767,42 @@ async def inv_id_factory():
                 await self.create_poll_or_cancel_coroutine(SingleUnresolvedFuture(invocation_id_handle))
             return await self.must_take_notification(invocation_id_handle)
 
-        return ServerCallDurableFuture(self, handle, self._create_fetch_result_coroutine(handle, serde), inv_id_factory)
+        future: ServerCallDurableFuture[T] = ServerCallDurableFuture(
+            self, handle, self._create_fetch_result_coroutine(handle, serde), inv_id_factory
+        )
+        if self.vm.explicit_cancellation:
+            self._calls_in_flight = [call for call in self._calls_in_flight if not call.is_completed()]
+            self._calls_in_flight.append(future)
+        return future
+
+    def cancellation(self) -> RestateDurableFuture[None]:
+        """Return a future that completes when Restate cancels this invocation: see restate._preview.cancellation."""
+        self._require_explicit_cancellation("cancellation")
+        if self._cancellation is None:
+
+            async def cancelled() -> None:
+                if not self.vm.is_completed(CANCEL_HANDLE):
+                    await self.create_poll_or_cancel_coroutine(SingleUnresolvedFuture(CANCEL_HANDLE))
+                await self.must_take_notification(CANCEL_HANDLE)
+
+            self._cancellation = ServerDurableFuture(self, CANCEL_HANDLE, cancelled)
+        return self._cancellation
+
+    async def cancel_previous_calls(self) -> List[str]:
+        """Cancel this attempt's calls that are still in flight, and return their invocation ids."""
+        self._require_explicit_cancellation("cancel_previous_calls")
+        calls, self._calls_in_flight = self._calls_in_flight, []
+        cancelled: List[str] = []
+        for call in calls:
+            if not call.is_completed():
+                invocation_id = await call.invocation_id()
+                self.cancel_invocation(invocation_id)
+                cancelled.append(invocation_id)
+        return cancelled
+
+    def _require_explicit_cancellation(self, operation: str) -> None:
+        if not self.vm.explicit_cancellation:
+            raise TerminalError(f"{operation}() needs a handler with explicit_cancellation=True")
 
     def get(
         self,

python/restate/service.py

@@ -69,6 +69,8 @@ class Service:
             otherwise the service discovery will fail.
         invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this service.
         invocation_context_managers (List[Callable[[], AsyncContextManager[None]]], optional): List of context manager factories that will be applied to each invocation of this service.
+        explicit_cancellation (bool, optional): (Preview) Explicit cancellation for all handlers of this service:
+            see the handler option of the same name. NOTE: in preview, it might change without notice.
     """
 
     def __init__(
@@ -85,6 +87,7 @@ def __init__(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> None:
         self.service_tag = ServiceTag("service", name, description, metadata)
         self.handlers: typing.Dict[str, Handler] = {}
@@ -96,6 +99,7 @@ def __init__(
         self.invocation_retry_policy = invocation_retry_policy
         self.context_managers = invocation_context_managers
         self.explicit_recovery = False
+        self.explicit_cancellation = explicit_cancellation
 
     @property
     def name(self):
@@ -121,6 +125,7 @@ def handler(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> typing.Callable[[T], T]:
         """
         Decorator for defining a handler function.
@@ -159,6 +164,9 @@ def handler(
                 otherwise the service discovery will fail.
             invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this handler.
             invocation_context_managers: List of context manager factories that will be applied to each invocation of this handler.
+            explicit_cancellation: (Preview) If true, Restate cancelling an invocation neither fails the handler's
+                awaits nor cancels its calls: the handler awaits `restate._preview.cancellation(ctx)` and decides
+                what to stop. NOTE: in preview, it might change without notice.
 
         Returns:
             Callable: The decorated function.
@@ -207,6 +215,7 @@ def wrapped(*args, **kwargs):
                 ingress_private,
                 invocation_retry_policy,
                 combined_context_managers,
+                explicit_cancellation=explicit_cancellation,
             )
             if handler.name in self.handlers:
                 raise ValueError(f"Handler {handler.name} already exists in {self.name}")

python/restate/vm.py

@@ -210,8 +210,10 @@ def __init__(
         headers: typing.List[typing.Tuple[str, str]],
         disable_payload_checks: bool = False,
         explicit_recovery: bool = False,
+        explicit_cancellation: bool = False,
     ):
-        self.vm = PyVM(headers, disable_payload_checks, explicit_recovery)
+        self.vm = PyVM(headers, disable_payload_checks, explicit_recovery, explicit_cancellation)
+        self.explicit_cancellation = explicit_cancellation
 
     def get_response_head(self) -> typing.Tuple[int, typing.Iterable[typing.Tuple[str, str]]]:
         """

python/restate/workflow.py

@@ -81,6 +81,8 @@ class Workflow:
             otherwise the service discovery will fail.
         invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this workflow.
         invocation_context_managers (List[Callable[[], AsyncContextManager[None]]], optional): List of context manager factories that will be applied to each invocation of this workflow.
+        explicit_cancellation (bool, optional): (Preview) Explicit cancellation for all handlers of this workflow:
+            see the handler option of the same name. NOTE: in preview, it might change without notice.
     """
 
     handlers: typing.Dict[str, Handler[typing.Any, typing.Any]]
@@ -101,6 +103,7 @@ def __init__(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ):
         self.service_tag = ServiceTag("workflow", name, description, metadata)
         self.handlers = {}
@@ -115,6 +118,7 @@ def __init__(
         self.invocation_retry_policy = invocation_retry_policy
         self.context_managers = invocation_context_managers
         self.explicit_recovery = False
+        self.explicit_cancellation = explicit_cancellation
 
     @property
     def name(self):
@@ -142,6 +146,7 @@ def main(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> typing.Callable[[T], T]:
         """
         Mark this handler as a workflow entry point.
@@ -186,6 +191,9 @@ def main(
                 NOTE: You can set this field only if you register this service against restate-server >= 1.4,
                 otherwise the service discovery will fail.
             invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this workflow.
+            explicit_cancellation: (Preview) If true, Restate cancelling an invocation neither fails the handler's
+                awaits nor cancels its calls: the handler awaits `restate._preview.cancellation(ctx)` and decides
+                what to stop. NOTE: in preview, it might change without notice.
         """
         return self._add_handler(
             name,
@@ -205,6 +213,7 @@ def main(
             ingress_private=ingress_private,
             invocation_retry_policy=invocation_retry_policy,
             invocation_context_managers=invocation_context_managers,
+            explicit_cancellation=explicit_cancellation,
         )
 
     def handler(
@@ -226,6 +235,7 @@ def handler(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> typing.Callable[[T], T]:
         """
         Decorator for defining a handler function.
@@ -271,6 +281,9 @@ def handler(
                 otherwise the service discovery will fail.
             invocation_retry_policy (InvocationRetryPolicy, optional): Retry policy applied for all invocations to this handler.
             invocation_context_managers (List[Callable[[], AsyncContextManager[None]]], optional): List of context manager factories that will be applied to each invocation of this handler.
+            explicit_cancellation: (Preview) If true, Restate cancelling an invocation neither fails the handler's
+                awaits nor cancels its calls: the handler awaits `restate._preview.cancellation(ctx)` and decides
+                what to stop. NOTE: in preview, it might change without notice.
         """
         return self._add_handler(
             name,
@@ -290,6 +303,7 @@ def handler(
             ingress_private,
             invocation_retry_policy,
             invocation_context_managers,
+            explicit_cancellation=explicit_cancellation,
         )
 
     # pylint: disable=R0914
@@ -314,6 +328,7 @@ def _add_handler(
         invocation_context_managers: typing.Optional[
             typing.List[typing.Callable[[], typing.AsyncContextManager[None]]]
         ] = None,
+        explicit_cancellation: bool = False,
     ) -> typing.Callable[[T], T]:
         """
         Decorator for defining a handler function.
@@ -408,6 +423,7 @@ def wrapped(*args, **kwargs):
                 ingress_private=ingress_private,
                 invocation_retry_policy=invocation_retry_policy,
                 context_managers=combined_context_managers,
+                explicit_cancellation=explicit_cancellation,
             )
             if handler.name in self.handlers:
                 raise ValueError(f"Handler {handler.name} already exists in {self.name}")