diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py index 89912b97..aead9011 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/core.py @@ -95,6 +95,11 @@ def apply( ``now`` is the checkpoint's single source of "now", resolved from the execution's clock, so all timestamps stamped by this apply advance with modeled time under a skip clock. + + Before computing the delta, completes every wait and step retry + whose scheduled moment has passed, so those transitions are + delivered in this response and a suspended branch resumes within + the running invocation instead of waiting for the next one. """ effects: list[CheckpointEffect] = [] if updates: @@ -109,6 +114,11 @@ def apply( new_token_sequence: int = execution.advance_token_sequence() + # Pick up scheduler-due completions before pinning the paginator: + # each transition touches its operation, so it lands in this + # response's unseen delta and the handler sees it immediately. + execution.complete_due_operations(now) + paginator: OperationPaginatorState = OperationPaginatorState.pin(execution) # The checkpoint response returns the full unseen delta in a single # response. Advance handler_seen_seq to cover every returned op so 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 b03b45f3..da976f5b 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 @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from dataclasses import dataclass, replace from datetime import datetime from enum import Enum @@ -35,6 +36,8 @@ CheckpointToken, ) +logger = logging.getLogger(__name__) + class ExecutionStatus(Enum): """Execution status for API responses.""" @@ -500,6 +503,54 @@ def complete_retry(self, operation_id: str) -> Operation: self._record_updated_operation(operation_id) return updated_operation + def complete_due_operations(self, now: datetime) -> bool: + """Complete every operation whose scheduled moment has passed. + + Transitions due WAIT operations (``STARTED`` with a + ``scheduled_end_timestamp`` at or before ``now``) via + :meth:`complete_wait` and due STEP retries (``PENDING`` with a + ``next_attempt_timestamp`` at or before ``now``) via + :meth:`complete_retry`. A transition failure is logged and + skipped so one bad operation cannot block the rest. + + Returns True when at least one operation completed. + """ + completed_any: bool = False + for op in list(self.operations): + if ( + op.operation_type is OperationType.WAIT + and op.status is OperationStatus.STARTED + and op.wait_details is not None + and op.wait_details.scheduled_end_timestamp is not None + and op.wait_details.scheduled_end_timestamp <= now + ): + try: + self.complete_wait(op.operation_id, now=now) + completed_any = True + except Exception: # noqa: BLE001 — skip one bad op, keep the rest + logger.exception( + "[%s] complete_wait failed for due operation %s", + self.durable_execution_arn, + op.operation_id, + ) + elif ( + op.operation_type is OperationType.STEP + and op.status is OperationStatus.PENDING + and op.step_details is not None + and op.step_details.next_attempt_timestamp is not None + and op.step_details.next_attempt_timestamp <= now + ): + try: + self.complete_retry(op.operation_id) + completed_any = True + except Exception: # noqa: BLE001 — skip one bad op, keep the rest + logger.exception( + "[%s] complete_retry failed for due operation %s", + self.durable_execution_arn, + op.operation_id, + ) + return completed_any + def complete_callback_success( self, callback_id: str, 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 aa43add7..2b303b2a 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 @@ -1036,41 +1036,7 @@ def _fire_due_operations(self, execution_arn: str) -> bool: if execution.is_complete: return False - now = self._clock.now() - completed_any = False - for op in list(execution.operations): - if ( - op.operation_type == OperationType.WAIT - and op.status == OperationStatus.STARTED - and op.wait_details is not None - and op.wait_details.scheduled_end_timestamp is not None - and op.wait_details.scheduled_end_timestamp <= now - ): - try: - execution.complete_wait(op.operation_id, now=now) - completed_any = True - except Exception: # noqa: BLE001 - logger.exception( - "[%s] earliest-pending: complete_wait failed for %s", - execution_arn, - op.operation_id, - ) - elif ( - op.operation_type == OperationType.STEP - and op.status == OperationStatus.PENDING - and op.step_details is not None - and op.step_details.next_attempt_timestamp is not None - and op.step_details.next_attempt_timestamp <= now - ): - try: - execution.complete_retry(op.operation_id) - completed_any = True - except Exception: # noqa: BLE001 - logger.exception( - "[%s] earliest-pending: complete_retry failed for %s", - execution_arn, - op.operation_id, - ) + completed_any = execution.complete_due_operations(self._clock.now()) if completed_any: self._store.update(execution) return completed_any diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py index bf1ee527..5a4e811f 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/checkpoint/processor_test.py @@ -1,15 +1,19 @@ """Unit tests for CheckpointProcessor.""" +from datetime import UTC, datetime, timedelta from unittest.mock import Mock, patch import pytest from aws_durable_execution_sdk_python.lambda_service import ( CheckpointOutput, CheckpointUpdatedExecutionState, + Operation, OperationAction, + OperationStatus, OperationType, OperationUpdate, StateOutput, + WaitDetails, ) from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( @@ -354,3 +358,53 @@ def test_checkpoint_from_superseded_invocation_is_rejected(): # The current invocation's checkpoint is accepted. result = processor.process_checkpoint(current_token, updates, "client-token") assert isinstance(result, CheckpointOutput) + + +def test_process_checkpoint_delivers_due_wait_completion() -> None: + """A wait whose scheduled end has passed is completed by the + checkpoint and returned in the same response delta.""" + store = InMemoryExecutionStore() + scheduler = Mock(spec=Scheduler) + processor = CheckpointProcessor(store, scheduler) + + start_input = StartDurableExecutionInput( + account_id="123456789012", + function_name="test-function", + function_qualifier="$LATEST", + execution_name="test-execution", + execution_timeout_seconds=300, + execution_retention_period_days=7, + invocation_id="test-inv-id", + ) + execution = Execution.new(start_input) + execution.start() + + past: datetime = datetime.now(UTC) - timedelta(seconds=5) + execution.operations.append( + Operation( + operation_id="wait-1", + parent_id=None, + name="due-wait", + start_timestamp=past, + operation_type=OperationType.WAIT, + status=OperationStatus.STARTED, + wait_details=WaitDetails(scheduled_end_timestamp=past), + ) + ) + store.save(execution) + + token: str = CheckpointToken( + execution_arn=execution.durable_execution_arn, token_sequence=0 + ).to_str() + + result: CheckpointOutput = processor.process_checkpoint(token, [], None) + + returned = {op.operation_id: op for op in result.new_execution_state.operations} + assert "wait-1" in returned + assert returned["wait-1"].status is OperationStatus.SUCCEEDED + + persisted = store.load(execution.durable_execution_arn) + persisted_wait = next( + op for op in persisted.operations if op.operation_id == "wait-1" + ) + assert persisted_wait.status is OperationStatus.SUCCEEDED diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_inline_completion_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_inline_completion_test.py new file mode 100644 index 00000000..d36f4e66 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_inline_completion_test.py @@ -0,0 +1,123 @@ +"""End-to-end test proving a wait inside a parallel branch completes inline. + +Demonstrates that the local runner matches the real service: a wait +started in one parallel branch completes while its sibling is still +working, within a single invocation. Without the fix, the wait would +only complete after the invocation returns and a second invocation +re-checks it. +""" + +from __future__ import annotations + +import json +import time +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.types import StepContext + +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) + + +# Wall-clock seconds branch B keeps working. Must exceed the modeled +# wait plus the SDK's branch-retry overhead so branch A resumes while +# B is still running. +_SIBLING_WORK_SECONDS: float = 3.0 + +# Modeled wait duration for branch A (minimum 1 second). +_WAIT_SECONDS: int = 1 + + +@durable_step +def do_work(step_context: StepContext, seconds: float) -> str: # noqa: ARG001 + """Simulate long-running work by sleeping.""" + time.sleep(seconds) + return "work-done" + + +# Timeline collected during the execution for post-hoc assertions. +_timeline: list[tuple[float, str]] = [] + + +def _record(start: float, msg: str) -> None: + _timeline.append((time.time() - start, msg)) + + +@durable_execution +def parallel_wait_and_work(event: Any, context: DurableContext) -> list[str]: # noqa: ARG001 + """Two parallel branches: A waits, B does work.""" + start: float = time.time() + _record(start, "INVOCATION START") + + def branch_a(ctx: DurableContext) -> str: + _record(start, "A: before wait") + ctx.wait(Duration.from_seconds(_WAIT_SECONDS)) + _record(start, "A: after wait") + return "a-done" + + def branch_b(ctx: DurableContext) -> str: + _record(start, "B: start work") + ctx.step(do_work(_SIBLING_WORK_SECONDS)) + _record(start, "B: finished work") + return "b-done" + + batch = context.parallel(functions=[branch_a, branch_b], name="test-parallel") + results: list[str] = batch.get_results() + _record(start, "parallel returned") + return results + + +@pytest.mark.parametrize("skip_time", [True, False]) +def test_wait_completes_inline_during_parallel_invocation(skip_time: bool) -> None: # noqa: FBT001 + """A wait in one parallel branch completes while the sibling is still + working, within a single invocation. + + The wait's completion is delivered through a checkpoint response + while the invocation is in flight, so the waiting branch resumes + without an extra invocation. Runs on both the skip clock and the + wall clock. + """ + _timeline.clear() + + with DurableFunctionTestRunner( + handler=parallel_wait_and_work, + execution_timeout=15, + skip_time=skip_time, + ) as runner: + result: DurableFunctionTestResult = runner.run(input="x") + + assert result.status is InvocationStatus.SUCCEEDED + assert result.result == json.dumps(["a-done", "b-done"]) + + # Branch A must resume from its wait BEFORE branch B finishes its + # work. This proves the wait completed inline while the sibling was + # still working. + a_after_wait_times = [t for t, msg in _timeline if msg == "A: after wait"] + b_finished_times = [t for t, msg in _timeline if msg == "B: finished work"] + + assert len(a_after_wait_times) >= 1, "Branch A never resumed after wait" + assert len(b_finished_times) >= 1, "Branch B never finished" + + a_after_wait: float = a_after_wait_times[0] + b_finished: float = b_finished_times[0] + + assert a_after_wait < b_finished, ( + f"Branch A resumed at {a_after_wait:.2f}s but B finished at " + f"{b_finished:.2f}s — wait did not complete inline" + ) + + # Only ONE invocation should be recorded. Look for exactly one + # "INVOCATION START" event. + invocation_starts = [msg for _, msg in _timeline if msg == "INVOCATION START"] + assert len(invocation_starts) == 1, ( + f"Expected 1 invocation, got {len(invocation_starts)}" + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py new file mode 100644 index 00000000..755e4e16 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/wait_suspend_replay_test.py @@ -0,0 +1,96 @@ +"""End-to-end test: a wait completed while suspended is delivered as new. + +A top-level wait suspends the execution. The wait completes between +invocations, and the next invocation must observe the completion as a +NEW operation update rather than a replayed one, finishing the +execution in exactly two invocations. Guards the plugin-visible +contract: ``on_operation_end`` fires once for the wait with +``is_replayed`` False and an end time set. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +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 ( + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, + OperationEndInfo, +) + +from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + DurableFunctionTestRunner, +) + +_WAIT_NAME: str = "suspend-wait" + + +class RecordingWaitPlugin(DurableInstrumentationPlugin): + """Records end notifications for the wait and counts invocations.""" + + invocation_count: ClassVar[int] = 0 + wait_end_infos: ClassVar[list[OperationEndInfo]] = [] + + @classmethod + def reset(cls) -> None: + cls.invocation_count = 0 + cls.wait_end_infos.clear() + + def on_invocation_start(self, info: InvocationStartInfo) -> None: # noqa: ARG002 + type(self).invocation_count += 1 + + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.operation_type is OperationType.WAIT and info.name == _WAIT_NAME: + self.wait_end_infos.append(info) + + +def _wait_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + """Suspend on a top-level wait, then finish.""" + context.wait(Duration.from_seconds(1), name=_WAIT_NAME) + return "done" + + +wait_handler = durable_execution(_wait_handler, plugins=[RecordingWaitPlugin()]) + + +def test_wait_completed_during_suspend_is_delivered_as_new() -> None: + """The resumed invocation sees the wait completion as a new update. + + The runner takes exactly one suspend and one resume invocation, and + the plugin observes exactly one SUCCEEDED end notification for the + wait with ``is_replayed`` False. + """ + RecordingWaitPlugin.reset() + + with DurableFunctionTestRunner( + handler=wait_handler, execution_timeout=15 + ) as runner: + result: DurableFunctionTestResult = runner.run(input="{}") + + assert result.status is InvocationStatus.SUCCEEDED + + wait_op = result.get_wait(_WAIT_NAME) + assert wait_op.status is OperationStatus.SUCCEEDED + + # Exactly one suspending invocation and one resuming invocation. + assert RecordingWaitPlugin.invocation_count == 2 + + succeeded_infos: list[OperationEndInfo] = [ + info + for info in RecordingWaitPlugin.wait_end_infos + if info.status is OperationStatus.SUCCEEDED + ] + assert len(succeeded_infos) == 1 + assert succeeded_infos[0].is_replayed is False + assert succeeded_infos[0].end_time is not None diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/execution_wait_retry_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/execution_wait_retry_test.py index 53e7a1d7..211916ee 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/execution_wait_retry_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/execution_wait_retry_test.py @@ -2,13 +2,14 @@ import threading from concurrent.futures import ThreadPoolExecutor, as_completed -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from aws_durable_execution_sdk_python.lambda_service import ( Operation, OperationStatus, OperationType, StepDetails, + WaitDetails, ) from aws_durable_execution_sdk_python_testing.execution import Execution @@ -80,3 +81,127 @@ def complete_retry(): # that only advances on accepted checkpoint calls). assert execution.seq_counter == 2 assert execution.token_sequence == 0 + + +def _make_execution() -> Execution: + """Create a bare execution for due-operation tests.""" + input_data = StartDurableExecutionInput( + account_id="123456789012", + function_name="test-function", + function_qualifier="$LATEST", + execution_name="test-execution", + execution_timeout_seconds=300, + execution_retention_period_days=7, + invocation_id="test-inv-id", + input='{"test": "data"}', + ) + return Execution.new(input_data) + + +def test_complete_due_operations_completes_due_wait_and_retry() -> None: + """A due wait transitions to SUCCEEDED and a due step retry to READY.""" + execution = _make_execution() + now: datetime = datetime.now(UTC) + past: datetime = now - timedelta(seconds=5) + + execution.operations.append( + Operation( + operation_id="wait-1", + parent_id=None, + name="test-wait", + start_timestamp=past, + operation_type=OperationType.WAIT, + status=OperationStatus.STARTED, + wait_details=WaitDetails(scheduled_end_timestamp=past), + ) + ) + execution.operations.append( + Operation( + operation_id="step-1", + parent_id=None, + name="test-step", + start_timestamp=past, + operation_type=OperationType.STEP, + status=OperationStatus.PENDING, + step_details=StepDetails(next_attempt_timestamp=past), + ) + ) + + completed_any: bool = execution.complete_due_operations(now) + + assert completed_any is True + + wait_op = next(op for op in execution.operations if op.operation_id == "wait-1") + assert wait_op.status is OperationStatus.SUCCEEDED + assert wait_op.end_timestamp == now + + step_op = next(op for op in execution.operations if op.operation_id == "step-1") + assert step_op.status is OperationStatus.READY + assert step_op.step_details is not None + assert step_op.step_details.next_attempt_timestamp is None + + # Each completion touches its operation, so the transitions land in + # the next checkpoint response delta. + assert "wait-1" in execution.operation_last_touched_seq + assert "step-1" in execution.operation_last_touched_seq + + +def test_complete_due_operations_ignores_operations_not_yet_due() -> None: + """Operations scheduled in the future are left untouched.""" + execution = _make_execution() + now: datetime = datetime.now(UTC) + future: datetime = now + timedelta(seconds=60) + + execution.operations.append( + Operation( + operation_id="wait-1", + parent_id=None, + name="test-wait", + start_timestamp=now, + operation_type=OperationType.WAIT, + status=OperationStatus.STARTED, + wait_details=WaitDetails(scheduled_end_timestamp=future), + ) + ) + execution.operations.append( + Operation( + operation_id="step-1", + parent_id=None, + name="test-step", + start_timestamp=now, + operation_type=OperationType.STEP, + status=OperationStatus.PENDING, + step_details=StepDetails(next_attempt_timestamp=future), + ) + ) + + completed_any: bool = execution.complete_due_operations(now) + + assert completed_any is False + assert execution.operations[0].status is OperationStatus.STARTED + assert execution.operations[1].status is OperationStatus.PENDING + assert execution.operation_last_touched_seq == {} + + +def test_complete_due_operations_ignores_non_pending_statuses() -> None: + """Terminal and unstarted operations never match the due predicate.""" + execution = _make_execution() + now: datetime = datetime.now(UTC) + past: datetime = now - timedelta(seconds=5) + + execution.operations.append( + Operation( + operation_id="wait-1", + parent_id=None, + name="test-wait", + start_timestamp=past, + operation_type=OperationType.WAIT, + status=OperationStatus.SUCCEEDED, + wait_details=WaitDetails(scheduled_end_timestamp=past), + ) + ) + + completed_any: bool = execution.complete_due_operations(now) + + assert completed_any is False + assert execution.operation_last_touched_seq == {} diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py index 6771895d..c789dcd7 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py @@ -25,6 +25,7 @@ OperationUpdate, StepDetails, WaitDetails, + WaitOptions, ) from aws_durable_execution_sdk_python_testing.exceptions import ( @@ -172,7 +173,9 @@ def test_async_completion_between_polls_surfaces_on_next_poll(): executor, store, execution, token_0 = _make_executor_with_started_execution() arn = execution.durable_execution_arn - # First checkpoint: handler declares a 1-second wait. + # First checkpoint: handler declares a 60-second wait. The duration + # must be far enough out that the checkpoint's own due-operation + # pickup does not complete it immediately. r1 = executor.checkpoint_execution( execution_arn=arn, checkpoint_token=token_0, @@ -182,6 +185,7 @@ def test_async_completion_between_polls_surfaces_on_next_poll(): operation_type=OperationType.WAIT, action=OperationAction.START, name="wait-1", + wait_options=WaitOptions(wait_seconds=60), ), ], )