diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index c0ada9a6..8ca5ef90 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -299,6 +299,12 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), is_first_invocation=not has_prior_operations, execution_input=input_event, + # Read the map through a callable rather than snapshotting it + # here: the invocation-end hook needs the state as of the end of + # the invocation, and neither hook pays for the conversion until + # a plugin actually reads it. + operations_provider=lambda: execution_state.operations, + updated_operation_ids=invocation_input.updated_operation_ids, ) # Thread 1: Run background checkpoint processing executor.submit(execution_state.checkpoint_batches_forever) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index 7ccb8d50..ffc5a65f 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -98,7 +98,7 @@ def from_operation( start_time=operation.start_timestamp, end_time=operation.end_timestamp, result=_extract_result(operation), - error=_extract_error(operation), + error=_copy_error(_extract_error(operation)), attempt=( operation.step_details.attempt if operation.step_details else None ), @@ -107,6 +107,42 @@ def from_operation( ) +def _copy_error(error: ErrorObject | None) -> ErrorObject | None: + """Return a plugin-owned copy of an operation error. + + The checkpointed ``ErrorObject`` is handed straight to user code on replay, + and its ``stack_trace`` is a mutable list. Without a copy a plugin reading + ``info.operations`` could append to (or clear) that list and change the error + the execution later raises. Only the list needs cloning -- the other fields + are immutable strings -- so this is cheaper than a full deep copy. + """ + if error is None: + return None + return ErrorObject( + message=error.message, + type=error.type, + data=error.data, + stack_trace=( + list(error.stack_trace) if error.stack_trace is not None else None + ), + ) + + +def _to_operation_info_map( + operations: Mapping[str, Operation], +) -> dict[str, OperationInfo]: + """Convert a map of checkpointed operations to the plugin ``OperationInfo`` view. + + ``is_replayed`` is left at its default ``False``: these entries describe the + stored state of an operation, not a replay event for it. Replay is signalled + through the dedicated operation hooks. + """ + return { + operation_id: OperationInfo.from_operation(operation) + for operation_id, operation in operations.items() + } + + @dataclass(frozen=True) class OperationStartInfo(OperationInfo): pass @@ -204,11 +240,54 @@ class InvocationInfo: without it); ``durable_execution()`` always populates it with the deserialized input payload, which is ``{}`` when the payload is empty. """ + operations: dict[str, OperationInfo] = field( + default_factory=dict, + kw_only=True, + repr=False, + compare=False, + hash=False, + metadata={"experimental": True}, + ) + """EXPERIMENTAL: Checkpointed operations for this execution, keyed by id. + + A point-in-time view of the execution's operation map: as observed at the + start of the invocation on ``on_invocation_start``, and as observed at the + end of the invocation on ``on_invocation_end``. + + Not a reliable signal of whether this is the first invocation: the initial + execution state already carries the ``EXECUTION`` operation, so even a first + invocation-start sees a non-empty map. Use + :attr:`is_first_invocation` for that. What a first invocation lacks is prior + non-execution operations. + + Excluded from ``repr``, ``__eq__`` and ``__hash__`` for the same reasons as + :attr:`execution_input`: the entries carry operation results and errors that + instrumentation would otherwise log wholesale, and a mapping-valued field + would make a previously hashable info unhashable. + """ @dataclass(frozen=True) class InvocationStartInfo(InvocationInfo): - pass + updated_operations: dict[str, OperationInfo] = field( + default_factory=dict, + kw_only=True, + repr=False, + compare=False, + hash=False, + metadata={"experimental": True}, + ) + """EXPERIMENTAL: Operations updated externally while this execution was suspended. + + A wait timer that expired, a callback that was delivered, or a chained + invoke that completed between the previous invocation and this one. This is + the subset of :attr:`InvocationInfo.operations` named by the durable + invocation input's ``UpdatedOperationIds``, so it is empty on the first + invocation. + + Excluded from ``repr``, ``__eq__`` and ``__hash__`` like + :attr:`InvocationInfo.operations`. + """ @dataclass(frozen=True) @@ -244,6 +323,7 @@ def from_durable_execution_invocation_output( cls, invocation_start_info: InvocationStartInfo, output: "DurableExecutionInvocationOutput", + operations: dict[str, OperationInfo] | None = None, ): return InvocationEndInfo( request_id=invocation_start_info.request_id, @@ -251,6 +331,13 @@ def from_durable_execution_invocation_output( is_first_invocation=invocation_start_info.is_first_invocation, execution_start_time=invocation_start_info.execution_start_time, execution_input=invocation_start_info.execution_input, + # Default to the start-of-invocation view when the caller has no + # fresher snapshot to offer. + operations=( + operations + if operations is not None + else invocation_start_info.operations + ), status=output.status, error=output.error, execution_result=output.result, @@ -341,6 +428,7 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None + self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None @contextlib.contextmanager def run(self): @@ -353,6 +441,7 @@ def run(self): yield finally: self._invocation_status = None + self._operations_provider = None # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @@ -393,6 +482,35 @@ def execute_plugins(self, info, sync): # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) + def _snapshot_operation_infos( + self, + operations_provider: Callable[[], Mapping[str, Operation]] | None, + ) -> dict[str, OperationInfo]: + """Build the plugin ``OperationInfo`` view of the current operation map. + + Returns a plain ``dict``, matching :class:`OperationChangeInfo`. That + matters beyond consistency: ``dataclasses.asdict()`` and ``pickle`` only + traverse real dicts, so a custom ``Mapping`` here would leave the + enclosing hook info unserializable for the very plugins these fields + exist to serve. + + Built eagerly, which also pins the point in time the hook reports: a + plugin that stashes the info and reads it later still sees the state as + of its own hook. + + Skipped entirely when no plugins are registered -- ``durable_execution()`` + passes a provider unconditionally, so without this gate a plugin-free + execution would pay for a view nothing can read. + """ + if not self._plugins or operations_provider is None: + return {} + try: + return _to_operation_info_map(operations_provider()) + except Exception: + # A plugin-facing view must never break the execution. + logger.exception("Failed to snapshot operations for plugin hook") + return {} + def on_invocation_start( self, execution_arn: str, @@ -400,14 +518,37 @@ def on_invocation_start( execution_start_time: datetime.datetime | None, lambda_context: LambdaContext | None, execution_input: Any = None, + operations_provider: Callable[[], Mapping[str, Operation]] | None = None, + updated_operation_ids: Sequence[str] | None = None, ) -> None: + """Fire the invocation-start hook. + + Args: + execution_arn: ARN of the durable execution. + is_first_invocation: False when prior operations exist (a replay). + execution_start_time: Start timestamp of the execution operation. + lambda_context: Lambda context, for the request id. + execution_input: The deserialized execution input event. + operations_provider: Returns the current checkpointed operation map, + converted here into the plugin's ``OperationInfo`` view. + updated_operation_ids: Operation ids from the invocation input's + ``UpdatedOperationIds`` -- those updated while suspended. + """ aws_request_id = lambda_context.aws_request_id if lambda_context else None + self._operations_provider = operations_provider if self._plugins else None + operations = self._snapshot_operation_infos(operations_provider) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, request_id=aws_request_id, is_first_invocation=is_first_invocation, execution_start_time=execution_start_time, execution_input=self._snapshot_execution_input(execution_input), + operations=operations, + updated_operations={ + operation_id: operations[operation_id] + for operation_id in (updated_operation_ids or []) + if operation_id in operations + }, ) self.execute_plugins(self._invocation_status, sync=True) @@ -448,9 +589,13 @@ def on_invocation_end( # on_invocation_start not called, skip return + # Re-read the operation map so the end hook sees the state as of the end + # of this invocation, not the snapshot taken at its start. invocation_end_info = ( InvocationEndInfo.from_durable_execution_invocation_output( - self._invocation_status, output + self._invocation_status, + output, + operations=self._snapshot_operation_infos(self._operations_provider), ) ) self.execute_plugins(invocation_end_info, sync=True) @@ -629,10 +774,7 @@ def on_operation_update( operation.operation_id: OperationInfo.from_operation(operation) for operation in changed_operations }, - operations={ - operation_id: OperationInfo.from_operation(operation) - for operation_id, operation in operations.items() - }, + operations=_to_operation_info_map(operations), ), sync=True, ) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py new file mode 100644 index 00000000..26f1b77a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -0,0 +1,258 @@ +"""Integration tests for the plugin invocation operation maps. + +Exercises `InvocationInfo.operations` and +`InvocationStartInfo.updated_operations` through complete +`durable_execution()` invocations -- across the decorator, invocation-input +parsing, the checkpoint path and the invocation hooks -- including a +suspend/replay pair where the replay reports the externally-completed wait via +`UpdatedOperationIds`. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import Mock, patch + +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from tests.test_helpers import operation_id_sequence + + +class _MapRecordingPlugin(DurableInstrumentationPlugin): + """Records the operation maps seen on each invocation hook.""" + + def __init__(self) -> None: + self.starts: list[tuple[list[str], list[str]]] = [] + self.ends: list[tuple[str, list[str]]] = [] + + def on_invocation_start(self, info) -> None: + self.starts.append((sorted(info.operations), sorted(info.updated_operations))) + + def on_invocation_end(self, info) -> None: + self.ends.append((info.status.value, sorted(info.operations))) + + +def _lambda_context() -> Mock: + ctx = Mock() + ctx.aws_request_id = "test-request-id" + ctx.client_context = None + ctx.identity = None + ctx._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + ctx.invoked_function_arn = "test-arn" + ctx.tenant_id = None + return ctx + + +def _event( + extra_operations: list[dict] | None = None, + updated_operation_ids: list[str] | None = None, +) -> dict: + execution_operation = { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": '{"name": "World"}'}, + } + event: dict[str, Any] = { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [execution_operation, *(extra_operations or [])], + "NextMarker": "", + }, + "LocalRunner": True, + } + if updated_operation_ids is not None: + event["UpdatedOperationIds"] = updated_operation_ids + return event + + +def _tracking_checkpoint(initial_operations: list[Operation] | None = None): + """Checkpoint mock that accumulates operations, as the service would.""" + operations: list[Operation] = list(initial_operations or []) + + def mock_checkpoint( + durable_execution_arn, # noqa: ARG001 + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + for update in updates: + operations.append( + Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + ) + ) + return CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=operations.copy() + ), + ) + + return mock_checkpoint + + +def test_operation_maps_on_a_completing_invocation(): + """The start map holds the prior state; the end map sees the step added.""" + plugin = _MapRecordingPlugin() + + @durable_execution(plugins=[plugin]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + return context.step(lambda _ctx: "stepped", name="greet") + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + result = my_handler(_event(), _lambda_context()) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + + # First invocation: only the EXECUTION operation exists at start, and + # nothing was updated externally. + (start_operations, start_updated) = plugin.starts[0] + assert start_operations == ["execution-1"] + assert start_updated == [] + + # The end hook re-reads the map, so it sees the step checkpointed during + # this invocation -- that is the point of re-reading rather than reusing + # the start snapshot. + status, end_operations = plugin.ends[0] + assert status == InvocationStatus.SUCCEEDED.value + assert len(end_operations) > len(start_operations) + assert "execution-1" in end_operations + + +def test_operation_maps_across_suspend_and_replay(): + """The replay start hook reports the externally-completed wait. + + Invocation 1 suspends on a wait. Invocation 2 replays with the wait already + SUCCEEDED and its id in ``UpdatedOperationIds``, which is exactly what + ``updated_operations`` is derived from. + """ + wait_id = next(operation_id_sequence()) + + # --- Invocation 1: the wait starts and the execution suspends. + first = _MapRecordingPlugin() + + @durable_execution(plugins=[first]) + def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + context.wait(Duration.from_seconds(60)) + return "done" + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + first_result = suspending_handler(_event(), _lambda_context()) + + assert first_result["Status"] == InvocationStatus.PENDING.value + start_operations, start_updated = first.starts[0] + assert start_operations == ["execution-1"] + assert start_updated == [] + # The suspending end hook already sees the wait that was just checkpointed. + _, end_operations = first.ends[0] + assert wait_id in end_operations + + # --- Invocation 2: replay with the wait completed externally. + replay = _MapRecordingPlugin() + + @durable_execution(plugins=[replay]) + def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + context.wait(Duration.from_seconds(60)) + return "done" + + completed_wait = { + "Id": wait_id, + "Type": OperationType.WAIT.value, + "Status": OperationStatus.SUCCEEDED.value, + } + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + replay_result = replayed_handler( + _event(extra_operations=[completed_wait], updated_operation_ids=[wait_id]), + _lambda_context(), + ) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + + start_operations, start_updated = replay.starts[0] + # The replay start map carries the prior state, including the wait. + assert sorted(["execution-1", wait_id]) == start_operations + # And updated_operations is the UpdatedOperationIds subset of it. + assert start_updated == [wait_id] + + status, end_operations = replay.ends[0] + assert status == InvocationStatus.SUCCEEDED.value + assert wait_id in end_operations + + +def test_updated_operations_ignores_ids_absent_from_the_map(): + """An id the execution state does not carry must not appear in the subset.""" + plugin = _MapRecordingPlugin() + + @durable_execution(plugins=[plugin]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + return "ok" + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + my_handler( + _event(updated_operation_ids=["execution-1", "never-checkpointed"]), + _lambda_context(), + ) + + _, start_updated = plugin.starts[0] + assert start_updated == ["execution-1"] + + +def test_plugin_free_execution_still_completes(): + """The provider gate must not disturb an execution without plugins.""" + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + return "ok" + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = _tracking_checkpoint() + mock_client_class.initialize_client.return_value = mock_client + + result = my_handler(_event(), _lambda_context()) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 89de9c68..c33c370a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,8 +1,10 @@ import datetime import logging +import pickle import unittest -from dataclasses import fields -from unittest.mock import MagicMock +from copy import deepcopy +from dataclasses import asdict, fields +from unittest.mock import MagicMock, patch from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -229,6 +231,42 @@ def test_payload_fields_are_excluded_from_equality(self): ), ) + def test_operation_map_fields_are_additive(self): + """The operation maps must be as additive as the payload fields. + + They are mapping-valued, so including them in the generated methods + would make a previously hashable info unhashable, and their entries + carry operation results and errors that instrumentation logs wholesale. + """ + start_fields = {f.name: f for f in fields(InvocationStartInfo)} + end_fields = {f.name: f for f in fields(InvocationEndInfo)} + + for holder, name in ( + (start_fields, "operations"), + (start_fields, "updated_operations"), + (end_fields, "operations"), + ): + self.assertFalse(holder[name].repr, name) + self.assertFalse(holder[name].compare, name) + self.assertIs(holder[name].hash, False, name) + + base = { + "request_id": "req-1", + "execution_arn": "arn:test", + "is_first_invocation": True, + } + with_maps = InvocationStartInfo( + **base, + operations={"op-1": OPERATION_END_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + + # Hashable despite the mapping fields, and equal to the map-free info. + self.assertEqual(hash(InvocationStartInfo(**base)), hash(with_maps)) + self.assertEqual(InvocationStartInfo(**base), with_maps) + # Operation results do not leak through the info's repr. + self.assertNotIn("op-1", repr(with_maps)) + def test_payload_fields_are_declared_non_compare(self): """Pin the declarations, not just the observed behaviour.""" start_fields = {f.name: f for f in fields(InvocationStartInfo)} @@ -330,6 +368,69 @@ def test_invocation_end_info_from_invocation_output_carries_input_and_result(sel self.assertEqual(end_info.status, InvocationStatus.SUCCEEDED) self.assertIsNone(end_info.error) + def test_invocation_info_operation_maps_default_to_empty(self): + """Both maps default to empty so existing constructor calls keep working.""" + self.assertEqual({}, INVOCATION_START_INFO.operations) + self.assertEqual({}, INVOCATION_START_INFO.updated_operations) + self.assertEqual({}, INVOCATION_END_INFO.operations) + + def test_invocation_start_info_carries_operation_maps(self): + info = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=False, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + info.operations, + ) + self.assertEqual({"op-1": OPERATION_END_INFO}, info.updated_operations) + + def test_invocation_end_info_inherits_start_operations(self): + """The end factory reuses the start snapshot when given no override.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_END_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output(start, output) + + self.assertEqual({"op-1": OPERATION_END_INFO}, end.operations) + # updated_operations is a start-hook-only surface. + self.assertFalse(hasattr(end, "updated_operations")) + + def test_invocation_end_info_accepts_fresher_operations(self): + """An explicit map wins, letting the end hook see end-of-invocation state.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_START_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output( + start, + output, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + end.operations, + ) + def test_user_function_start_info(self): self.assertEqual(USER_FUNCTION_START_INFO.operation_id, "op-1") self.assertEqual(USER_FUNCTION_START_INFO.operation_type, OperationType.STEP) @@ -706,6 +807,438 @@ def test_pending_fires_invocation_end(self): self.assertIn("invocation_end:req-1", self.plugin.calls) +class TestInvocationHookOperationMaps(unittest.TestCase): + """Tests for the operations / updated_operations maps on invocation hooks.""" + + def setUp(self): + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + @staticmethod + def _operation(operation_id, status=OperationStatus.SUCCEEDED): + return Operation( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=f"name-{operation_id}", + parent_id="root", + status=status, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + + def _start(self, *, operations=None, updated_operation_ids=None): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=(lambda: operations) + if operations is not None + else None, + updated_operation_ids=updated_operation_ids, + ) + + def test_start_info_converts_operations_to_operation_info(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start(operations=operations) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.operations)) + converted = start_info.operations["op-1"] + self.assertIsInstance(converted, OperationInfo) + self.assertEqual("op-1", converted.operation_id) + self.assertEqual(OperationStatus.SUCCEEDED, converted.status) + self.assertEqual(START_TS, converted.start_time) + self.assertEqual(END_TS, converted.end_time) + self.assertEqual(1, converted.attempt) + # Stored state, not a replay event for the operation. + self.assertFalse(converted.is_replayed) + + def test_updated_operations_is_the_subset_named_by_updated_ids(self): + operations = { + "op-1": self._operation("op-1"), + "op-2": self._operation("op-2"), + } + + with self.executor.run(): + self._start(operations=operations, updated_operation_ids=["op-2"]) + + (start_info,) = self.captured + self.assertEqual(2, len(start_info.operations)) + self.assertEqual(["op-2"], list(start_info.updated_operations)) + # Same converted instance as in the full map, not a second conversion. + self.assertIs( + start_info.operations["op-2"], start_info.updated_operations["op-2"] + ) + + def test_updated_ids_absent_from_the_map_are_skipped(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start( + operations=operations, updated_operation_ids=["op-1", "op-missing"] + ) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.updated_operations)) + + def test_first_invocation_has_empty_maps(self): + with self.executor.run(): + self._start(operations={}, updated_operation_ids=[]) + + (start_info,) = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + + def test_missing_provider_yields_empty_maps(self): + """Callers that supply no provider still get a usable, empty map.""" + with self.executor.run(): + self._start() + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + self.assertEqual({}, end_info.operations) + + def test_plugin_cannot_mutate_checkpointed_error(self): + """A plugin must not be able to alter the error replayed to user code. + + ``ErrorObject.stack_trace`` is a mutable list and the checkpointed error + is handed to user code on replay, so the plugin-facing view gets its own + copy. + """ + checkpoint_error = ErrorObject( + message="boom", type="Error", data=None, stack_trace=["frame-A"] + ) + operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="failing", + parent_id="root", + status=OperationStatus.FAILED, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1, error=checkpoint_error), + ) + + info = OperationInfo.from_operation(operation) + + # A distinct ErrorObject, and a distinct stack_trace list. + self.assertIsNot(checkpoint_error, info.error) + self.assertIsNot(checkpoint_error.stack_trace, info.error.stack_trace) + self.assertEqual(["frame-A"], info.error.stack_trace) + + # Mutating the plugin's view leaves the checkpointed error untouched. + info.error.stack_trace.append("injected-by-plugin") + info.error.stack_trace.clear() + self.assertEqual(["frame-A"], checkpoint_error.stack_trace) + + def test_no_plugins_never_invokes_the_operations_provider(self): + """durable_execution() always passes a provider; an empty executor must + not call it, at either hook.""" + calls = [] + + def provider(): + calls.append(1) + return {} + + executor = PluginExecutor(plugins=[]) + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=provider, + updated_operation_ids=["op-1"], + ) + executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + # The retained provider is dropped too, not just the snapshot. + self.assertIsNone(executor._operations_provider) # noqa: SLF001 + + self.assertEqual([], calls) + + def test_snapshot_is_pinned_against_a_live_provider_mapping(self): + """The snapshot must pin the point in time, whatever the provider returns. + + The SDK's own provider returns a fresh copy, but the eager snapshot is + what actually pins the hook's view, so it must not depend on that. + """ + live = {"op-1": self._operation("op-1")} + seen: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + seen.append(info) + + executor = PluginExecutor(plugins=[_CapturingPlugin()]) + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=lambda: live, + ) + # Mutating the provider's own mapping after the hook must not change + # what that hook reported. + live["op-2"] = self._operation("op-2") + + (start_info,) = seen + self.assertEqual(["op-1"], list(start_info.operations)) + + def test_maps_are_plain_dicts_for_serialization(self): + """The maps must be real dicts, not a custom Mapping. + + ``dataclasses.asdict()`` and ``pickle`` only traverse real dicts, so a + custom Mapping here would leave the hook info unserializable for exactly + the plugins these fields exist to serve. + """ + with self.executor.run(): + self._start( + operations={"op-1": self._operation("op-1")}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + for info in (start_info, end_info): + self.assertIs(dict, type(info.operations)) + self.assertIs(dict, type(start_info.updated_operations)) + + # asdict recurses all the way down, so the values become nested dicts. + as_dict = asdict(start_info) + self.assertIs(dict, type(as_dict["operations"])) + self.assertIs(dict, type(as_dict["operations"]["op-1"])) + self.assertEqual("op-1", as_dict["operations"]["op-1"]["operation_id"]) + self.assertIs(dict, type(as_dict["updated_operations"]["op-1"])) + + def test_hook_infos_survive_a_pickle_round_trip(self): + """Plugins hand infos to multiprocessing queues and caches.""" + with self.executor.run(): + self._start( + operations={"op-1": self._operation("op-1")}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + for info in (start_info, end_info): + restored = pickle.loads(pickle.dumps(info)) # noqa: S301 + self.assertEqual(["op-1"], list(restored.operations)) + self.assertEqual("op-1", restored.operations["op-1"].operation_id) + + # And with the empty maps a plugin-free-style info carries. + empty = InvocationStartInfo( + request_id="req-1", execution_arn="arn:test", is_first_invocation=True + ) + self.assertEqual({}, pickle.loads(pickle.dumps(empty)).operations) # noqa: S301 + + def test_conversion_happens_at_hook_time(self): + """The map is built during the hook, pinning the point in time.""" + conversions: list[str] = [] + real_from_operation = OperationInfo.from_operation + + def counting(operation, **kwargs): + conversions.append(operation.operation_id) + return real_from_operation(operation, **kwargs) + + with ( + patch.object(OperationInfo, "from_operation", staticmethod(counting)), + self.executor.run(), + ): + self._start(operations={"op-1": self._operation("op-1")}) + # Converted while the hook ran, not on first read afterwards. + self.assertEqual(["op-1"], conversions) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.operations)) + # Reading again does not reconvert. + self.assertEqual(["op-1"], conversions) + + def test_end_info_reflects_operations_added_during_the_invocation(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=lambda: operations, + ) + operations["op-2"] = self._operation("op-2") + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual(["op-1", "op-2"], sorted(end_info.operations)) + # The start map is snapshotted at its hook, so a later checkpoint does + # not retroactively change what the start hook reported. + self.assertEqual(["op-1"], sorted(start_info.operations)) + + def test_failing_provider_degrades_to_an_empty_map(self): + def provider(): + raise RuntimeError("boom") + + with ( + self.executor.run(), + self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level="ERROR" + ) as logs, + ): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=provider, + updated_operation_ids=["op-1"], + ) + + self.assertTrue(any("Failed to snapshot operations" in m for m in logs.output)) + (start_info,) = self.captured + self.assertEqual(0, len(start_info.operations)) + self.assertEqual(0, len(start_info.updated_operations)) + + def test_provider_is_released_when_the_run_scope_exits(self): + with self.executor.run(): + self._start(operations={"op-1": self._operation("op-1")}) + self.assertIsNotNone(self.executor._operations_provider) # noqa: SLF001 + + self.assertIsNone(self.executor._operations_provider) # noqa: SLF001 + + +class TestInvocationInfoCopySafety(unittest.TestCase): + """Invocation infos carrying the lazy operation map must stay copyable. + + ``dataclasses.asdict()`` deep-copies any field that is not a dict/list/tuple + or dataclass, so anything uncopyable embedded in the lazy map would make + every plugin that serializes an info raise -- and because plugin exceptions + are swallowed, the hook would be silently dropped. + """ + + def setUp(self): + self.operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="name-op-1", + parent_id="root", + status=OperationStatus.SUCCEEDED, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + def _fire_hooks(self): + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=lambda: {"op-1": self.operation}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + return self.captured + + def test_asdict_on_invocation_start_info(self): + start_info, _ = self._fire_hooks() + + as_dict = asdict(start_info) + + # Both lazy maps materialize as ordinary dicts. + self.assertIsInstance(as_dict["operations"], dict) + self.assertIsInstance(as_dict["updated_operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + self.assertEqual(["op-1"], list(as_dict["updated_operations"])) + + def test_asdict_on_invocation_end_info(self): + _, end_info = self._fire_hooks() + + as_dict = asdict(end_info) + + self.assertIsInstance(as_dict["operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + + def test_deepcopy_on_invocation_infos(self): + start_info, end_info = self._fire_hooks() + + for info in (start_info, end_info): + copied = deepcopy(info) + self.assertIsInstance(copied.operations, dict) + self.assertEqual(["op-1"], list(copied.operations)) + + def test_deepcopy_is_independent_of_the_original(self): + start_info, _ = self._fire_hooks() + + copied = deepcopy(start_info) + copied.operations["op-2"] = OPERATION_END_INFO + + self.assertEqual(["op-1"], list(start_info.operations)) + + def test_unresolved_map_is_still_copyable(self): + """Copying must not require the map to have been read first.""" + start_info, _ = self._fire_hooks() + fresh = InvocationStartInfo( + request_id=start_info.request_id, + execution_arn=start_info.execution_arn, + is_first_invocation=start_info.is_first_invocation, + operations=start_info.operations, + ) + + self.assertIsInstance(asdict(fresh)["operations"], dict) + + class TestPluginExecutorOnOperationAction(unittest.TestCase): """Tests for PluginExecutor.on_operation_action."""