From 0d3f5a4068592871f9c98531d26029872dd19139 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:59:53 +0100 Subject: [PATCH 1/2] Add explicit cancellation (preview) With explicit cancellation, Restate cancelling an invocation neither fails the handler's awaits nor cancels the calls it made: the handler finds out by awaiting cancellation(ctx), and decides what to stop. @explicit_cancellation @greeter.handler() async def greet(ctx, name): done, _ = await restate.select(work=ctx.service_call(slow, arg=name), cancelled=cancellation(ctx)) if done == "cancelled": await cancel_previous_calls(ctx) # the calls still in flight raise TerminalError("cancelled", status_code=409) In restate._preview, next to explicit_recovery: on a handler, or a whole service, object or workflow. The VM disables shared-core's implicit cancellation; the context keeps the attempt's calls that may be in flight (only with the option on), which cancel_previous_calls cancels. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/_preview.py | 61 +++++++++++ python/restate/handler.py | 3 + python/restate/object.py | 1 + python/restate/server.py | 1 + python/restate/server_context.py | 43 +++++++- python/restate/service.py | 1 + python/restate/vm.py | 4 +- python/restate/workflow.py | 1 + src/lib.rs | 18 +++- tests/explicit_cancellation.py | 169 +++++++++++++++++++++++++++++++ 10 files changed, 294 insertions(+), 8 deletions(-) create mode 100644 tests/explicit_cancellation.py diff --git a/python/restate/_preview.py b/python/restate/_preview.py index eda5a524..2e5f9e3f 100644 --- a/python/restate/_preview.py +++ b/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,61 @@ async def greet(ctx: Context, name: str) -> str: else: handler_from_callable(target).explicit_recovery = True # type: ignore return target + + +def explicit_cancellation(target: T) -> T: + """ + Run the handlers with explicit cancellation: when Restate cancels an invocation, no await fails and none + of its calls is cancelled. The handler finds out with `cancellation(ctx)` and decides what to stop. + + Apply it to a service, virtual object or workflow to enable it for all its handlers, + or as a decorator on top of a handler to enable it for that handler only. + + Example: + @explicit_cancellation + @greeter.handler() + 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) + ... + + NOTE: This API is in preview and might change without notice. + """ + if isinstance(target, (Service, VirtualObject, Workflow)): + target.explicit_cancellation = True + else: + handler_from_callable(target).explicit_cancellation = 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`. The handler must run + with `explicit_cancellation`. + + 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`. + + 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 diff --git a/python/restate/handler.py b/python/restate/handler.py index 3253b7e6..e2f58881 100644 --- a/python/restate/handler.py +++ b/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, + see restate._preview.explicit_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 diff --git a/python/restate/object.py b/python/restate/object.py index ffde1d4f..70c6ddfa 100644 --- a/python/restate/object.py +++ b/python/restate/object.py @@ -109,6 +109,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False + self.explicit_cancellation = False @property def name(self): diff --git a/python/restate/server.py b/python/restate/server.py index 723858a7..77cc1830 100644 --- a/python/restate/server.py +++ b/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), diff --git a/python/restate/server_context.py b/python/restate/server_context.py index 6a996bd4..a2ff9129 100644 --- a/python/restate/server_context.py +++ b/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 (restate._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 restate._preview.explicit_cancellation") def get( self, diff --git a/python/restate/service.py b/python/restate/service.py index ea695430..e736a7d6 100644 --- a/python/restate/service.py +++ b/python/restate/service.py @@ -96,6 +96,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False + self.explicit_cancellation = False @property def name(self): diff --git a/python/restate/vm.py b/python/restate/vm.py index 96e8a9b8..85be9830 100644 --- a/python/restate/vm.py +++ b/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]]]: """ diff --git a/python/restate/workflow.py b/python/restate/workflow.py index e1d0056b..030d1813 100644 --- a/python/restate/workflow.py +++ b/python/restate/workflow.py @@ -115,6 +115,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False + self.explicit_cancellation = False @property def name(self): diff --git a/src/lib.rs b/src/lib.rs index 57ecc9e4..5f18c254 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,10 +6,10 @@ use pyo3::prelude::*; use pyo3::types::{PyBytes, PyNone, PyString}; use restate_sdk_shared_core::fmt::{set_error_formatter, ErrorFormatter}; use restate_sdk_shared_core::{ - AwaitResponse, AwakeableHandle, CallHandle, CoreVM, Error, Header, IdentityVerifier, Input, - JournalReplay, NonDeterministicChecksOption, NonEmptyValue, NotificationHandle, OnMaxAttempts, - ResponseHead, RetryPolicy, RunExitResult, RunHandle, Target, TerminalFailure, UnresolvedFuture, - VMOptions, Value, CANCEL_NOTIFICATION_HANDLE, VM, + AwaitResponse, AwakeableHandle, CallHandle, CoreVM, Error, Header, IdentityVerifier, + ImplicitCancellationOption, Input, JournalReplay, NonDeterministicChecksOption, NonEmptyValue, + NotificationHandle, OnMaxAttempts, ResponseHead, RetryPolicy, RunExitResult, RunHandle, Target, + TerminalFailure, UnresolvedFuture, VMOptions, Value, CANCEL_NOTIFICATION_HANDLE, VM, }; use std::fmt; use std::time::{Duration, SystemTime}; @@ -373,11 +373,12 @@ struct PyVM { #[pymethods] impl PyVM { #[new] - #[pyo3(signature = (headers, disable_payload_checks=false, explicit_recovery=false))] + #[pyo3(signature = (headers, disable_payload_checks=false, explicit_recovery=false, explicit_cancellation=false))] fn new( headers: Vec<(String, String)>, disable_payload_checks: bool, explicit_recovery: bool, + explicit_cancellation: bool, ) -> Result { // When a journal value codec is configured, the serialized payloads written to the journal // may legitimately differ between attempts (e.g. encryption with a random nonce), so we @@ -393,6 +394,13 @@ impl PyVM { } else { JournalReplay::Enabled }, + // With explicit cancellation, a cancellation neither fails the handler's awaits nor + // cancels its calls: the handler awaits CANCEL_NOTIFICATION_HANDLE and decides. + implicit_cancellation: if explicit_cancellation { + ImplicitCancellationOption::Disabled + } else { + VMOptions::default().implicit_cancellation + }, ..Default::default() }; Ok(Self { diff --git a/tests/explicit_cancellation.py b/tests/explicit_cancellation.py new file mode 100644 index 00000000..900d7037 --- /dev/null +++ b/tests/explicit_cancellation.py @@ -0,0 +1,169 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Explicit cancellation (preview): a cancelled handler is told, and decides what to stop.""" + +# pylint: disable=missing-function-docstring + +import asyncio +from datetime import timedelta + +import httpx +import pytest + +import restate +from restate import Context, HarnessEnvironment, Service, VirtualObject, Workflow +from restate._preview import cancel_previous_calls, cancellation, explicit_cancellation + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +# --- Configuration ------------------------------------------------------------------------------ + + +def test_it_is_off_by_default(): + greeter = Service("greeter") + + @greeter.handler() + async def greet(ctx: Context, name: str) -> str: + return name + + assert not greeter.explicit_cancellation + assert not greeter.handlers["greet"].explicit_cancellation + + +def test_it_applies_to_a_whole_service_or_to_one_handler(): + counter = explicit_cancellation(VirtualObject("counter")) + assert counter.explicit_cancellation + + flow = Workflow("flow") + + @explicit_cancellation + @flow.main() + async def run(ctx: restate.WorkflowContext) -> None: + pass + + assert flow.handlers["run"].explicit_cancellation + + +def test_cancellation_needs_the_context_of_a_handler(): + with pytest.raises(restate.TerminalError, match="takes the context a handler is invoked with"): + cancellation(object()) # type: ignore[arg-type] + + +# --- Against a Restate server ------------------------------------------------------------------- + +sleeper = Service("sleeper") + + +@sleeper.handler() +async def nap(ctx: Context) -> str: + await ctx.sleep(timedelta(minutes=5)) + return "rested" + + +@sleeper.handler() +async def echo(ctx: Context, n: int) -> int: + return n + + +explicit = explicit_cancellation(Service("explicit")) + + +@explicit.handler() +async def wait_for_nap(ctx: Context) -> list[str]: + """Wait for a nap; when cancelled, cancel the nap and return the ids of the calls cancelled.""" + assert cancellation(ctx) is cancellation(ctx) + done, _ = await restate.select(nap=ctx.service_call(nap, arg=None), cancelled=cancellation(ctx)) + assert done == "cancelled" + return await cancel_previous_calls(ctx) + + +@explicit.handler() +async def echo_twice(ctx: Context) -> list[str]: + """Make two calls that complete: there is nothing left to cancel.""" + await ctx.service_call(echo, arg=1) + await ctx.service_call(echo, arg=2) + return await cancel_previous_calls(ctx) + + +implicit = Service("implicit") + + +@implicit.handler() +async def nap_implicitly(ctx: Context) -> None: + await ctx.service_call(nap, arg=None) + + +@implicit.handler() +async def ask_for_cancellation(ctx: Context) -> None: + await cancellation(ctx) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + restate.app([sleeper, explicit, implicit]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +async def cancel_once_running(harness: HarnessEnvironment, invocation_id: str) -> None: + async with httpx.AsyncClient() as admin: + for _ in range(100): + status = await admin.post( + f"{harness.admin_api_url}/query", + json={"query": f"SELECT status FROM sys_invocation WHERE id = '{invocation_id}'"}, + headers={"Accept": "application/json"}, + ) + rows = status.json().get("rows", []) + if rows and rows[0]["status"] in ("running", "suspended"): + break + await asyncio.sleep(0.05) + (await admin.patch(f"{harness.admin_api_url}/invocations/{invocation_id}/cancel")).raise_for_status() + + +async def result_of(harness: HarnessEnvironment, invocation_id: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=30) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation_id}/attach") + + +@pytest.mark.anyio +async def test_a_cancelled_handler_goes_on_and_cancels_its_calls_itself(harness: HarnessEnvironment): + sent = await harness.client.service_send(wait_for_nap, arg=None) + await cancel_once_running(harness, sent.invocation_id) + + answer = await result_of(harness, sent.invocation_id) + assert answer.status_code == 200, answer.text + [nap_invocation] = answer.json() + assert (await result_of(harness, nap_invocation)).status_code == 409 + + +@pytest.mark.anyio +async def test_calls_that_completed_are_not_cancelled(harness: HarnessEnvironment): + assert await harness.client.service_call(echo_twice, arg=None) == [] + + +@pytest.mark.anyio +async def test_without_it_the_cancellation_fails_the_handler(harness: HarnessEnvironment): + sent = await harness.client.service_send(nap_implicitly, arg=None) + await cancel_once_running(harness, sent.invocation_id) + assert (await result_of(harness, sent.invocation_id)).status_code == 409 + + +@pytest.mark.anyio +async def test_without_it_cancellation_is_refused(harness: HarnessEnvironment): + sent = await harness.client.service_send(ask_for_cancellation, arg=None) + failed = await result_of(harness, sent.invocation_id) + assert failed.status_code == 500 + assert "needs a handler with restate._preview.explicit_cancellation" in failed.text From 5f2ab582807b5fc1ce06812231a2626387b5a7ac Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 19:12:49 +0100 Subject: [PATCH 2/2] explicit_cancellation is a key of the handler and service options Like every other option: @svc.handler(explicit_cancellation=True), a workflow's main()/handler() too, and Service/VirtualObject/Workflow(..., explicit_cancellation=True) for all their handlers. The restate._preview.explicit_cancellation decorator is gone; cancellation() and cancel_previous_calls() stay in restate._preview. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/_preview.py | 31 +++++++------------------------ python/restate/handler.py | 7 +++++-- python/restate/object.py | 10 +++++++++- python/restate/server_context.py | 4 ++-- python/restate/service.py | 10 +++++++++- python/restate/workflow.py | 17 ++++++++++++++++- tests/explicit_cancellation.py | 31 ++++++++++++++++++++++--------- 7 files changed, 70 insertions(+), 40 deletions(-) diff --git a/python/restate/_preview.py b/python/restate/_preview.py index 2e5f9e3f..78bbb6c1 100644 --- a/python/restate/_preview.py +++ b/python/restate/_preview.py @@ -58,17 +58,13 @@ async def greet(ctx: Context, name: str) -> str: return target -def explicit_cancellation(target: T) -> T: +def cancellation(ctx: Context) -> RestateDurableFuture[None]: """ - Run the handlers with explicit cancellation: when Restate cancels an invocation, no await fails and none - of its calls is cancelled. The handler finds out with `cancellation(ctx)` and decides what to stop. + Return a future that completes when Restate cancels the invocation of `ctx`. - Apply it to a service, virtual object or workflow to enable it for all its handlers, - or as a decorator on top of a handler to enable it for that handler only. + Await it together with the handler's other work, for example with `restate.select`:: - Example: - @explicit_cancellation - @greeter.handler() + @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)) @@ -77,21 +73,8 @@ async def greet(ctx: Context, name: str) -> str: raise TerminalError("cancelled", status_code=409) ... - NOTE: This API is in preview and might change without notice. - """ - if isinstance(target, (Service, VirtualObject, Workflow)): - target.explicit_cancellation = True - else: - handler_from_callable(target).explicit_cancellation = 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`. The handler must run - with `explicit_cancellation`. + 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. """ @@ -103,7 +86,7 @@ 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`. + ids of the cancelled calls. The handler must run with ``explicit_cancellation=True``. NOTE: This API is in preview and might change without notice. """ diff --git a/python/restate/handler.py b/python/restate/handler.py index e2f58881..b274ba51 100644 --- a/python/restate/handler.py +++ b/python/restate/handler.py @@ -256,8 +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, - see restate._preview.explicit_cancellation. + explicit_cancellation: If true, the handler runs with explicit cancellation (preview): + see restate._preview.cancellation. """ service_tag: ServiceTag @@ -304,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. @@ -328,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 @@ -363,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 diff --git a/python/restate/object.py b/python/restate/object.py index 70c6ddfa..06b0b6ed 100644 --- a/python/restate/object.py +++ b/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,7 +112,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False - self.explicit_cancellation = False + self.explicit_cancellation = explicit_cancellation @property def name(self): @@ -139,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. @@ -184,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. @@ -229,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}") diff --git a/python/restate/server_context.py b/python/restate/server_context.py index a2ff9129..e7e0dc5f 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -534,7 +534,7 @@ def __init__( self.request_finished_event = asyncio.Event() self.tasks = Tasks() self.extension_data: Dict[str, Any] = {} - # Explicit cancellation (restate._preview): this attempt's calls that may be in flight, and the + # 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 @@ -802,7 +802,7 @@ async def cancel_previous_calls(self) -> List[str]: def _require_explicit_cancellation(self, operation: str) -> None: if not self.vm.explicit_cancellation: - raise TerminalError(f"{operation}() needs a handler with restate._preview.explicit_cancellation") + raise TerminalError(f"{operation}() needs a handler with explicit_cancellation=True") def get( self, diff --git a/python/restate/service.py b/python/restate/service.py index e736a7d6..f9eb7827 100644 --- a/python/restate/service.py +++ b/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,7 +99,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False - self.explicit_cancellation = False + self.explicit_cancellation = explicit_cancellation @property def name(self): @@ -122,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. @@ -160,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. @@ -208,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}") diff --git a/python/restate/workflow.py b/python/restate/workflow.py index 030d1813..4a20173d 100644 --- a/python/restate/workflow.py +++ b/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,7 +118,7 @@ def __init__( self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers self.explicit_recovery = False - self.explicit_cancellation = False + self.explicit_cancellation = explicit_cancellation @property def name(self): @@ -143,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. @@ -187,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, @@ -206,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( @@ -227,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. @@ -272,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, @@ -291,6 +303,7 @@ def handler( ingress_private, invocation_retry_policy, invocation_context_managers, + explicit_cancellation=explicit_cancellation, ) # pylint: disable=R0914 @@ -315,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. @@ -409,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}") diff --git a/tests/explicit_cancellation.py b/tests/explicit_cancellation.py index 900d7037..80b485ad 100644 --- a/tests/explicit_cancellation.py +++ b/tests/explicit_cancellation.py @@ -20,7 +20,7 @@ import restate from restate import Context, HarnessEnvironment, Service, VirtualObject, Workflow -from restate._preview import cancel_previous_calls, cancellation, explicit_cancellation +from restate._preview import cancel_previous_calls, cancellation @pytest.fixture(scope="session") @@ -42,18 +42,28 @@ async def greet(ctx: Context, name: str) -> str: assert not greeter.handlers["greet"].explicit_cancellation -def test_it_applies_to_a_whole_service_or_to_one_handler(): - counter = explicit_cancellation(VirtualObject("counter")) +def test_it_is_a_key_of_a_whole_service_or_of_one_handler(): + assert Service("greeter", explicit_cancellation=True).explicit_cancellation + counter = VirtualObject("counter", explicit_cancellation=True) assert counter.explicit_cancellation + @counter.handler(explicit_cancellation=True) + async def add(ctx: restate.ObjectContext, n: int) -> None: + pass + flow = Workflow("flow") - @explicit_cancellation - @flow.main() + @flow.main(explicit_cancellation=True) async def run(ctx: restate.WorkflowContext) -> None: pass + @flow.handler(explicit_cancellation=True) + async def poke(ctx: restate.WorkflowSharedContext) -> None: + pass + + assert counter.handlers["add"].explicit_cancellation assert flow.handlers["run"].explicit_cancellation + assert flow.handlers["poke"].explicit_cancellation def test_cancellation_needs_the_context_of_a_handler(): @@ -77,7 +87,7 @@ async def echo(ctx: Context, n: int) -> int: return n -explicit = explicit_cancellation(Service("explicit")) +explicit = Service("explicit", explicit_cancellation=True) @explicit.handler() @@ -89,7 +99,10 @@ async def wait_for_nap(ctx: Context) -> list[str]: return await cancel_previous_calls(ctx) -@explicit.handler() +on_one_handler = Service("on_one_handler") + + +@on_one_handler.handler(explicit_cancellation=True) async def echo_twice(ctx: Context) -> list[str]: """Make two calls that complete: there is nothing left to cancel.""" await ctx.service_call(echo, arg=1) @@ -113,7 +126,7 @@ async def ask_for_cancellation(ctx: Context) -> None: @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([sleeper, explicit, implicit]), restate_image="ghcr.io/restatedev/restate:latest" + restate.app([sleeper, explicit, on_one_handler, implicit]), restate_image="ghcr.io/restatedev/restate:latest" ) as harness: yield harness @@ -166,4 +179,4 @@ async def test_without_it_cancellation_is_refused(harness: HarnessEnvironment): sent = await harness.client.service_send(ask_for_cancellation, arg=None) failed = await result_of(harness, sent.invocation_id) assert failed.status_code == 500 - assert "needs a handler with restate._preview.explicit_cancellation" in failed.text + assert "needs a handler with explicit_cancellation=True" in failed.text