diff --git a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md index 0a33e4ec1..a92eba592 100644 --- a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md +++ b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md @@ -167,9 +167,9 @@ the handler would wait on a lane it needs in order to make progress. 2. Invoke. The blocking call to the function runs off the lane, bounded by the invocation timeout. While it blocks, the handler's checkpoint and get state tasks run on the now free lane. -3. Finish. A task on the lane records the response and releases the - gate. It moves `PENDING` to `PRE_INVOKE`, or a terminal result to - `COMPLETED`. +3. Finish. A task on the lane validates the response, acknowledges its input + snapshot, and releases the gate. It moves `PENDING` to `PRE_INVOKE`, or a + terminal result to `COMPLETED`. ```mermaid sequenceDiagram @@ -234,10 +234,17 @@ list in a checkpoint response. The delta for the next response is the operations whose `operation_last_touched_seq` is greater than `handler_seen_seq`. -The watermark advances only when a checkpoint response is built. It never -advances when invocation input is built. That is what makes retries -correct. A failed invocation never moves the watermark. So the retry -delivers the same state again. +The watermark advances when a checkpoint response is built or a valid +invocation response acknowledges its input snapshot, including a `PENDING` +response that sent no checkpoint. Constructing invocation input does not +acknowledge delivery: a failed delivery leaves its updates available for retry. + +Acknowledgment removes `UpdatedOperationIds` only when their last-touched +sequence is covered by that snapshot. Newer versions, including another +version of the same operation, remain pending. An invocation's acknowledgment +never retreats a watermark already advanced by a checkpoint. Legacy update +IDs without a stored version are treated as sequence zero and can be consumed +by an equal-sequence acknowledgment. ``` seq_counter handler_seen_seq delta next checkpoint diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 5190b9b69..4ffa24c6f 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -479,15 +479,11 @@ def record_invocation_completion( ) def mark_state_delivered(self) -> None: - """Reset the list of changed operations once an invocation's input - carries them. - - The service reports, on each invocation, the operations that - changed after the last state the handler observed. The runner - delivers state in the invocation input. So the list is reset when - that input is built, not when the invocation completes: an - operation that completes while the handler is still running is - reported on the next invocation. + """Explicitly clear all pending operation IDs after their delivery. + + The executor uses versioned ``advance_handler_seen`` acknowledgments + so failed invocations retain their input updates and changes newer + than the acknowledged snapshot remain pending. """ self.updated_operation_ids = [] @@ -860,10 +856,9 @@ class OperationPaginatorState: * :meth:`unseen_operations` — strict ``> handler_seen_seq`` delta filter, used by ``checkpoint_execution`` to decide which ops go back to the handler. - * :meth:`advance_handler_seen` — monotonic forward update of the - delivery watermark on the wrapped :class:`Execution`, called - after a checkpoint response is constructed and its ops are - known. + * :meth:`advance_handler_seen` — acknowledge delivered operation versions + without retreating the watermark, after constructing a checkpoint + response or validating an invocation's response. Pinning ``token_sequence`` and ``operations`` at construction @@ -920,11 +915,18 @@ def unseen_operations(self) -> list[Operation]: ] def advance_handler_seen(self, seq: int) -> None: - """Advance the wrapped execution's ``handler_seen_seq`` to - ``seq`` if that represents forward progress. Monotonic: - smaller or equal values are ignored.""" + """Acknowledge delivered updates without retreating the watermark. + + Equal acknowledgments also consume legacy, unversioned update IDs. + Only versions covered by this delivery may be removed. + """ if seq > self.execution.handler_seen_seq: self.execution.handler_seen_seq = seq + self.execution.updated_operation_ids = [ + operation_id + for operation_id in self.execution.updated_operation_ids + if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq + ] # --- internals ------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py index 4c3a0a71f..d2056daa5 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py @@ -1551,7 +1551,8 @@ def _begin_invocation( self._set_invocation_gate(execution_arn, InvocationState.INVOKING) execution.begin_new_invocation() invocation_input = self._invoker.create_invocation_input(execution=execution) - execution.mark_state_delivered() + # Retain input updates until a valid response acknowledges this snapshot. + # A failed invocation must be able to receive the same updates on retry. self._store.save(execution) return execution, invocation_input, execution.seq_counter @@ -1618,6 +1619,12 @@ def _finish_invocation( # A clean invocation resets the retry counter. reset_target = self._store.load(execution_arn) reset_target.consecutive_failed_invocation_attempts = 0 + if invocation_seq is not None: + # A valid response acknowledges input even without a checkpoint. + # Preserve changes newer than the input and any newer checkpoint ack. + OperationPaginatorState.pin(reset_target).advance_handler_seen( + invocation_seq + ) self._store.save(reset_target) # Release the gate: terminal -> COMPLETED, otherwise PRE_INVOKE diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py new file mode 100644 index 000000000..c8a399fef --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py @@ -0,0 +1,183 @@ +"""A wait update survives failed delivery and a checkpoint-free resume.""" + +import json +import time +from pathlib import Path +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python import ( + DurableContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import CallbackConfig, Duration +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationInput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + OperationChangeInfo, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.execution import Execution +from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("failed_delivery", [False, True]) +@pytest.mark.parametrize("legacy_format", [False, True]) +def test_wait_input_acknowledgment_survives_checkpoint_free_resume( + filesystem: bool, + failed_delivery: bool, + legacy_format: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store: ExecutionStore = ( + FileSystemExecutionStore.create(tmp_path / "store") + if filesystem + else InMemoryExecutionStore() + ) + starts: list[InvocationStartInfo] = [] + ends: list[InvocationEndInfo] = [] + wait_deliveries: list[int] = [] + attempts: list[int] = [] + step_calls: list[str] = [] + legacy_restores: list[dict[str, Any]] = [] + original_begin = Executor._begin_invocation # noqa: SLF001 + + def begin_with_legacy_state( + executor: Executor, arn: str + ) -> tuple[Execution, DurableExecutionInvocationInput, int] | None: + if ( + legacy_format + and not legacy_restores + and executor._invocation_gate(arn) is InvocationState.PRE_INVOKE # noqa: SLF001 + ): + execution = store.load(arn) + if execution.updated_operation_ids: + # Rehydrate an actual wait completion in the older store format. + snapshot = execution.to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + restored = Execution.from_json_dict(snapshot) + assert restored.seq_counter == restored.handler_seen_seq == 0 + assert restored.updated_operation_ids == execution.updated_operation_ids + store.save(restored) + legacy_restores.append(snapshot) + return original_begin(executor, arn) + + monkeypatch.setattr(Executor, "_begin_invocation", begin_with_legacy_state) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info: InvocationStartInfo) -> None: + starts.append(info) + if any( + op.name == "delivered-wait" and op.status.value == "SUCCEEDED" + for op in info.updated_operations.values() + ): + wait_deliveries.append(len(starts)) + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + ends.append(info) + + def on_operation_change(self, info: OperationChangeInfo) -> None: + if any( + op.name == "delivered-wait" and op.status.value == "SUCCEEDED" + for op in info.updated_operations.values() + ): + wait_deliveries.append(len(starts)) + + class ObserveFactory: + def create_plugin( + self, info: InvocationStartInfo + ) -> DurableInstrumentationPlugin: + return Observe() + + @durable_step + def marker(_context: Any, value: str) -> str: + step_calls.append(value) + return value + + def durable_handler(_event: Any, context: DurableContext) -> list[str]: + before = context.step(marker("before"), name="before-wait") + barrier = context.create_callback( + name="barrier", config=CallbackConfig(serdes=JsonSerDes()) + ) + context.wait(Duration.from_seconds(1), name="delivered-wait") + value = barrier.result() + assert isinstance(value, str) + return [before, context.step(marker(value), name="after-barrier")] + + handler = durable_execution(durable_handler, plugins=[ObserveFactory()]) + + def entrypoint(event: Any, context: Any) -> Any: + attempts.append(1) + if failed_delivery and len(attempts) == 2: + raise RuntimeError("delivery failed before the SDK received its input") + return handler(event, context) + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + callback_id = runner.wait_for_callback(arn, name="barrier", timeout=5) + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn).events + if ( + len(ends) >= 2 + and sum(event.event_type == "InvocationCompleted" for event in events) + >= 2 + ): + break + time.sleep(0.01) + else: + raise AssertionError("The checkpoint-free resumed invocation did not end") + saved = store.load(arn) + wait = next(op for op in saved.operations if op.name == "delivered-wait") + acknowledged = ( + list(saved.updated_operation_ids), + saved.handler_seen_seq, + saved.operation_last_touched_seq.get(wait.operation_id, 0), + ) + runner.send_callback_success(callback_id, result=b'"after"') + result = runner.wait_for_result(arn, timeout=10) + + assert result.status is InvocationStatus.SUCCEEDED + assert result.result is not None + assert json.loads(result.result) == ["before", "after"] + pending_ids, seen_sequence, wait_sequence = acknowledged + assert pending_ids == [] + assert seen_sequence >= wait_sequence + assert [info.status.value for info in ends] == [ + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert set(wait_deliveries) == {2} + assert step_calls == ["before", "after"] + assert len(attempts) == (4 if failed_delivery else 3) + assert len(legacy_restores) == int(legacy_format) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py index 21993649e..628cf8956 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py @@ -736,6 +736,9 @@ def test_pending_response_is_valid_when_an_operation_completed_after_the_handler mock_execution.consecutive_failed_invocation_attempts = 0 mock_execution.seq_counter = 4 mock_execution.handler_seen_seq = 2 + mock_execution.operations = [] + mock_execution.updated_operation_ids = [] + mock_execution.operation_last_touched_seq = {} mock_execution.has_pending_operations.return_value = False mock_execution.has_changes_after.side_effect = lambda seq: seq == 4 mock_execution.has_unseen_changes.return_value = True diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py new file mode 100644 index 000000000..225f087b3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py @@ -0,0 +1,239 @@ +"""Only successful state delivery acknowledges pending operation versions.""" + +from collections.abc import Iterator +from datetime import UTC, datetime +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationOutput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + Operation, + OperationAction, + OperationStatus, + OperationType, + OperationUpdate, +) +from aws_durable_execution_sdk_python_testing.exceptions import ( + InvalidParameterValueException, +) +from aws_durable_execution_sdk_python_testing.execution import ( + Execution, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState +from aws_durable_execution_sdk_python_testing.invoker import ( + InProcessInvoker, + InvokeResponse, +) +from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput +from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) + + +@pytest.fixture(params=["memory", "filesystem"]) +def stored_execution( + request: pytest.FixtureRequest, tmp_path: Path +) -> Iterator[tuple[Executor, ExecutionStore, str]]: + store: ExecutionStore = ( + FileSystemExecutionStore.create(tmp_path / "store") + if request.param == "filesystem" + else InMemoryExecutionStore() + ) + execution = Execution.new( + StartDurableExecutionInput( + account_id="123456789012", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=30, + execution_retention_period_days=1, + invocation_id="invocation", + ) + ) + execution.start() + for name in ("delivered", "later", "checkpoint"): + execution.operations.append( + Operation( + operation_id=name, + operation_type=OperationType.WAIT, + status=OperationStatus.STARTED, + ) + ) + execution.touch_operation(name) + execution.handler_seen_seq = execution.seq_counter + execution.complete_wait("delivered") + store.save(execution) + executor = Executor(store, Mock(), InProcessInvoker(Mock(), Mock()), Mock()) + try: + yield executor, store, execution.durable_execution_arn + finally: + executor.shutdown() + executor._registry.shutdown() # noqa: SLF001 + + +def test_building_input_does_not_acknowledge_delivery( + stored_execution: tuple[Executor, ExecutionStore, str], +) -> None: + executor, store, arn = stored_execution + before = store.load(arn).handler_seen_seq + claim = executor._begin_invocation(arn) # noqa: SLF001 + assert claim is not None + _, invocation_input, sequence = claim + assert invocation_input.updated_operation_ids == ["delivered"] + saved = store.load(arn) + assert saved.updated_operation_ids == ["delivered"] + assert saved.handler_seen_seq == before < sequence + + +@pytest.mark.parametrize("status", list(InvocationStatus)) +def test_only_valid_response_acknowledges_input( + stored_execution: tuple[Executor, ExecutionStore, str], + status: InvocationStatus, +) -> None: + executor, store, arn = stored_execution + execution = store.load(arn) + sequence = execution.seq_counter + before = execution.handler_seen_seq + executor._finish_invocation( # noqa: SLF001 + arn, + InvokeResponse( + invocation_output=DurableExecutionInvocationOutput(status=status), + request_id="request", + ), + datetime.now(UTC), + datetime.now(UTC), + sequence, + ) + saved = store.load(arn) + valid = status is not InvocationStatus.RETRY + assert saved.handler_seen_seq == (sequence if valid else before) + assert saved.updated_operation_ids == ([] if valid else ["delivered"]) + + +@pytest.mark.parametrize("checkpoint_during_invocation", [False, True]) +def test_input_ack_preserves_newer_completion_and_checkpoint_watermark( + stored_execution: tuple[Executor, ExecutionStore, str], + checkpoint_during_invocation: bool, +) -> None: + executor, store, arn = stored_execution + execution = store.load(arn) + input_sequence = execution.seq_counter + if checkpoint_during_invocation: + execution.complete_wait("checkpoint") + OperationPaginatorState.pin(execution).advance_handler_seen( + execution.seq_counter + ) + already_seen = execution.handler_seen_seq + execution.complete_wait("later") + execution.needs_reinvoke = True + store.save(execution) + with patch.object(executor, "_invoke_execution") as reinvoke: + executor._finish_invocation( # noqa: SLF001 + arn, + InvokeResponse( + invocation_output=DurableExecutionInvocationOutput( + status=InvocationStatus.PENDING + ), + request_id="request", + ), + datetime.now(UTC), + datetime.now(UTC), + input_sequence, + ) + saved = store.load(arn) + assert saved.handler_seen_seq == max(already_seen, input_sequence) + assert saved.updated_operation_ids == ["later"] + assert [ + op.operation_id for op in OperationPaginatorState.pin(saved).unseen_operations() + ] == ["later"] + reinvoke.assert_called_once_with(arn) + + +@pytest.mark.parametrize("seen,ack", [(0, 0), (5, 5), (5, 3)]) +def test_acknowledgment_prunes_only_covered_versions_in_legacy_state( + stored_execution: tuple[Executor, ExecutionStore, str], seen: int, ack: int +) -> None: + _, store, arn = stored_execution + snapshot = store.load(arn).to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + snapshot["UpdatedOperationIds"] = ["legacy", "covered", "newer"] + restored = Execution.from_json_dict(snapshot) + restored.handler_seen_seq = seen + restored.operation_last_touched_seq.update(covered=ack, newer=ack + 1) + store.save(restored) + execution = store.load(arn) + paginator = OperationPaginatorState.pin(execution) + paginator.advance_handler_seen(ack) + paginator.advance_handler_seen(ack) + assert execution.updated_operation_ids == ["newer"] + assert execution.handler_seen_seq == seen + assert execution.seq_counter == 0 + + +def test_older_acknowledgment_preserves_newer_version_of_same_operation( + stored_execution: tuple[Executor, ExecutionStore, str], +) -> None: + _, store, arn = stored_execution + execution = store.load(arn) + sequence = execution.seq_counter + # A new version of an already-pending ID must survive the old snapshot's ack. + execution.touch_operation("delivered") + OperationPaginatorState.pin(execution).advance_handler_seen(sequence) + assert execution.updated_operation_ids == ["delivered"] + OperationPaginatorState.pin(execution).advance_handler_seen(execution.seq_counter) + assert execution.updated_operation_ids == [] + + +@pytest.mark.parametrize("rejection", ["token", "stale", "operation"]) +def test_rejected_checkpoint_retains_undelivered_update( + stored_execution: tuple[Executor, ExecutionStore, str], rejection: str +) -> None: + executor, store, arn = stored_execution + execution = store.load(arn) + token = execution.get_new_checkpoint_token() + if rejection == "stale": + execution.advance_token_sequence() + store.save(execution) + before = ( + execution.token_sequence, + execution.handler_seen_seq, + execution.seq_counter, + ) + executor._set_invocation_gate(arn, InvocationState.INVOKING) # noqa: SLF001 + with pytest.raises(InvalidParameterValueException): + executor.checkpoint_execution( + execution_arn=arn, + checkpoint_token="invalid" if rejection == "token" else token, + updates=( + [ + OperationUpdate( + operation_id="delivered", + operation_type=OperationType.WAIT, + action=OperationAction.FAIL, + error=ErrorObject.from_message("invalid wait failure"), + ) + ] + if rejection == "operation" + else [] + ), + ) + saved = store.load(arn) + assert saved.updated_operation_ids == ["delivered"] + assert (saved.token_sequence, saved.handler_seen_seq, saved.seq_counter) == before