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 856fc68c6..6030c4ed2 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 @@ -471,6 +471,11 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] + @property + def has_plugins(self) -> bool: + """Return True when at least one plugin is registered.""" + return bool(self._plugins) + @contextlib.contextmanager def run(self): if self._plugins: @@ -805,7 +810,9 @@ def on_operation_update( operations: full operation map after the update, when available. previous_operations: operation map before the update, when available. """ - if operation_or_operations is None: + # Only plugins read the info built below, and the change info converts + # every operation. So without a plugin, return before building any of it. + if operation_or_operations is None or not self.has_plugins: return updated_operations: list[Operation] = ( diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py index f5ce7214a..de396b11e 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py @@ -776,11 +776,15 @@ def create_checkpoint( ) if operation_update is not None: + # Read the one entry under the lock. self.operations would copy the + # whole map on every checkpoint to look up a single operation. + with self._operations_lock: + previous_operation = self._operations.get(operation_update.operation_id) # Dispatch before queueing so START strictly precedes any user # function attempt, regardless of checkpoint synchronization mode. self._plugin_executor.on_operation_action( operation_update, - previous_operation=self.operations.get(operation_update.operation_id), + previous_operation=previous_operation, ) # Create wrapper object for queue @@ -1036,7 +1040,11 @@ def checkpoint_batches_forever(self) -> None: "checkpoint." ) - previous_operations = self.operations + # Each read of self.operations copies the whole map, and only + # plugins use these copies. So take the snapshots, and call + # the hook, only when a plugin is registered. + has_plugins = self._plugin_executor.has_plugins + previous_operations = self.operations if has_plugins else None # Fetch new operations from the API before unblocking sync # waiters. On completion the token is consumed, so skip @@ -1052,11 +1060,12 @@ def checkpoint_batches_forever(self) -> None: current_checkpoint_token, fetch_marker, ) - self._plugin_executor.on_operation_update( - updated_operations, - self.operations, - previous_operations, - ) + if has_plugins: + self._plugin_executor.on_operation_update( + updated_operations, + self.operations, + previous_operations, + ) # Signal completion for any synchronous operations for queued_op in batch: 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 70d266e3a..db5a74cd1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -540,6 +540,13 @@ def test_init_with_plugins(self): executor = PluginExecutor(plugins=[p1, p2]) self.assertEqual(len(executor._plugins), 2) + def test_has_plugins_is_false_without_plugins(self) -> None: + self.assertFalse(PluginExecutor(plugins=None).has_plugins) + self.assertFalse(PluginExecutor(plugins=[]).has_plugins) + + def test_has_plugins_is_true_with_a_plugin(self) -> None: + self.assertTrue(PluginExecutor(plugins=[_NoOpPlugin()]).has_plugins) + class TestPluginExecutor(unittest.TestCase): def test_no_thread_pool_when_plugins_is_none(self): @@ -1662,6 +1669,43 @@ def test_operation_change_without_invocation_start_is_noop(self): self.assertEqual(self.plugin.calls, []) + def test_operation_update_without_plugins_builds_no_info(self) -> None: + """Without a plugin, an update converts no operation and builds no info.""" + operation = Operation( + operation_id="op-1", + operation_type=ServiceOperationType.STEP, + status=OperationStatus.SUCCEEDED, + step_details=StepDetails(attempt=1, result='"done"', error=None), + ) + executor = PluginExecutor(plugins=None) + plugin_module = "aws_durable_execution_sdk_python.plugin" + + with executor.run(): + # The invocation status and the terminal, changed operation reach + # every info constructor when the update does not return early. + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + ) + with ( + patch.object(OperationInfo, "from_operation") as from_operation, + patch(f"{plugin_module}._to_operation_info_map") as to_info_map, + patch(f"{plugin_module}.OperationEndInfo") as end_info, + patch(f"{plugin_module}.OperationChangeInfo") as change_info, + ): + executor.on_operation_update( + [operation], + {"op-1": operation}, + previous_operations={}, + ) + + from_operation.assert_not_called() + to_info_map.assert_not_called() + end_info.assert_not_called() + change_info.assert_not_called() + class TestPluginExecutorExtractError(unittest.TestCase): """Tests for PluginExecutor._extract_error static method.""" diff --git a/packages/aws-durable-execution-sdk-python/tests/state_test.py b/packages/aws-durable-execution-sdk-python/tests/state_test.py index 0634cb2a3..d67cfb895 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -12,7 +12,7 @@ import time import unittest.mock from concurrent.futures import ThreadPoolExecutor -from unittest.mock import Mock, call, create_autospec, patch +from unittest.mock import Mock, PropertyMock, call, create_autospec, patch import pytest @@ -4765,6 +4765,60 @@ def test_existing_operation_start_is_reported_as_replayed(): assert plugin.operation_starts[0].is_replayed is True +@pytest.mark.parametrize( + ("operations", "expected_is_replayed"), + [ + pytest.param({}, False, id="new-operation"), + pytest.param( + { + "step-1": Operation( + operation_id="step-1", + operation_type=OperationType.STEP, + status=OperationStatus.READY, + ) + }, + True, + id="existing-operation", + ), + ], +) +def test_create_checkpoint_reads_previous_operation_without_copying_map( + operations: dict[str, Operation], expected_is_replayed: bool +) -> None: + """START replay detection reads one entry instead of copying the map.""" + plugin = _RecordingPlugin() + plugin_executor = PluginExecutor(plugins=[plugin]) + with ( + plugin_executor.run(), + patch.object( + ExecutionState, "operations", new_callable=PropertyMock + ) as operations_property, + ): + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations=operations, + service_client=create_autospec(LambdaClient), + plugin_executor=plugin_executor, + ) + # An async checkpoint dispatches START before it returns, so the test + # needs no background thread. + state.create_checkpoint( + OperationUpdate( + operation_id="step-1", + operation_type=OperationType.STEP, + action=OperationAction.START, + name="my-step", + ), + is_sync=False, + ) + + operations_property.assert_not_called() + assert [info.is_replayed for info in plugin.operation_starts] == [ + expected_is_replayed + ] + + def test_plugin_executor_on_operation_update_called_for_terminal_operations(): """Test that plugin_executor.on_operation_update is called for terminal operations.""" mock_client = create_autospec(LambdaClient) @@ -5008,6 +5062,62 @@ def test_plugin_executor_on_operation_change_called_for_status_changes(): assert "operation_change:wait-1" not in plugin.calls +def test_checkpoint_loop_without_plugins_skips_operation_snapshots() -> None: + """Without plugins, a sync checkpoint copies no map and skips the hook.""" + mock_client = create_autospec(spec=LambdaClient) + step_op = Operation( + operation_id="step-1", + operation_type=OperationType.STEP, + status=OperationStatus.SUCCEEDED, + step_details=StepDetails(attempt=1, result='"done"'), + ) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=[step_op], + next_marker=None, + ), + ) + plugin_executor = PluginExecutor(plugins=None) + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=mock_client, + plugin_executor=plugin_executor, + ) + + with ( + patch.object( + ExecutionState, "operations", new_callable=PropertyMock + ) as operations_property, + patch.object(plugin_executor, "on_operation_update") as on_operation_update, + ): + executor = ThreadPoolExecutor(max_workers=1) + executor.submit(state.checkpoint_batches_forever) + try: + # Returns only after the loop has processed the batch, so the + # assertions below see the whole round trip. + state.create_checkpoint( + OperationUpdate( + operation_id="step-1", + operation_type=OperationType.STEP, + action=OperationAction.SUCCEED, + name="my-step", + payload='"done"', + ), + is_sync=True, + ) + finally: + state.stop_checkpointing() + executor.shutdown(wait=True) + + operations_property.assert_not_called() + on_operation_update.assert_not_called() + # The loop still stores the operations from the response. + assert state.operations == {"step-1": step_op} + + def test_operation_start_plugin_hook_fires_before_checkpoint_failure(): """START is reported before queueing even when checkpointing later fails.""" mock_client = create_autospec(spec=LambdaClient)