diff --git a/CHANGELOG.md b/CHANGELOG.md index 6230f2754..e15538fe5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ to include examples, links to docs, or any other relevant information. - Added the experimental `Worker` `patch_activation_callback` option, allowing workers to decide whether a first non-replay `workflow.patched` call should activate a patch during rolling deployments. +- Added external storage support to Nexus task handling. ### Changed diff --git a/scripts/gen_payload_visitor.py b/scripts/gen_payload_visitor.py index c2bc15837..a5940d992 100644 --- a/scripts/gen_payload_visitor.py +++ b/scripts/gen_payload_visitor.py @@ -12,6 +12,7 @@ sys.path.insert(0, str(base_dir)) from temporalio.api.common.v1.message_pb2 import Payload, Payloads, SearchAttributes +from temporalio.bridge.proto.nexus import NexusTaskCompletion from temporalio.bridge.proto.workflow_activation.workflow_activation_pb2 import ( WorkflowActivation, ) @@ -426,10 +427,11 @@ def write_bridge_visitors() -> None: out_path = base_dir / "temporalio" / "bridge" / "_visitor.py" # Build root descriptors: WorkflowActivation, WorkflowActivationCompletion, - # and all messages from selected API modules + # NexusTaskCompletion, and all messages from selected API modules roots: list[Descriptor] = [ WorkflowActivation.DESCRIPTOR, WorkflowActivationCompletion.DESCRIPTOR, + NexusTaskCompletion.DESCRIPTOR, ] code = VisitorGenerator().generate(roots) diff --git a/temporalio/bridge/_visitor.py b/temporalio/bridge/_visitor.py index a9956e12b..6f3af2aea 100644 --- a/temporalio/bridge/_visitor.py +++ b/temporalio/bridge/_visitor.py @@ -549,3 +549,55 @@ async def _visit_coresdk_workflow_completion_WorkflowActivationCompletion( await self._visit_coresdk_workflow_completion_Success(fs, o.successful) elif o.HasField("failed"): await self._visit_coresdk_workflow_completion_Failure(fs, o.failed) + + async def _visit_temporal_api_nexus_v1_StartOperationResponse_Sync( + self, fs: VisitorFunctions, o: Any + ): + if o.HasField("payload"): + await self._visit_temporal_api_common_v1_Payload(fs, o.payload) + + async def _visit_temporal_api_nexus_v1_Failure(self, fs: VisitorFunctions, o: Any): + if o.HasField("cause"): + await self._visit_temporal_api_nexus_v1_Failure(fs, o.cause) + + async def _visit_temporal_api_nexus_v1_UnsuccessfulOperationError( + self, fs: VisitorFunctions, o: Any + ): + if o.HasField("failure"): + await self._visit_temporal_api_nexus_v1_Failure(fs, o.failure) + + async def _visit_temporal_api_nexus_v1_StartOperationResponse( + self, fs: VisitorFunctions, o: Any + ): + if o.HasField("sync_success"): + await self._visit_temporal_api_nexus_v1_StartOperationResponse_Sync( + fs, o.sync_success + ) + elif o.HasField("operation_error"): + await self._visit_temporal_api_nexus_v1_UnsuccessfulOperationError( + fs, o.operation_error + ) + elif o.HasField("failure"): + await self._visit_temporal_api_failure_v1_Failure(fs, o.failure) + + async def _visit_temporal_api_nexus_v1_Response(self, fs: VisitorFunctions, o: Any): + if o.HasField("start_operation"): + await self._visit_temporal_api_nexus_v1_StartOperationResponse( + fs, o.start_operation + ) + + async def _visit_temporal_api_nexus_v1_HandlerError( + self, fs: VisitorFunctions, o: Any + ): + if o.HasField("failure"): + await self._visit_temporal_api_nexus_v1_Failure(fs, o.failure) + + async def _visit_coresdk_nexus_NexusTaskCompletion( + self, fs: VisitorFunctions, o: Any + ): + if o.HasField("completed"): + await self._visit_temporal_api_nexus_v1_Response(fs, o.completed) + elif o.HasField("error"): + await self._visit_temporal_api_nexus_v1_HandlerError(fs, o.error) + elif o.HasField("failure"): + await self._visit_temporal_api_failure_v1_Failure(fs, o.failure) diff --git a/temporalio/converter/_data_converter.py b/temporalio/converter/_data_converter.py index 6425d6b61..8604ea196 100644 --- a/temporalio/converter/_data_converter.py +++ b/temporalio/converter/_data_converter.py @@ -13,9 +13,9 @@ import temporalio.api.common.v1 import temporalio.api.failure.v1 import temporalio.common -from temporalio.api.sdk.v1.external_storage_pb2 import ExternalStorageReference from temporalio.converter._extstore import ( _REFERENCE_ENCODING, + _REFERENCE_MESSAGE_TYPE, ExternalStorage, StorageDriverStoreContext, ) @@ -35,8 +35,6 @@ WithSerializationContext, ) -_REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode() - def _is_reference_payload(p: temporalio.api.common.v1.Payload) -> bool: """Return True if *p* is an external-storage reference payload.""" diff --git a/temporalio/converter/_extstore.py b/temporalio/converter/_extstore.py index c31424acf..a946b2c0f 100644 --- a/temporalio/converter/_extstore.py +++ b/temporalio/converter/_extstore.py @@ -27,6 +27,7 @@ _T = TypeVar("_T") _REFERENCE_ENCODING = b"json/external-storage-reference" +_REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode() @dataclass @@ -455,8 +456,6 @@ async def _store_payload_sequence( def _decode_reference(self, payload: Payload) -> ExternalStorageReference | None: """Decode an external storage reference from a payload.""" - if len(payload.external_payloads) == 0: - return None encoding = payload.metadata.get("encoding", b"") if encoding == _REFERENCE_ENCODING: legacy = self._legacy_claim_converter.from_payload( @@ -468,6 +467,11 @@ def _decode_reference(self, payload: Payload) -> ExternalStorageReference | None driver_name=legacy.driver_name, claim_data=legacy.driver_claim.claim_data, ) + if not ( + encoding == b"json/protobuf" + and payload.metadata.get("messageType") == _REFERENCE_MESSAGE_TYPE + ): + return None ref = self._claim_converter.from_payload(payload, ExternalStorageReference) return ref if isinstance(ref, ExternalStorageReference) else None diff --git a/temporalio/worker/_nexus.py b/temporalio/worker/_nexus.py index 08ecd2f81..131e50862 100644 --- a/temporalio/worker/_nexus.py +++ b/temporalio/worker/_nexus.py @@ -6,7 +6,7 @@ import concurrent.futures import contextvars import threading -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from functools import reduce @@ -30,6 +30,8 @@ import temporalio.common import temporalio.converter import temporalio.nexus +from temporalio.bridge._visitor import PayloadVisitor +from temporalio.bridge._visitor_functions import PayloadSequence, VisitorFunctions from temporalio.bridge.worker import PollShutdownError from temporalio.exceptions import ( ApplicationError, @@ -216,6 +218,19 @@ async def _complete_task( ): await asyncio.shield(self._bridge_worker().complete_nexus_task(completion)) + async def _encode_completion( + self, completion: temporalio.bridge.proto.nexus.NexusTaskCompletion + ) -> None: + """Apply the payload codec then external storage to the completion's payloads.""" + dc = self._data_converter + await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit( + _PayloadTransformVisitor(dc._encode_payload_sequence), completion + ) + await PayloadVisitor(skip_search_attributes=True).visit( + _PayloadTransformVisitor(dc._external_store_payload_sequence), + completion, + ) + # TODO(nexus-preview): stack trace pruning. See sdk-typescript NexusHandler.execute # "Any call up to this function and including this one will be trimmed out of stack traces."" @@ -260,6 +275,14 @@ async def _handle_cancel_operation_task( try: try: await self._handler.cancel_operation(ctx, request.operation_token) + completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( + task_token=task_token, + completed=temporalio.api.nexus.v1.Response( + cancel_operation=temporalio.api.nexus.v1.CancelOperationResponse() + ), + ) + # No-op but keeps the cancel covered if it ever carries a payload. + await self._encode_completion(completion) except asyncio.CancelledError: completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, @@ -271,16 +294,12 @@ async def _handle_cancel_operation_task( completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, ) - await self._data_converter.encode_failure( - handler_error, completion.failure - ) - else: - completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( - task_token=task_token, - completed=temporalio.api.nexus.v1.Response( - cancel_operation=temporalio.api.nexus.v1.CancelOperationResponse() - ), + self._data_converter.failure_converter.to_failure( + handler_error, + self._data_converter.payload_converter, + completion.failure, ) + await self._encode_completion(completion) await self._complete_task(completion) except Exception: logger.exception("Failed to send Nexus task completion") @@ -315,6 +334,13 @@ async def _handle_start_operation_task( request_deadline, endpoint, ) + completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( + task_token=task_token, + completed=temporalio.api.nexus.v1.Response( + start_operation=start_response + ), + ) + await self._encode_completion(completion) except asyncio.CancelledError: completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( task_token=task_token, @@ -326,19 +352,15 @@ async def _handle_start_operation_task( task_token=task_token, ) handler_error = _exception_to_handler_error(err) - await self._data_converter.encode_failure( - handler_error, completion.failure + self._data_converter.failure_converter.to_failure( + handler_error, + self._data_converter.payload_converter, + completion.failure, ) if isinstance(err, concurrent.futures.BrokenExecutor): self._fail_worker_exception_queue.put_nowait(err) - else: - completion = temporalio.bridge.proto.nexus.NexusTaskCompletion( - task_token=task_token, - completed=temporalio.api.nexus.v1.Response( - start_operation=start_response - ), - ) + await self._encode_completion(completion) await self._complete_task(completion) except Exception: @@ -417,7 +439,9 @@ async def _start_operation( ) ) elif isinstance(result, nexusrpc.handler.StartOperationResultSync): - [payload] = await self._data_converter.encode([result.value]) + [payload] = self._data_converter.payload_converter.to_payloads( + [result.value] + ) return temporalio.api.nexus.v1.StartOperationResponse( sync_success=temporalio.api.nexus.v1.StartOperationResponse.Sync( payload=payload, @@ -446,10 +470,41 @@ async def _start_operation( ) from err.__cause__ except FailureError as new_err: response = temporalio.api.nexus.v1.StartOperationResponse() - await self._data_converter.encode_failure(new_err, response.failure) + self._data_converter.failure_converter.to_failure( + new_err, + self._data_converter.payload_converter, + response.failure, + ) return response +class _PayloadTransformVisitor(VisitorFunctions): + """Adapts a payload-sequence transform for use with :class:`PayloadVisitor`.""" + + def __init__( + self, + f: Callable[ + [Sequence[temporalio.api.common.v1.Payload]], + Awaitable[list[temporalio.api.common.v1.Payload]], + ], + ) -> None: + self._f = f + + async def visit_payload(self, payload: temporalio.api.common.v1.Payload) -> None: + new_payload = (await self._f([payload]))[0] + if new_payload is not payload: + payload.CopyFrom(new_payload) + + async def visit_payloads(self, payloads: PayloadSequence) -> None: + if len(payloads) == 0: + return + new_payloads = await self._f(payloads) + if new_payloads is payloads: + return + del payloads[:] + payloads.extend(new_payloads) + + @dataclass class _DummyPayloadSerializer: data_converter: temporalio.converter.DataConverter @@ -465,18 +520,35 @@ async def deserialize( content: nexusrpc.Content, # type:ignore[reportUnusedParameter] as_type: type[Any] | None = None, ) -> Any: - payload = self.payload - if self.data_converter.payload_codec: - try: - [payload] = await self.data_converter.payload_codec.decode([payload]) - except Exception as err: - raise nexusrpc.HandlerError( - "Payload codec failed to decode Nexus operation input", - type=nexusrpc.HandlerErrorType.INTERNAL, - ) from err + dc = self.data_converter + # The visitor mutates in place, so work on a copy to leave the request + # payload untouched. + payload = temporalio.api.common.v1.Payload() + payload.CopyFrom(self.payload) + try: + await PayloadVisitor(skip_search_attributes=True).visit( + _PayloadTransformVisitor(dc._external_retrieve_payload_sequence), + payload, + ) + except Exception as err: + raise nexusrpc.HandlerError( + "Failed to retrieve Nexus operation input from external storage", + type=nexusrpc.HandlerErrorType.INTERNAL, + retryable_override=True, + ) from err + + try: + await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit( + _PayloadTransformVisitor(dc._decode_payload_sequence), payload + ) + except Exception as err: + raise nexusrpc.HandlerError( + "Payload codec failed to decode Nexus operation input", + type=nexusrpc.HandlerErrorType.INTERNAL, + ) from err try: - [input] = self.data_converter.payload_converter.from_payloads( + [input] = dc.payload_converter.from_payloads( [payload], type_hints=[as_type] if as_type else None, ) diff --git a/tests/nexus/test_temporal_extstore.py b/tests/nexus/test_temporal_extstore.py new file mode 100644 index 000000000..863d44627 --- /dev/null +++ b/tests/nexus/test_temporal_extstore.py @@ -0,0 +1,268 @@ +"""Integration tests for external storage with Nexus task processing. + +Mirrors sdk-typescript's test-integration-extstore-nexus. A caller workflow +invokes a Nexus operation whose input and/or result is large enough to be +offloaded to external storage. Verifies that the offloaded input is retrieved +before the handler runs, that a large synchronous result is offloaded when +completing the task (and retrieved by the caller), and that a transient +storage-driver failure fails the Nexus task retryably and then recovers. +""" + +from __future__ import annotations + +import dataclasses +import uuid +from collections.abc import Sequence +from datetime import timedelta + +import nexusrpc +import pytest +from nexusrpc.handler import ( + StartOperationContext, + service_handler, + sync_operation, +) + +import temporalio.converter +from temporalio import workflow +from temporalio.api.common.v1 import Payload +from temporalio.client import Client, WorkflowFailureError +from temporalio.converter import ( + ExternalStorage, + StorageDriverClaim, + StorageDriverRetrieveContext, + StorageDriverStoreContext, +) +from temporalio.exceptions import ApplicationError, NexusOperationError +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import UnsandboxedWorkflowRunner, Worker +from tests.helpers.nexus import make_nexus_endpoint_name +from tests.test_extstore import InMemoryTestDriver + +PAYLOAD_SIZE = 4096 +PAYLOAD_SIZE_THRESHOLD = 1024 +_STORE_FAILURE_MESSAGE = "external storage store failed" + + +@nexusrpc.service +class ExtStoreNexusService: + size_op: nexusrpc.Operation[str, int] + big_result_op: nexusrpc.Operation[int, str] + + +@service_handler(service=ExtStoreNexusService) +class ExtStoreNexusServiceHandler: + @sync_operation + async def size_op(self, _ctx: StartOperationContext, data: str) -> int: + return len(data) + + @sync_operation + async def big_result_op(self, _ctx: StartOperationContext, size: int) -> str: + return "x" * size + + +@workflow.defn +class SizeOpCallerWorkflow: + """Calls ``size_op`` with a large input (offloaded) and returns its length.""" + + @workflow.run + async def run(self, task_queue: str) -> int: + nexus_client = workflow.create_nexus_client( + service=ExtStoreNexusService, + endpoint=make_nexus_endpoint_name(task_queue), + ) + return await nexus_client.execute_operation( + ExtStoreNexusService.size_op, "x" * PAYLOAD_SIZE + ) + + +@workflow.defn +class BigResultOpCallerWorkflow: + """Calls ``big_result_op`` (whose result is offloaded) and returns its length.""" + + @workflow.run + async def run(self, task_queue: str) -> int: + nexus_client = workflow.create_nexus_client( + service=ExtStoreNexusService, + endpoint=make_nexus_endpoint_name(task_queue), + ) + result = await nexus_client.execute_operation( + ExtStoreNexusService.big_result_op, PAYLOAD_SIZE + ) + return len(result) + + +class TransientFailureDriver(InMemoryTestDriver): + """In-memory driver that fails the first store and/or retrieve call, then + behaves normally. Simulates a transient storage-driver outage.""" + + def __init__( + self, + *, + fail_first_store: bool = False, + fail_first_retrieve: bool = False, + driver_name: str = "test-driver", + ): + super().__init__(driver_name=driver_name) + self._fail_first_store = fail_first_store + self._fail_first_retrieve = fail_first_retrieve + self.store_attempts = 0 + self.retrieve_attempts = 0 + + async def store( + self, + context: StorageDriverStoreContext, + payloads: Sequence[Payload], + ) -> list[StorageDriverClaim]: + self.store_attempts += 1 + if self._fail_first_store and self.store_attempts == 1: + raise RuntimeError("transient store failure") + return await super().store(context, payloads) + + async def retrieve( + self, + context: StorageDriverRetrieveContext, + claims: Sequence[StorageDriverClaim], + ) -> list[Payload]: + self.retrieve_attempts += 1 + if self._fail_first_retrieve and self.retrieve_attempts == 1: + raise RuntimeError("transient retrieve failure") + return await super().retrieve(context, claims) + + +def _client_with_extstore( + env: WorkflowEnvironment, driver: InMemoryTestDriver +) -> Client: + config = env.client.config() + config["data_converter"] = dataclasses.replace( + temporalio.converter.default(), + external_storage=ExternalStorage( + drivers=[driver], + payload_size_threshold=PAYLOAD_SIZE_THRESHOLD, + ), + ) + return Client(**config) + + +async def _run_caller( + env: WorkflowEnvironment, + driver: InMemoryTestDriver, + caller: type, +) -> int: + client = _client_with_extstore(env, driver) + task_queue = str(uuid.uuid4()) + async with Worker( + client, + task_queue=task_queue, + workflows=[SizeOpCallerWorkflow, BigResultOpCallerWorkflow], + nexus_service_handlers=[ExtStoreNexusServiceHandler()], + workflow_runner=UnsandboxedWorkflowRunner(), + ): + await env.create_nexus_endpoint( + make_nexus_endpoint_name(task_queue), task_queue + ) + return await client.execute_workflow( + caller.run, + task_queue, + id=str(uuid.uuid4()), + task_queue=task_queue, + execution_timeout=timedelta(seconds=30), + ) + + +def _cause_chain(err: BaseException) -> list[BaseException]: + chain: list[BaseException] = [] + e: BaseException | None = err + while e is not None: + chain.append(e) + e = e.__cause__ + return chain + + +async def test_nexus_operation_input_offloaded_and_retrieved(env: WorkflowEnvironment): + """The offloaded operation input is retrieved before the handler runs.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + driver = InMemoryTestDriver() + result = await _run_caller(env, driver, SizeOpCallerWorkflow) + + assert result == PAYLOAD_SIZE + assert driver._store_calls >= 1 + assert driver._retrieve_calls >= 1 + + +async def test_nexus_operation_sync_result_offloaded_and_retrieved( + env: WorkflowEnvironment, +): + """A large synchronous result is offloaded and retrieved by the caller.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + driver = InMemoryTestDriver() + result = await _run_caller(env, driver, BigResultOpCallerWorkflow) + + assert result == PAYLOAD_SIZE + assert driver._store_calls >= 1 + assert driver._retrieve_calls >= 1 + + +async def test_nexus_operation_transient_retrieve_failure_recovers( + env: WorkflowEnvironment, +): + """A transient retrieve failure fails the task retryably; it then recovers.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + driver = TransientFailureDriver(fail_first_retrieve=True) + result = await _run_caller(env, driver, SizeOpCallerWorkflow) + + assert result == PAYLOAD_SIZE + assert driver.retrieve_attempts >= 2 + + +async def test_nexus_operation_transient_store_failure_recovers( + env: WorkflowEnvironment, +): + """A transient store failure fails the task retryably; it then recovers.""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + driver = TransientFailureDriver(fail_first_store=True) + result = await _run_caller(env, driver, BigResultOpCallerWorkflow) + + assert result == PAYLOAD_SIZE + assert driver.store_attempts >= 2 + + +class PermanentFailStoreDriver(InMemoryTestDriver): + """Store always fails non-retryably, so the Nexus operation fails permanently.""" + + async def store( + self, + context: StorageDriverStoreContext, + payloads: Sequence[Payload], + ) -> list[StorageDriverClaim]: + raise ApplicationError(_STORE_FAILURE_MESSAGE, non_retryable=True) + + +async def test_nexus_operation_store_failure_fails_operation( + env: WorkflowEnvironment, +): + """A non-retryable store failure fails the operation and surfaces the driver + error to the caller (deterministically, with no retries).""" + if env.supports_time_skipping: + pytest.skip("Nexus tests don't work with the Java test server") + + driver = PermanentFailStoreDriver() + with pytest.raises(WorkflowFailureError) as exc_info: + await _run_caller(env, driver, BigResultOpCallerWorkflow) + + causes = _cause_chain(exc_info.value) + assert [type(c) for c in causes] == [ + WorkflowFailureError, + NexusOperationError, + nexusrpc.HandlerError, + ApplicationError, + ] + assert _STORE_FAILURE_MESSAGE in str(causes[-1])