diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 95f2b1f2..d335b2b5 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -68,7 +68,7 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): Operation IDs are converted into deterministic span IDs. The first observed span for an operation uses that deterministic ID; later continuation spans use newly generated span IDs and link back to the deterministic span ID so - trace viewers can relate retries and replay-created terminal spans to the + trace viewers can relate retries and cross-invocation completions to the original logical operation. Args: @@ -388,7 +388,7 @@ def on_operation_start(self, info: OperationStartInfo) -> None: operation_id=info.operation_id, name=info.name or info.operation_id, attributes=attributes, - start_time=info.start_time, + start_time=datetime.datetime.now(datetime.UTC), parent_span=parent_span, existed=info.is_replayed, ) @@ -407,7 +407,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # Context operations are tracked using on_user_function_end. return span = self._get_span(info.operation_id) - if not span: + if span is None: # the span was not started in the current invocation, so we need to # create a new one that links to the previous one parent_span = self._resolve_parent_span(info.parent_id) @@ -416,7 +416,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: operation_id=info.operation_id, name=info.name or info.operation_id, attributes=attributes, - start_time=info.start_time, + start_time=datetime.datetime.now(datetime.UTC), parent_span=parent_span, existed=True, ) @@ -431,10 +431,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: else: span.set_status(StatusCode.OK) - end_timestamp = info.end_time - if end_timestamp is not None and end_timestamp == info.start_time: - end_timestamp += datetime.timedelta(microseconds=1) - self._end_span(info.operation_id, end_timestamp) + self._end_span(info.operation_id) def on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index bfc6fd4d..1dd00570 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -209,7 +209,7 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - # Operation end with no matching start (started in a prior invocation). + # Backend-updated completion for an operation started in a prior invocation. plugin.on_operation_end( OperationEndInfo( operation_id="step-earlier", @@ -218,7 +218,7 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id(): name="earlier-step", parent_id=None, start_time=START_TIME, - is_replayed=True, + is_replayed=False, status=OperationStatus.SUCCEEDED, end_time=END_TIME, error=None, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index e906e681..116fa556 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime @@ -218,6 +219,9 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id(): ) active_wait_span = plugin._get_span(operation_id) assert active_wait_span is not None + invocation_span = plugin._get_span(None) + assert invocation_span is not None + assert invocation_span.start_time <= active_wait_span.start_time assert ( active_wait_span.attributes["durable.operation.status"] == OperationStatus.STARTED.value @@ -250,6 +254,12 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id(): EXECUTION_ARN, operation_id ) assert wait_span.parent.span_id == invocation_span.context.span_id + assert ( + invocation_span.start_time + <= wait_span.start_time + <= wait_span.end_time + <= invocation_span.end_time + ) assert wait_span.attributes["durable.operation.id"] == operation_id assert wait_span.attributes["durable.operation.type"] == OperationType.WAIT.value assert ( @@ -297,10 +307,13 @@ def test_operation_end_without_start_emits_continuation_span_with_link(): ) -def test_continuation_span_uses_recorded_start_and_end_times(): - """Continuation spans use the recorded operation start/end times.""" +def test_continuation_span_uses_current_start_and_end_times(): + """Continuation spans use current times within the invocation.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + before_callback = time.time_ns() plugin.on_operation_end( OperationEndInfo( @@ -316,59 +329,11 @@ def test_continuation_span_uses_recorded_start_and_end_times(): error=None, ) ) + after_callback = time.time_ns() span = exporter.get_finished_spans()[0] - expected_start = int(START_TIME.timestamp() * 1_000_000_000) - expected_end = int(END_TIME.timestamp() * 1_000_000_000) - assert span.start_time == expected_start - assert span.end_time == expected_end - # Duration must be non-negative. - assert span.end_time >= span.start_time - - -def test_replayed_operation_start_emits_continuation_span_with_link(): - """Replayed operation spans should not reuse the original deterministic span ID.""" - plugin, exporter = _create_plugin() - plugin.on_invocation_start(_invocation_start_info()) - operation_id = "wait-replayed" - random_span_id = int("abcdef1234567890", 16) - plugin._id_generator._fallback_id_generator.generate_span_id = lambda: ( - random_span_id - ) - - plugin.on_operation_start( - OperationStartInfo( - operation_id=operation_id, - operation_type=OperationType.WAIT, - sub_type=OperationSubType.WAIT, - name="replayed-wait", - parent_id=None, - start_time=START_TIME, - is_replayed=True, - status=OperationStatus.SUCCEEDED, - ) - ) - plugin.on_operation_end( - OperationEndInfo( - operation_id=operation_id, - operation_type=OperationType.WAIT, - sub_type=OperationSubType.WAIT, - name="replayed-wait", - parent_id=None, - start_time=START_TIME, - is_replayed=True, - status=OperationStatus.SUCCEEDED, - end_time=END_TIME, - error=None, - ) - ) - - span = exporter.get_finished_spans()[0] - assert span.name == "replayed-wait" - assert span.context.span_id == random_span_id - assert span.links[0].context.span_id == operation_id_to_span_id( - EXECUTION_ARN, operation_id - ) + assert invocation_span.start_time <= span.start_time + assert before_callback <= span.start_time <= span.end_time <= after_callback def test_retried_operation_start_emits_continuation_span_with_link(): 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 5294f1c7..c2d24750 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 @@ -213,9 +213,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def on_operation_start(self, info: OperationStartInfo) -> None: """ Called before an operation's START checkpoint is queued, or when a - prior operation is replayed. This guarantees that it strictly precedes - ``on_user_function_start``. This is called NOT within the thread that - runs operation. + prior non-terminal operation is replayed. This guarantees that it + strictly precedes ``on_user_function_start``. This is called NOT within + the thread that runs operation. Args: info: Information about the operation. @@ -225,8 +225,9 @@ def on_operation_start(self, info: OperationStartInfo) -> None: def on_operation_end(self, info: OperationEndInfo) -> None: """ - Called when an operation checkpoints a terminal status, or when a prior - terminal operation is replayed. This is called NOT within the thread that runs operation. + Called when an operation checkpoints a terminal status. Terminal + operations are not emitted again during replay. This is called NOT + within the thread that runs operation. Args: info: Information about the operation. @@ -407,7 +408,10 @@ def on_operation_action( ) def on_operation_replay(self, operation: Operation) -> None: - """Execute plugins for a checkpointed operation observed during replay.""" + """Execute plugins for a non-terminal operation observed during replay.""" + if self._is_terminal_status(operation.status): + return + start_info = OperationStartInfo( operation_id=operation.operation_id, operation_type=operation.operation_type, @@ -420,29 +424,6 @@ def on_operation_replay(self, operation: Operation) -> None: ) self.execute_plugins(start_info, sync=True) - if self._is_terminal_status(operation.status): - self.execute_plugins( - OperationEndInfo( - operation_id=operation.operation_id, - operation_type=operation.operation_type, - sub_type=operation.sub_type, - name=operation.name, - parent_id=operation.parent_id, - start_time=operation.start_timestamp, - end_time=operation.end_timestamp, - result=_extract_result(operation), - status=operation.status, - error=self._extract_error(operation), - attempt=( - operation.step_details.attempt - if operation.step_details - else None - ), - is_replayed=True, - ), - sync=True, - ) - def on_operation_update( self, operation_or_operations: Operation | Sequence[Operation] | None, 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 b1b9a6a1..0dff5f5b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -642,6 +642,48 @@ def test_fail_action_does_not_fire(self): self.assertEqual(self.plugin.calls, []) +class TestPluginExecutorOnOperationReplay(unittest.TestCase): + """Tests for PluginExecutor.on_operation_replay.""" + + def test_terminal_operation_does_not_fire_callbacks(self): + terminal_statuses = [ + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.CANCELLED, + OperationStatus.STOPPED, + ] + + for status in terminal_statuses: + with self.subTest(status=status): + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin]) + operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + status=status, + ) + + with executor.run(): + executor.on_operation_replay(operation) + + self.assertEqual(plugin.calls, []) + + def test_non_terminal_operation_fires_operation_start(self): + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin]) + operation = Operation( + operation_id="op-1", + operation_type=OperationType.WAIT, + status=OperationStatus.STARTED, + ) + + with executor.run(): + executor.on_operation_replay(operation) + + self.assertEqual(plugin.calls, ["operation_start:op-1"]) + + class TestPluginExecutorOnOperationUpdate(unittest.TestCase): """Tests for PluginExecutor.on_operation_update.""" 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 f0a576e4..94d6d56f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -4921,14 +4921,26 @@ def test_plugin_executor_not_called_for_pending_operations(): assert len(operation_end_calls) == 0 -def test_emit_operation_replay_hook_fires_start_and_end_for_terminal_operation(): - """emit_operation_replay_hook emits start+end (is_replayed=True) for terminal ops.""" +@pytest.mark.parametrize( + "terminal_status", + [ + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.CANCELLED, + OperationStatus.STOPPED, + ], +) +def test_emit_operation_replay_hook_skips_terminal_operation( + terminal_status: OperationStatus, +): + """Terminal operations do not emit plugin callbacks during replay.""" start_time = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) end_time = datetime.datetime(2025, 1, 2, tzinfo=datetime.UTC) operation = Operation( operation_id="step-1", operation_type=OperationType.STEP, - status=OperationStatus.SUCCEEDED, + status=terminal_status, parent_id="parent-1", name="my-step", start_timestamp=start_time, @@ -4959,10 +4971,7 @@ def on_operation_end(self, info): state.emit_operation_replay_hook(operation) state.emit_operation_replay_hook(operation) - assert captured == [ - ("start", "step-1", True, OperationStatus.SUCCEEDED), - ("end", "step-1", True, OperationStatus.SUCCEEDED), - ] + assert captured == [] def test_emit_operation_replay_hook_fires_only_start_for_non_terminal_operation():