From 7bb1de1ce46b29a4dfa890aba8a758f8a4d98b29 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Wed, 21 Oct 2026 09:07:22 +0200 Subject: [PATCH] Add entry names and explicit recovery preview --- Cargo.lock | 4 +- Cargo.toml | 2 +- python/restate/_preview.py | 55 +++++++++++++++ python/restate/context.py | 59 +++++++++++++--- python/restate/handler.py | 2 + python/restate/object.py | 1 + python/restate/server.py | 6 +- python/restate/server_context.py | 24 ++++--- python/restate/service.py | 1 + python/restate/vm.py | 34 ++++++---- python/restate/workflow.py | 1 + src/lib.rs | 54 +++++++++++---- tests/explicit_recovery.py | 112 +++++++++++++++++++++++++++++++ 13 files changed, 306 insertions(+), 49 deletions(-) create mode 100644 python/restate/_preview.py create mode 100644 tests/explicit_recovery.py diff --git a/Cargo.lock b/Cargo.lock index 7201e4e..1f36c52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -814,9 +814,9 @@ dependencies = [ [[package]] name = "restate-sdk-shared-core" -version = "7.0.3" +version = "8.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4a8731358ce185d1c9663d3ac97414acf7c0d642e5724864ade03fcccb85134" +checksum = "508098a994e131960dfa3ba4306a50605fdc1cc281875dfef28926c5b88c3b9b" dependencies = [ "base64", "bs58", diff --git a/Cargo.toml b/Cargo.toml index ce23f1c..8339993 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,4 +14,4 @@ doc = false [dependencies] pyo3 = { version = "0.25.1", features = ["extension-module"] } tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } -restate-sdk-shared-core = { version = "7.0.3", features = ["request_identity", "sha2_random_seed", "rust_crypto"] } +restate-sdk-shared-core = { version = "8.1.0", features = ["request_identity", "sha2_random_seed", "rust_crypto"] } diff --git a/python/restate/_preview.py b/python/restate/_preview.py new file mode 100644 index 0000000..eda5a52 --- /dev/null +++ b/python/restate/_preview.py @@ -0,0 +1,55 @@ +# +# Copyright (c) 2023-2024 - 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 +# +"""Preview apis: they might change or be removed without notice.""" + +import typing + +from .handler import handler_from_callable +from .service import Service +from .object import VirtualObject +from .workflow import Workflow + +T = typing.TypeVar("T") + + +def explicit_recovery(target: T) -> T: + """ + Run the handlers with explicit recovery: a retried or resumed invocation doesn't replay the journal of its + earlier attempts, but runs the handler again from the start, executing every context operation anew, + e.g. reading the current state. + + The handler takes care of recovering from the effects of the earlier attempts, which might repeat (at least once): + + * Each attempt's calls and sends are new requests. To reach the invocation an earlier attempt started, + give it an idempotency key that is the same on every attempt, such as one made from `ctx.request().id`. + * `ctx.run` executes its action again on every attempt. + * Awakeables are not supported. + + 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: + counter = explicit_recovery(VirtualObject("counter")) + + greeter = Service("greeter") + + @explicit_recovery + @greeter.handler() + async def greet(ctx: Context, name: str) -> str: + ... + + NOTE: This API is in preview and might change without notice. + """ + if isinstance(target, (Service, VirtualObject, Workflow)): + target.explicit_recovery = True + else: + handler_from_callable(target).explicit_recovery = True # type: ignore + return target diff --git a/python/restate/context.py b/python/restate/context.py index feb0086..0ca0ffb 100644 --- a/python/restate/context.py +++ b/python/restate/context.py @@ -178,7 +178,11 @@ class KeyValueStore(abc.ABC): @abc.abstractmethod def get( - self, name: str, serde: Serde[T] = DefaultSerde(), type_hint: Optional[typing.Type[T]] = None + self, + name: str, + serde: Serde[T] = DefaultSerde(), + type_hint: Optional[typing.Type[T]] = None, + entry_name: Optional[str] = None, ) -> Awaitable[Optional[T]]: """ Retrieves the value associated with the given name. @@ -188,23 +192,48 @@ def get( serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type. See also 'type_hint'. type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer. + entry_name: The name of the journal entry, shown in the Restate journal for observability. """ @abc.abstractmethod - def state_keys(self) -> Awaitable[List[str]]: - """Returns the list of keys in the store.""" + def state_keys(self, entry_name: Optional[str] = None) -> Awaitable[List[str]]: + """ + Returns the list of keys in the store. + + Args: + entry_name: The name of the journal entry, shown in the Restate journal for observability. + """ @abc.abstractmethod - def set(self, name: str, value: T, serde: Serde[T] = DefaultSerde()) -> None: - """set the value associated with the given name.""" + def set(self, name: str, value: T, serde: Serde[T] = DefaultSerde(), entry_name: Optional[str] = None) -> None: + """ + set the value associated with the given name. + + Args: + name: The state name + value: The value to set + serde: The serialization/deserialization mechanism. + entry_name: The name of the journal entry, shown in the Restate journal for observability. + """ @abc.abstractmethod - def clear(self, name: str) -> None: - """clear the value associated with the given name.""" + def clear(self, name: str, entry_name: Optional[str] = None) -> None: + """ + clear the value associated with the given name. + + Args: + name: The state name + entry_name: The name of the journal entry, shown in the Restate journal for observability. + """ @abc.abstractmethod - def clear_all(self) -> None: - """clear all the values in the store.""" + def clear_all(self, entry_name: Optional[str] = None) -> None: + """ + clear all the values in the store. + + Args: + entry_name: The name of the journal entry, shown in the Restate journal for observability. + """ # pylint: disable=R0903 @@ -797,7 +826,11 @@ def key(self) -> str: @abc.abstractmethod def get( - self, name: str, serde: Serde[T] = DefaultSerde(), type_hint: Optional[typing.Type[T]] = None + self, + name: str, + serde: Serde[T] = DefaultSerde(), + type_hint: Optional[typing.Type[T]] = None, + entry_name: Optional[str] = None, ) -> RestateDurableFuture[Optional[T]]: """ Retrieves the value associated with the given name. @@ -807,12 +840,16 @@ def get( serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type. See also 'type_hint'. type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer. + entry_name: The name of the journal entry, shown in the Restate journal for observability. """ @abc.abstractmethod - def state_keys(self) -> Awaitable[List[str]]: + def state_keys(self, entry_name: Optional[str] = None) -> Awaitable[List[str]]: """ Returns the list of keys in the store. + + Args: + entry_name: The name of the journal entry, shown in the Restate journal for observability. """ diff --git a/python/restate/handler.py b/python/restate/handler.py index 91a55a7..306e12c 100644 --- a/python/restate/handler.py +++ b/python/restate/handler.py @@ -239,6 +239,7 @@ class Handler(Generic[I, O]): enable_lazy_state: If true, lazy state is enabled. 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. """ service_tag: ServiceTag @@ -258,6 +259,7 @@ class Handler(Generic[I, O]): ingress_private: Optional[bool] = None invocation_retry_policy: Optional[InvocationRetryPolicy] = None context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None + explicit_recovery: bool = False # disable too many arguments warning diff --git a/python/restate/object.py b/python/restate/object.py index 14d5ef5..175c30d 100644 --- a/python/restate/object.py +++ b/python/restate/object.py @@ -101,6 +101,7 @@ def __init__( self.ingress_private = ingress_private self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers + self.explicit_recovery = False @property def name(self): diff --git a/python/restate/server.py b/python/restate/server.py index c5e233f..84d3c7c 100644 --- a/python/restate/server.py +++ b/python/restate/server.py @@ -320,7 +320,11 @@ async def app(scope: Scope, receive: Receive, send: Send): try: journal_codec = await get_codec() await process_invocation_to_completion( - VMWrapper(request_headers, disable_payload_checks=disable_payload_checks), + VMWrapper( + request_headers, + disable_payload_checks=disable_payload_checks, + explicit_recovery=handler.explicit_recovery or service.explicit_recovery, + ), handler, dict(request_headers), receive_channel, diff --git a/python/restate/server_context.py b/python/restate/server_context.py index 80fbec3..6a996bd 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -766,33 +766,37 @@ async def inv_id_factory(): return ServerCallDurableFuture(self, handle, self._create_fetch_result_coroutine(handle, serde), inv_id_factory) def get( - self, name: str, serde: Serde[T] = DefaultSerde(), type_hint: Optional[typing.Type[T]] = None + self, + name: str, + serde: Serde[T] = DefaultSerde(), + type_hint: Optional[typing.Type[T]] = None, + entry_name: Optional[str] = None, ) -> Awaitable[Optional[T]]: - handle = self.vm.sys_get_state(name) + handle = self.vm.sys_get_state(name, entry_name) update_restate_context_is_replaying(self.vm) if isinstance(serde, DefaultSerde): serde = serde.with_maybe_type(type_hint) return self.create_future(handle, serde) # type: ignore - def state_keys(self) -> Awaitable[List[str]]: - handle = self.vm.sys_get_state_keys() + def state_keys(self, entry_name: Optional[str] = None) -> Awaitable[List[str]]: + handle = self.vm.sys_get_state_keys(entry_name) update_restate_context_is_replaying(self.vm) return self.create_future(handle) - def set(self, name: str, value: T, serde: Serde[T] = DefaultSerde()) -> None: + def set(self, name: str, value: T, serde: Serde[T] = DefaultSerde(), entry_name: Optional[str] = None) -> None: """Set the value associated with the given name.""" if isinstance(serde, DefaultSerde): serde = serde.with_maybe_type(type(value)) buffer = self._encode(serde.serialize(value)) - self.vm.sys_set_state(name, bytes(buffer)) + self.vm.sys_set_state(name, bytes(buffer), entry_name) update_restate_context_is_replaying(self.vm) - def clear(self, name: str) -> None: - self.vm.sys_clear_state(name) + def clear(self, name: str, entry_name: Optional[str] = None) -> None: + self.vm.sys_clear_state(name, entry_name) update_restate_context_is_replaying(self.vm) - def clear_all(self) -> None: - self.vm.sys_clear_all_state() + def clear_all(self, entry_name: Optional[str] = None) -> None: + self.vm.sys_clear_all_state(entry_name) update_restate_context_is_replaying(self.vm) def request(self) -> Request: diff --git a/python/restate/service.py b/python/restate/service.py index cecd90f..357c33b 100644 --- a/python/restate/service.py +++ b/python/restate/service.py @@ -95,6 +95,7 @@ def __init__( self.ingress_private = ingress_private self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers + self.explicit_recovery = False @property def name(self): diff --git a/python/restate/vm.py b/python/restate/vm.py index 7b7b91f..96e8a9b 100644 --- a/python/restate/vm.py +++ b/python/restate/vm.py @@ -205,8 +205,13 @@ class VMWrapper: It provides a type-friendly interface to our shared vm. """ - def __init__(self, headers: typing.List[typing.Tuple[str, str]], disable_payload_checks: bool = False): - self.vm = PyVM(headers, disable_payload_checks) + def __init__( + self, + headers: typing.List[typing.Tuple[str, str]], + disable_payload_checks: bool = False, + explicit_recovery: bool = False, + ): + self.vm = PyVM(headers, disable_payload_checks, explicit_recovery) def get_response_head(self) -> typing.Tuple[int, typing.Iterable[typing.Tuple[str, str]]]: """ @@ -345,47 +350,52 @@ def sys_write_output_failure(self, output: Failure): res = PyFailure(output.code, output.message, metadata=metadata) self.vm.sys_write_output_failure(res) - def sys_get_state(self, name) -> int: + def sys_get_state(self, name, entry_name: typing.Optional[str] = None) -> int: """ Retrieves a key-value binding. Args: name: The name of the value to be retrieved. + entry_name: The name of the journal entry. Returns: The value associated with the given name. """ - return self.vm.sys_get_state(name) + return self.vm.sys_get_state(name, entry_name) - def sys_get_state_keys(self) -> int: + def sys_get_state_keys(self, entry_name: typing.Optional[str] = None) -> int: """ Retrieves all keys. + Args: + entry_name: The name of the journal entry. + Returns: The state keys """ - return self.vm.sys_get_state_keys() + return self.vm.sys_get_state_keys(entry_name) - def sys_set_state(self, name: str, value: bytes): + def sys_set_state(self, name: str, value: bytes, entry_name: typing.Optional[str] = None): """ Sets a key-value binding. Args: name: The name of the value to be set. value: The value to be set. + entry_name: The name of the journal entry. Returns: None """ - self.vm.sys_set_state(name, value) + self.vm.sys_set_state(name, value, entry_name) - def sys_clear_state(self, name: str): + def sys_clear_state(self, name: str, entry_name: typing.Optional[str] = None): """Clear the state associated with the given name.""" - self.vm.sys_clear_state(name) + self.vm.sys_clear_state(name, entry_name) - def sys_clear_all_state(self): + def sys_clear_all_state(self, entry_name: typing.Optional[str] = None): """Clear the state associated with the given name.""" - self.vm.sys_clear_all_state() + self.vm.sys_clear_all_state(entry_name) def sys_sleep(self, millis: int, name: typing.Optional[str] = None): """Ask to sleep for a given duration""" diff --git a/python/restate/workflow.py b/python/restate/workflow.py index b87a996..aa94039 100644 --- a/python/restate/workflow.py +++ b/python/restate/workflow.py @@ -107,6 +107,7 @@ def __init__( self.ingress_private = ingress_private self.invocation_retry_policy = invocation_retry_policy self.context_managers = invocation_context_managers + self.explicit_recovery = False @property def name(self): diff --git a/src/lib.rs b/src/lib.rs index 319281a..2f33781 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,9 +4,9 @@ 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, - NonDeterministicChecksOption, NonEmptyValue, NotificationHandle, OnMaxAttempts, ResponseHead, - RetryPolicy, RunExitResult, RunHandle, Target, TerminalFailure, UnresolvedFuture, VMOptions, - Value, CANCEL_NOTIFICATION_HANDLE, VM, + 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}; @@ -72,7 +72,6 @@ impl From for PyResponseHead { } } - type PyNotificationHandle = u32; #[pyclass] @@ -371,10 +370,11 @@ struct PyVM { #[pymethods] impl PyVM { #[new] - #[pyo3(signature = (headers, disable_payload_checks=false))] + #[pyo3(signature = (headers, disable_payload_checks=false, explicit_recovery=false))] fn new( headers: Vec<(String, String)>, disable_payload_checks: bool, + explicit_recovery: 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 @@ -385,6 +385,11 @@ impl PyVM { } else { NonDeterministicChecksOption::Enabled }, + journal_replay: if explicit_recovery { + JournalReplay::Disabled + } else { + JournalReplay::Enabled + }, ..Default::default() }; Ok(Self { @@ -468,6 +473,11 @@ impl PyVM { Ok(AwaitResponse::CancelSignalReceived) => { Ok(Bound::new(py, PyDoProgressCancelSignalReceived)?.into_any()) } + // The SDK doesn't issue ephemeral commands + Ok(AwaitResponse::EphemeralNotificationReady(_)) => Err(PyVMError::from(Error::new( + restate_sdk_shared_core::error::codes::INTERNAL, + "unexpected ephemeral notification", + )))?, } } @@ -510,44 +520,64 @@ impl PyVM { self_.vm.sys_input().map(Into::into).map_err(Into::into) } + #[pyo3(signature = (key, name=None))] fn sys_get_state( mut self_: PyRefMut<'_, Self>, key: String, + name: Option, ) -> Result { self_ .vm - .sys_state_get(key, Default::default()) + .sys_state_get(key, name, Default::default()) .map(Into::into) .map_err(Into::into) } + #[pyo3(signature = (name=None))] fn sys_get_state_keys( mut self_: PyRefMut<'_, Self>, + name: Option, ) -> Result { self_ .vm - .sys_state_get_keys() + .sys_state_get_keys(name) .map(Into::into) .map_err(Into::into) } + #[pyo3(signature = (key, buffer, name=None))] fn sys_set_state( mut self_: PyRefMut<'_, Self>, key: String, buffer: &Bound<'_, PyBytes>, + name: Option, ) -> Result<(), PyVMError> { self_ .vm - .sys_state_set(key, buffer.as_bytes().to_vec().into(), Default::default()) + .sys_state_set( + key, + buffer.as_bytes().to_vec().into(), + name, + Default::default(), + ) .map_err(Into::into) } - fn sys_clear_state(mut self_: PyRefMut<'_, Self>, key: String) -> Result<(), PyVMError> { - self_.vm.sys_state_clear(key).map_err(Into::into) + #[pyo3(signature = (key, name=None))] + fn sys_clear_state( + mut self_: PyRefMut<'_, Self>, + key: String, + name: Option, + ) -> Result<(), PyVMError> { + self_.vm.sys_state_clear(key, name).map_err(Into::into) } - fn sys_clear_all_state(mut self_: PyRefMut<'_, Self>) -> Result<(), PyVMError> { - self_.vm.sys_state_clear_all().map_err(Into::into) + #[pyo3(signature = (name=None))] + fn sys_clear_all_state( + mut self_: PyRefMut<'_, Self>, + name: Option, + ) -> Result<(), PyVMError> { + self_.vm.sys_state_clear_all(name).map_err(Into::into) } fn sys_sleep( diff --git a/tests/explicit_recovery.py b/tests/explicit_recovery.py new file mode 100644 index 0000000..6ea8cdc --- /dev/null +++ b/tests/explicit_recovery.py @@ -0,0 +1,112 @@ +# +# Copyright (c) 2023-2025 - 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 +# + +from datetime import timedelta + +import httpx +import pytest + +import restate +from restate import Context, ObjectContext, Service, VirtualObject, Workflow, WorkflowContext +from restate._preview import explicit_recovery +from restate.vm import VMWrapper + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +def test_is_off_by_default(): + greeter = Service("greeter") + + @greeter.handler() + async def greet(ctx: Context, name: str) -> str: + return name + + assert not greeter.explicit_recovery + assert not greeter.handlers["greet"].explicit_recovery + + +def test_applies_to_every_handler_from_the_service(): + counter = explicit_recovery(VirtualObject("counter")) + + @counter.handler() + async def add(ctx: ObjectContext, n: int) -> None: + pass + + assert counter.explicit_recovery + + +def test_applies_to_a_single_handler(): + flow = Workflow("flow") + + @explicit_recovery + @flow.main() + async def run(ctx: WorkflowContext) -> None: + pass + + @flow.handler() + async def query(ctx: WorkflowContext) -> None: + pass + + assert not flow.explicit_recovery + assert flow.handlers["run"].explicit_recovery + assert not flow.handlers["query"].explicit_recovery + + +def test_must_be_applied_on_top_of_the_handler_decorator(): + async def greet(ctx: Context, name: str) -> str: + return name + + with pytest.raises(ValueError): + explicit_recovery(greet) + + +def test_the_vm_takes_it(): + headers = [("content-type", "application/vnd.restate.invocation.v6")] + VMWrapper(headers, explicit_recovery=True) + VMWrapper(headers) + + +@pytest.mark.anyio +async def test_a_retry_runs_the_handler_again_and_entries_are_named(): + calls = 0 + counter = explicit_recovery(VirtualObject("counter")) + + def step() -> int: + nonlocal calls + calls += 1 + return calls + + @counter.handler(journal_retention=timedelta(hours=1)) + async def add(ctx: ObjectContext) -> str: + attempt = await ctx.run_typed("step", step) # replayed without explicit recovery + ctx.set("attempt", attempt, entry_name="store attempt") + if attempt == 1: + raise ValueError("retry me") + return ctx.request().id + + async with restate.create_test_harness( + restate.app([counter]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + invocation_id = await harness.client.object_call(add, key="a", arg=None) + assert calls == 2 + + async with httpx.AsyncClient() as http: + response = await http.post( + f"{harness.admin_api_url}/query", + json={"query": f"SELECT entry_type, name FROM sys_journal WHERE id = '{invocation_id}'"}, + headers={"accept": "application/json"}, + ) + response.raise_for_status() + names = [row["name"] for row in response.json()["rows"] if row["entry_type"] == "Command: SetState"] + assert names == ["store attempt", "store attempt"]