From 57e5bdb422de7a8063009ec71db3d491cb55bf4f Mon Sep 17 00:00:00 2001 From: yaythomas Date: Fri, 16 Oct 2026 10:37:01 +0000 Subject: [PATCH] fix(core): skip plugin hook work without plugins When a checkpoint response changed an operation's status, on_operation_update converted every operation in the execution for the operation change hook, even when no plugin was registered. The checkpoint loop also copied the whole operation map twice for that hook, and create_checkpoint copied it to read one entry. So the cost of each step grew with the number of operations, and long sequential executions slowed down as they ran. Without a plugin, on_operation_update now returns early, and the checkpoint loop skips the copies and the hook call. create_checkpoint reads only the entry it needs. Plugins receive the same information as before. Fixes #780 --- .../plugin.py | 9 +- .../aws_durable_execution_sdk_python/state.py | 23 ++-- .../tests/plugin_test.py | 44 +++++++ .../tests/state_test.py | 112 +++++++++++++++++- 4 files changed, 179 insertions(+), 9 deletions(-) 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)