diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index f10c4bda..3f7c10fb 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -167,9 +167,25 @@ def _pop_span(self, key: str) -> Span | None: return self._operation_spans.pop(key, None) @staticmethod - def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: + def _attempt_key( + info: UserFunctionStartInfo | UserFunctionEndInfo, + ) -> str: return f"{info.operation_id}:attempt:{info.attempt or 1}" + @classmethod + def _user_function_key( + cls, + info: UserFunctionStartInfo | UserFunctionEndInfo, + ) -> str: + """Return the registry key a user function's span and scope are stored under. + + STEP user functions are attempts, so each attempt gets its own key; a + CONTEXT is entered once per invocation and uses the operation id. + """ + if info.operation_type is OperationType.STEP: + return cls._attempt_key(info) + return info.operation_id + # ------------------------------------------------------------------ # Context scope helpers # ------------------------------------------------------------------ @@ -511,12 +527,12 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: raise RuntimeError( "on_user_function_start only supports CONTEXT and STEP operations" ) + key = self._user_function_key(info) if info.operation_type is OperationType.STEP: parent = self._get_span(info.operation_id) or self._resolve_parent( info.parent_id ) name = f"{info.name or info.operation_id} attempt {info.attempt or 1}" - key = self._attempt_key(info) span = self._start_span( operation_id=info.operation_id, name=name, @@ -528,7 +544,6 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: ) else: # CONTEXT parent = self._resolve_parent(info.parent_id) - key = info.operation_id span = self._start_span( operation_id=info.operation_id, name=info.name or info.operation_id, @@ -548,17 +563,16 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end only supports CONTEXT and STEP operations" ) - key = ( - self._attempt_key(info) - if info.operation_type is OperationType.STEP - else info.operation_id - ) + key = self._user_function_key(info) span = self._get_span(key) if span is None: raise RuntimeError( "on_user_function_end without matching on_user_function_start" ) - if info.operation_type is OperationType.STEP: + if ( + info.operation_type is OperationType.STEP + and info.outcome is not UserFunctionOutcome.INCOMPLETE + ): span.set_attributes(self._operation_attributes(info)) if info.outcome is UserFunctionOutcome.FAILED: span.set_status( 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 182631b4..3b95a5d5 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 @@ -168,10 +168,26 @@ def _get_span(self, operation_id: str | None) -> Span | None: return self._operation_spans.get(operation_id) @staticmethod - def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: + def _attempt_span_key( + info: UserFunctionStartInfo | UserFunctionEndInfo, + ) -> str: """Return the registry key for a STEP attempt span.""" return f"{info.operation_id}:attempt:{info.attempt or 1}" + @classmethod + def _user_function_span_key( + cls, + info: UserFunctionStartInfo | UserFunctionEndInfo, + ) -> str: + """Return the registry key a user function's span and scope are stored under. + + STEP user functions are attempts, so each attempt gets its own key; a + CONTEXT is entered once per invocation and uses the operation id. + """ + if info.operation_type is OperationType.STEP: + return cls._attempt_span_key(info) + return info.operation_id + # ------------------------------------------------------------------ # Context scope helpers # ------------------------------------------------------------------ @@ -613,11 +629,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: span_name = info.name or info.operation_id if info.operation_type is OperationType.STEP: span_name = f"{span_name} attempt {info.attempt or 1}" - span_key = ( - self._attempt_span_key(info) - if info.operation_type is OperationType.STEP - else info.operation_id - ) + span_key = self._user_function_span_key(info) span = self._start_span( operation_id=info.operation_id, name=span_name, @@ -648,19 +660,17 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end should only be called for CONTEXT and STEP operations" ) - # key = f"{info.operation_id}-{int(info.start_time.timestamp())}" - span_key = ( - self._attempt_span_key(info) - if info.operation_type is OperationType.STEP - else info.operation_id - ) + span_key = self._user_function_span_key(info) span = self._get_span(span_key) if not span: raise RuntimeError( "on_user_function_end called without matching on_user_function_start" ) - if info.operation_type is OperationType.STEP: + if ( + info.operation_type is OperationType.STEP + and info.outcome is not UserFunctionOutcome.INCOMPLETE + ): span.set_attributes(self._extract_attributes(info)) if info.outcome is UserFunctionOutcome.FAILED: span.set_status( 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 63548eea..dc4a699c 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 @@ -613,6 +613,48 @@ def _step_end_info( ) +def _step_incomplete_info( + operation_id: str, + parent_id: str | None = None, + attempt: int = 1, +) -> UserFunctionEndInfo: + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=attempt, + outcome=UserFunctionOutcome.INCOMPLETE, + end_time=END_TIME, + error=None, + ) + + +def _context_incomplete_info( + operation_id: str, parent_id: str | None = None +) -> UserFunctionEndInfo: + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + outcome=UserFunctionOutcome.INCOMPLETE, + end_time=END_TIME, + error=None, + ) + + def _context_start_info( operation_id: str, parent_id: str | None = None ) -> UserFunctionStartInfo: @@ -882,75 +924,85 @@ def test_reentered_step_attempt_releases_the_previous_scope(): assert plugin._context_tokens == {} -def test_reentry_on_another_thread_leaves_the_originating_worker_dirty(): - """Pin what re-entry can and cannot clean up across threads. +def test_suspension_releases_the_scope_on_the_originating_worker(): + """Verify the suspending worker releases its own scope. - A resumed branch can land on a different pool thread than the one that - suspended. Re-entry drops the foreign token instead of resetting it, because - a context token can only be reset on its own thread, and it unwinds cleanly - on the thread that re-entered. The worker that suspended keeps the abandoned - span current: releasing it needs a hook invoked on that thread when the user - function fails to complete, which the SDK does not provide. The worker is - kept alive here so this limitation is asserted rather than hidden by pool - shutdown; the assertion flips once such a hook exists. + A suspended user function reports no outcome, so the SDK fires + on_user_function_end with INCOMPLETE on the thread that ran it -- the only thread that + can reset its context token. The worker is kept alive and probed to prove it + is left clean even though the resume lands on a different thread. """ - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) before_context = otel_context.get_current() span_key = "step-1:attempt:1" with ThreadPoolExecutor(max_workers=1) as worker: - # The suspending run happens on the worker and never reports an end. - worker.submit( - plugin.on_user_function_start, _step_start_info("step-1") - ).result() - abandoned_span = plugin._get_span(span_key) - assert abandoned_span is not None - foreign_thread_ident, _foreign_token = plugin._context_tokens[span_key] - assert foreign_thread_ident != threading.get_ident() - # The timed resume lands on this thread instead. + def suspend_on_worker() -> tuple[int, bool]: + plugin.on_user_function_start(_step_start_info("step-1")) + attached_span_id = trace.get_current_span().get_span_context().span_id + plugin.on_user_function_end(_step_incomplete_info("step-1")) + return ( + attached_span_id, + trace.get_current_span().get_span_context().is_valid, + ) + + attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + suspended_span = plugin._get_span(span_key) + + # The scope was released on the worker, and its span is left open. + assert attached_span_id != 0 + assert span_still_current is False + assert span_key not in plugin._context_tokens + assert suspended_span is not None + assert not exporter.get_finished_spans() + + # The timed resume lands on this thread, with nothing stale to unwind. plugin.on_user_function_start(_step_start_info("step-1")) assert plugin._context_tokens[span_key][0] == threading.get_ident() plugin.on_user_function_end(_step_end_info("step-1")) - - # This thread unwound to where it started. assert otel_context.get_current() == before_context - # The originating worker is still carrying the abandoned span. - worker_span_id = worker.submit( - lambda: trace.get_current_span().get_span_context().span_id + # The originating worker is still clean. + worker_span_valid = worker.submit( + lambda: trace.get_current_span().get_span_context().is_valid ).result() - assert worker_span_id == abandoned_span.get_span_context().span_id + assert worker_span_valid is False plugin.on_invocation_end(_invocation_end_info()) -def test_nested_reentry_restores_the_abandoned_outer_scope(): - """Pin nested re-entry: correct ids, but the abandoned outer span object. +def test_nested_suspension_unwinds_scopes_in_reverse_order(): + """Verify nested suspends release inner-first and resume without stale scopes. - When an outer child context and an inner one both suspend, re-entry releases - each scope in the order the operations are replayed, which is not the reverse - of the order they were attached. Ending the inner operation therefore - restores the scope captured for the abandoned outer span rather than the - resumed one. Deterministic CONTEXT span ids make the two indistinguishable - downstream -- same trace id and span id, so parenting and log correlation are - unaffected -- but the current span object is one that is never exported, so - anything an instrumentation library records on it is lost. Reverse-order - unwinding needs the SDK to report the suspension; this test documents the - current behaviour and flips when that lands. + The INCOMPLETE end callback fires as the exception propagates outward, so the inner + context's scope is released before its enclosing one. On resume, ending the + inner operation restores the resumed outer scope rather than the one captured + for the suspended run. """ - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) before_context = otel_context.get_current() - # Both contexts suspend, so neither reports an end. plugin.on_user_function_start(_context_start_info("ctx-outer")) - abandoned_outer = plugin._get_span("ctx-outer") + suspended_outer = plugin._get_span("ctx-outer") plugin.on_user_function_start( _context_start_info("ctx-inner", parent_id="ctx-outer") ) - assert abandoned_outer is not None + assert suspended_outer is not None + + # Both contexts suspend: the inner one unwinds first. + plugin.on_user_function_end( + _context_incomplete_info("ctx-inner", parent_id="ctx-outer") + ) + assert trace.get_current_span() is suspended_outer + + plugin.on_user_function_end(_context_incomplete_info("ctx-outer")) + assert otel_context.get_current() == before_context + assert set(plugin._context_tokens) == {"__invocation_context__"} + # Neither span is ended: both operations are still in flight. + assert not exporter.get_finished_spans() # The timed in-process resume replays both contexts, outer first. plugin.on_user_function_start(_context_start_info("ctx-outer")) @@ -960,27 +1012,14 @@ def test_nested_reentry_restores_the_abandoned_outer_scope(): ) resumed_inner = plugin._get_span("ctx-inner") assert resumed_outer is not None - assert resumed_inner is not None - assert resumed_outer is not abandoned_outer - - # Resumed inner code runs under the resumed inner span. + assert resumed_outer is not suspended_outer assert trace.get_current_span() is resumed_inner plugin.on_user_function_end(_context_end_info("ctx-inner", parent_id="ctx-outer")) - # The restored scope carries the abandoned outer span, whose ids match the - # resumed one because CONTEXT span ids are derived from the operation id. - assert trace.get_current_span() is abandoned_outer - assert ( - abandoned_outer.get_span_context().span_id - == resumed_outer.get_span_context().span_id - ) - assert ( - abandoned_outer.get_span_context().trace_id - == resumed_outer.get_span_context().trace_id - ) + # The resumed outer scope is restored, not the one from the suspended run. + assert trace.get_current_span() is resumed_outer - # Leaving the outer context still unwinds to where the invocation started. plugin.on_user_function_end(_context_end_info("ctx-outer")) assert otel_context.get_current() == before_context assert set(plugin._context_tokens) == {"__invocation_context__"} 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 c7901b5d..0f3bb62d 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 @@ -147,6 +147,30 @@ def _user_function_end_info( ) +def _user_function_incomplete_info( + operation_id: str, + attempt: int = 1, + parent_id: str | None = None, + operation_type: OperationType = OperationType.STEP, +) -> UserFunctionEndInfo: + """Create user function end info for an incomplete execution.""" + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=operation_type, + sub_type=None, + name=f"step-{operation_id}", + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=attempt, + outcome=UserFunctionOutcome.INCOMPLETE, + end_time=END_TIME, + error=None, + ) + + def test_extract_attributes_uses_structural_event_attributes(): plugin, _ = _create_plugin() @@ -1400,80 +1424,96 @@ def test_reentered_step_attempt_releases_the_previous_scope(): plugin.on_invocation_end(_invocation_end_info()) -def test_reentry_on_another_thread_leaves_the_originating_worker_dirty(): - """Pin what re-entry can and cannot clean up across threads. +def test_suspension_releases_the_scope_on_the_originating_worker(): + """Verify the suspending worker releases its own scope. - A resumed branch can land on a different pool thread than the one that - suspended. Re-entry drops the foreign token instead of resetting it, because - a context token can only be reset on its own thread, and it unwinds cleanly - on the thread that re-entered. The worker that suspended keeps the abandoned - span current: releasing it needs a hook invoked on that thread when the user - function fails to complete, which the SDK does not provide. The worker is - kept alive here so this limitation is asserted rather than hidden by pool - shutdown; the assertion flips once such a hook exists. + A suspended user function reports no outcome, so the SDK fires + on_user_function_end with INCOMPLETE on the thread that ran it -- the only thread that + can reset its context token. The worker is kept alive and probed to prove it + is left clean even though the resume lands on a different thread. """ - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) before_context = otel_context.get_current() operation_id = "step-1" span_key = "step-1:attempt:1" with ThreadPoolExecutor(max_workers=1) as worker: - # The suspending run happens on the worker and never reports an end. - worker.submit( - plugin.on_user_function_start, _user_function_start_info(operation_id) - ).result() - abandoned_span = plugin._get_span(span_key) - assert abandoned_span is not None - foreign_thread_ident, _foreign_token = plugin._context_tokens[span_key] - assert foreign_thread_ident != threading.get_ident() - # The timed resume lands on this thread instead. + def suspend_on_worker() -> tuple[int, bool]: + plugin.on_user_function_start(_user_function_start_info(operation_id)) + attached_span_id = trace.get_current_span().get_span_context().span_id + plugin.on_user_function_end(_user_function_incomplete_info(operation_id)) + return ( + attached_span_id, + trace.get_current_span().get_span_context().is_valid, + ) + + attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + suspended_span = plugin._get_span(span_key) + + # The scope was released on the worker, and its span is left open. + assert attached_span_id != 0 + assert span_still_current is False + assert span_key not in plugin._context_tokens + assert suspended_span is not None + assert not exporter.get_finished_spans() + + # The timed resume lands on this thread, with nothing stale to unwind. plugin.on_user_function_start(_user_function_start_info(operation_id)) assert plugin._context_tokens[span_key][0] == threading.get_ident() plugin.on_user_function_end(_user_function_end_info(operation_id)) - - # This thread unwound to where it started. assert otel_context.get_current() == before_context - # The originating worker is still carrying the abandoned span. - worker_span_id = worker.submit( - lambda: trace.get_current_span().get_span_context().span_id + # The originating worker is still clean. + worker_span_valid = worker.submit( + lambda: trace.get_current_span().get_span_context().is_valid ).result() - assert worker_span_id == abandoned_span.get_span_context().span_id + assert worker_span_valid is False plugin.on_invocation_end(_invocation_end_info()) -def test_nested_reentry_restores_the_abandoned_outer_scope(): - """Pin nested re-entry: correct ids, but the abandoned outer span object. +def test_nested_suspension_unwinds_scopes_in_reverse_order(): + """Verify nested suspends release inner-first and resume without stale scopes. - When an outer child context and an inner one both suspend, re-entry releases - each scope in the order the operations are replayed, which is not the reverse - of the order they were attached. Ending the inner operation therefore - restores the scope captured for the abandoned outer span rather than the - resumed one. Deterministic CONTEXT span ids make the two indistinguishable - downstream -- same trace id and span id, so parenting and log correlation are - unaffected -- but the current span object is one that is never exported, so - anything an instrumentation library records on it is lost. Reverse-order - unwinding needs the SDK to report the suspension; this test documents the - current behaviour and flips when that lands. + The INCOMPLETE end callback fires as the exception propagates outward, so the inner + context's scope is released before its enclosing one. On resume, ending the + inner operation restores the resumed outer scope rather than the one captured + for the suspended run. """ - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) before_context = otel_context.get_current() - # Both contexts suspend, so neither reports an end. plugin.on_user_function_start( _user_function_start_info("ctx-outer", operation_type=OperationType.CONTEXT) ) - abandoned_outer = plugin._get_span("ctx-outer") + suspended_outer = plugin._get_span("ctx-outer") plugin.on_user_function_start( _user_function_start_info( "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT ) ) - assert abandoned_outer is not None + assert suspended_outer is not None + + # Both contexts suspend: the inner one unwinds first. + plugin.on_user_function_end( + _user_function_incomplete_info( + "ctx-inner", parent_id="ctx-outer", operation_type=OperationType.CONTEXT + ) + ) + assert trace.get_current_span() is suspended_outer + + plugin.on_user_function_end( + _user_function_incomplete_info( + "ctx-outer", operation_type=OperationType.CONTEXT + ) + ) + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + # Neither span is ended: both operations are still in flight. + assert not exporter.get_finished_spans() # The timed in-process resume replays both contexts, outer first. plugin.on_user_function_start( @@ -1487,10 +1527,7 @@ def test_nested_reentry_restores_the_abandoned_outer_scope(): ) resumed_inner = plugin._get_span("ctx-inner") assert resumed_outer is not None - assert resumed_inner is not None - assert resumed_outer is not abandoned_outer - - # Resumed inner code runs under the resumed inner span. + assert resumed_outer is not suspended_outer assert trace.get_current_span() is resumed_inner plugin.on_user_function_end( @@ -1499,19 +1536,9 @@ def test_nested_reentry_restores_the_abandoned_outer_scope(): ) ) - # The restored scope carries the abandoned outer span, whose ids match the - # resumed one because CONTEXT span ids are derived from the operation id. - assert trace.get_current_span() is abandoned_outer - assert ( - abandoned_outer.get_span_context().span_id - == resumed_outer.get_span_context().span_id - ) - assert ( - abandoned_outer.get_span_context().trace_id - == resumed_outer.get_span_context().trace_id - ) + # The resumed outer scope is restored, not the one from the suspended run. + assert trace.get_current_span() is resumed_outer - # Leaving the outer context still unwinds to where the invocation started. plugin.on_user_function_end( _user_function_end_info("ctx-outer", operation_type=OperationType.CONTEXT) ) 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 daab80cc..e9549d30 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 @@ -191,6 +191,7 @@ class OperationChangeInfo: class UserFunctionOutcome(Enum): SUCCEEDED = "SUCCEEDED" FAILED = "FAILED" + INCOMPLETE = "INCOMPLETE" @classmethod def from_error(cls, error: ErrorObject | None) -> UserFunctionOutcome: @@ -215,7 +216,11 @@ class UserFunctionEndInfo(OperationInfo): @classmethod def from_start_info( - cls, start_info: UserFunctionStartInfo, error: ErrorObject | None + cls, + start_info: UserFunctionStartInfo, + error: ErrorObject | None, + *, + outcome: UserFunctionOutcome | None = None, ) -> UserFunctionEndInfo: return UserFunctionEndInfo( operation_id=start_info.operation_id, @@ -228,7 +233,11 @@ def from_start_info( status=start_info.status, is_replay_children=start_info.is_replay_children, attempt=start_info.attempt, - outcome=UserFunctionOutcome.from_error(error), + outcome=( + outcome + if outcome is not None + else UserFunctionOutcome.from_error(error) + ), end_time=datetime.datetime.now(datetime.UTC), error=error, ) @@ -650,10 +659,17 @@ def on_user_function_start( self.execute_plugins(start_info, sync=True) return start_info - def on_user_function_end(self, start_info: UserFunctionStartInfo, error) -> None: - """Execute any registered plugins for the operation when its user function finishes execution.""" + def on_user_function_end( + self, + start_info: UserFunctionStartInfo, + error, + *, + outcome: UserFunctionOutcome | None = None, + ) -> None: + """Execute plugins when a user function returns, fails, or is incomplete.""" self.execute_plugins( - UserFunctionEndInfo.from_start_info(start_info, error), sync=True + UserFunctionEndInfo.from_start_info(start_info, error, outcome=outcome), + sync=True, ) def on_operation_action( 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 26aefbe3..9f2391ce 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 @@ -21,7 +21,6 @@ DurableOperationError, GetExecutionStateError, OrphanedChildException, - SuspendExecution, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -37,6 +36,7 @@ ) from aws_durable_execution_sdk_python.plugin import ( PluginExecutor, + UserFunctionOutcome, ) from aws_durable_execution_sdk_python.threading import CompletionEvent @@ -1165,16 +1165,22 @@ def wrapper(*args, **kwargs): start_info = self._plugin_executor.on_user_function_start( operation_identifier, is_replay_children, attempt ) + outcome = UserFunctionOutcome.INCOMPLETE + error = None try: result = user_function(*args, **kwargs) - self._plugin_executor.on_user_function_end(start_info, None) - return result - except SuspendExecution: + except Exception as exception: + outcome = UserFunctionOutcome.FAILED + error = ErrorObject.from_exception(exception) raise - except Exception as e: + else: + outcome = UserFunctionOutcome.SUCCEEDED + return result + finally: + # Runs on the thread that executed the user function, which is + # the only thread that can release state bound to it. self._plugin_executor.on_user_function_end( - start_info, ErrorObject.from_exception(e) + start_info, error, outcome=outcome ) - raise return wrapper diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py new file mode 100644 index 00000000..7b30828d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_user_function_lifecycle_int_test.py @@ -0,0 +1,229 @@ +"""End-to-end coverage for user-function plugin lifecycle callbacks.""" + +from __future__ import annotations + +import dataclasses +import threading +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any +from unittest.mock import Mock, patch + +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 ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + UserFunctionEndInfo, + UserFunctionOutcome, + UserFunctionStartInfo, +) + + +@dataclass(frozen=True) +class _LifecycleEvent: + phase: str + operation_id: str + name: str | None + outcome: UserFunctionOutcome | None + thread_id: int + + +class _LifecycleRecordingPlugin(DurableInstrumentationPlugin): + def __init__(self) -> None: + self.events: list[_LifecycleEvent] = [] + + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + self.events.append( + _LifecycleEvent( + phase="start", + operation_id=info.operation_id, + name=info.name, + outcome=None, + thread_id=threading.get_ident(), + ) + ) + + def on_user_function_end(self, info: UserFunctionEndInfo) -> None: + self.events.append( + _LifecycleEvent( + phase="end", + operation_id=info.operation_id, + name=info.name, + outcome=info.outcome, + thread_id=threading.get_ident(), + ) + ) + + +def _lambda_context() -> Mock: + context = Mock() + context.aws_request_id = "test-request-id" + context.client_context = None + context.identity = None + context._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + context.invoked_function_arn = "test-arn" + context.tenant_id = None + return context + + +def _event( + extra_operations: Sequence[Mapping[str, Any]] | None = None, + updated_operation_ids: list[str] | None = None, +) -> dict[str, Any]: + event: dict[str, Any] = { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": OperationType.EXECUTION.value, + "Status": OperationStatus.STARTED.value, + "ExecutionDetails": {"InputPayload": "{}"}, + }, + *(extra_operations or []), + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + if updated_operation_ids is not None: + event["UpdatedOperationIds"] = updated_operation_ids + return event + + +def _tracking_checkpoint( + initial_operations: list[Operation] | None = None, +) -> tuple[Any, list[Operation]]: + operations = list(initial_operations or []) + if not operations: + operations.append( + Operation( + operation_id="execution-1", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + ) + ) + + def checkpoint( + durable_execution_arn, # noqa: ARG001 + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + for update in updates: + operations.append( + Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + ) + ) + return CheckpointOutput( + checkpoint_token="new-token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState( + operations=operations.copy() + ), + ) + + return checkpoint, operations + + +def test_child_user_function_lifecycle_across_suspend_and_replay() -> None: + """A child reports INCOMPLETE on suspend and SUCCEEDED after replay.""" + plugin = _LifecycleRecordingPlugin() + child_threads: list[int] = [] + + def child_function(context: DurableContext) -> str: + child_threads.append(threading.get_ident()) + context.wait(Duration.from_seconds(60)) + return "charged" + + def user_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + return context.run_in_child_context(child_function, name="charge") + + handler = durable_execution(user_handler, plugins=[plugin]) + + first_checkpoint, first_operations = _tracking_checkpoint() + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = first_checkpoint + mock_client_class.initialize_client.return_value = mock_client + first_result = handler(_event(), _lambda_context()) + + assert first_result["Status"] == InvocationStatus.PENDING.value + assert [(event.phase, event.name, event.outcome) for event in plugin.events] == [ + ("start", "charge", None), + ("end", "charge", UserFunctionOutcome.INCOMPLETE), + ] + first_start, first_end = plugin.events + assert first_start.operation_id == first_end.operation_id + assert first_start.thread_id == first_end.thread_id == child_threads[0] + + child_operation = next( + operation + for operation in first_operations + if operation.operation_type is OperationType.CONTEXT + ) + wait_operation = next( + operation + for operation in first_operations + if operation.operation_type is OperationType.WAIT + ) + assert wait_operation.parent_id == child_operation.operation_id + + replay_operations = [ + dataclasses.replace(operation, status=OperationStatus.SUCCEEDED) + if operation.operation_type is OperationType.WAIT + else operation + for operation in first_operations + ] + replay_event_operations = [ + operation.to_dict() + for operation in replay_operations + if operation.operation_type is not OperationType.EXECUTION + ] + replay_checkpoint, _ = _tracking_checkpoint(replay_operations) + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = replay_checkpoint + mock_client_class.initialize_client.return_value = mock_client + replay_result = handler( + _event( + extra_operations=replay_event_operations, + updated_operation_ids=[wait_operation.operation_id], + ), + _lambda_context(), + ) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + assert [(event.phase, event.name, event.outcome) for event in plugin.events] == [ + ("start", "charge", None), + ("end", "charge", UserFunctionOutcome.INCOMPLETE), + ("start", "charge", None), + ("end", "charge", UserFunctionOutcome.SUCCEEDED), + ] + replay_start, replay_end = plugin.events[2:] + assert replay_start.operation_id == first_start.operation_id + assert replay_end.operation_id == first_start.operation_id + assert replay_start.thread_id == replay_end.thread_id == child_threads[1] + assert len(child_threads) == 2 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 b693ab10..69cdfc50 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1841,7 +1841,7 @@ class TestUserFunctionOutcomeValues(unittest.TestCase): def test_outcome_values(self): self.assertEqual( {o.value for o in UserFunctionOutcome}, - {"SUCCEEDED", "FAILED"}, + {"SUCCEEDED", "FAILED", "INCOMPLETE"}, ) @@ -1856,6 +1856,16 @@ def test_error_is_failed(self): UserFunctionOutcome.from_error(ERROR), UserFunctionOutcome.FAILED ) + def test_explicit_incomplete_outcome_is_preserved(self): + info = UserFunctionEndInfo.from_start_info( + USER_FUNCTION_START_INFO, + None, + outcome=UserFunctionOutcome.INCOMPLETE, + ) + + self.assertEqual(info.outcome, UserFunctionOutcome.INCOMPLETE) + self.assertIsNone(info.error) + # endregion Suspend Outcome Tests 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 94d6d56f..7c33d3b0 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -20,6 +20,7 @@ GetExecutionStateError, OrphanedChildException, StepError, + SuspendExecution, TimedSuspendExecution, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -45,6 +46,7 @@ OperationStartInfo, PluginExecutor, UserFunctionEndInfo, + UserFunctionOutcome, ) from aws_durable_execution_sdk_python.state import ( CheckpointBatcherConfig, @@ -4302,6 +4304,7 @@ class _RecordingPlugin(DurableInstrumentationPlugin): def __init__(self) -> None: self.calls: list[str] = [] self.operation_starts: list[OperationStartInfo] = [] + self.user_function_ends: list[UserFunctionEndInfo] = [] def on_execution_start(self, info): self.calls.append("execution_start") @@ -4332,6 +4335,7 @@ def on_user_function_start(self, info): def on_user_function_end(self, info): self.calls.append(f"user_function_end:{info.operation_id}") + self.user_function_ends.append(info) def test_execution_state_accepts_plugin_executor_parameter(): @@ -4821,15 +4825,12 @@ def on_operation_end(self, info): executor.shutdown(wait=True) -def test_wrap_user_function_suspend_does_not_fire_end_hook(): - """A user function that suspends does not fire the end hook. +def test_wrap_user_function_suspend_fires_incomplete_end_hook(): + """A timed suspend reports an incomplete outcome through the end hook. - Regression: a timed suspend (TimedSuspendExecution) raised inside a wrapped - user function (e.g. a child context that waits) must not be surfaced to - plugins as a FAILED outcome. The suspend is normal durable control flow, - and the plugin observes it by absence (no end hook fires), with the - instrumentation plugin's own per-invocation span sweep closing any open - spans cleanly at invocation end. + Suspension is normal durable control flow rather than a user failure. The + callback still runs so plugins can release state bound to the user-function + thread, but it carries no error and must not be interpreted as completion. """ captured: list[UserFunctionEndInfo] = [] @@ -4858,7 +4859,9 @@ def suspends(_: object) -> None: with pytest.raises(TimedSuspendExecution): wrapped(None) - assert captured == [] + assert len(captured) == 1 + assert captured[0].outcome is UserFunctionOutcome.INCOMPLETE + assert captured[0].error is None def test_plugin_executor_not_called_for_pending_operations(): @@ -5124,3 +5127,134 @@ def reader(): writer_t.join(timeout=5) assert not errors, f"has_prior_operations raced with concurrent update: {errors}" + + +# region wrap_user_function incomplete notification + + +def _wrapping_state(plugin: _RecordingPlugin) -> ExecutionState: + """Build an ExecutionState whose plugin executor is running.""" + return ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=Mock(spec=LambdaClient), + plugin_executor=PluginExecutor(plugins=[plugin]), + ) + + +def _wrapped(state: ExecutionState, user_function): + return state.wrap_user_function( + user_function, + OperationIdentifier( + operation_id="step-1", + sub_type=OperationSubType.STEP, + name="fetch-user", + ), + attempt=1, + ) + + +@pytest.mark.parametrize( + "raised", + [ + SuspendExecution("suspended"), + TimedSuspendExecution("suspended until", 1.0), + OrphanedChildException("parent already completed", "step-1"), + BackgroundThreadError("checkpoint failed", RuntimeError("boom")), + SystemExit(1), + ], +) +def test_wrap_user_function_reports_incomplete_when_no_outcome(raised): + """A user function that reports no outcome notifies plugins instead.""" + plugin = _RecordingPlugin() + state = _wrapping_state(plugin) + + def user_function(): + raise raised + + with state._plugin_executor.run(), pytest.raises(type(raised)): + _wrapped(state, user_function)() + + assert plugin.calls == [ + "user_function_start:step-1", + "user_function_end:step-1", + ] + assert [info.outcome for info in plugin.user_function_ends] == [ + UserFunctionOutcome.INCOMPLETE + ] + assert plugin.user_function_ends[0].error is None + + +def test_wrap_user_function_does_not_report_incomplete_on_success(): + """A returning user function reports an end and nothing else.""" + plugin = _RecordingPlugin() + state = _wrapping_state(plugin) + + with state._plugin_executor.run(): + assert _wrapped(state, lambda: "done")() == "done" + + assert plugin.calls == [ + "user_function_start:step-1", + "user_function_end:step-1", + ] + assert [info.outcome for info in plugin.user_function_ends] == [ + UserFunctionOutcome.SUCCEEDED + ] + + +def test_wrap_user_function_does_not_report_incomplete_on_failure(): + """An ordinary exception reports an end, not an incomplete.""" + plugin = _RecordingPlugin() + state = _wrapping_state(plugin) + + def user_function(): + raise ValueError("boom") + + with state._plugin_executor.run(), pytest.raises(ValueError, match="boom"): + _wrapped(state, user_function)() + + assert plugin.calls == [ + "user_function_start:step-1", + "user_function_end:step-1", + ] + assert [info.outcome for info in plugin.user_function_ends] == [ + UserFunctionOutcome.FAILED + ] + + +def test_wrap_user_function_incomplete_runs_on_the_user_function_thread(): + """The notification must arrive on the thread that ran the user function.""" + hook_threads: list[int] = [] + + class _ThreadRecordingPlugin(DurableInstrumentationPlugin): + def on_user_function_end(self, info) -> None: + if info.outcome is UserFunctionOutcome.INCOMPLETE: + hook_threads.append(threading.get_ident()) + + plugin = _ThreadRecordingPlugin() + state = ExecutionState( + durable_execution_arn="test_arn", + initial_checkpoint_token="token123", # noqa: S106 + operations={}, + service_client=Mock(spec=LambdaClient), + plugin_executor=PluginExecutor(plugins=[plugin]), + ) + + def user_function(): + raise SuspendExecution("suspended") + + worker_threads: list[int] = [] + + def run_on_worker() -> None: + worker_threads.append(threading.get_ident()) + with contextlib.suppress(SuspendExecution): + _wrapped(state, user_function)() + + with state._plugin_executor.run(), ThreadPoolExecutor(max_workers=1) as worker: + worker.submit(run_on_worker).result() + + assert hook_threads == worker_threads + + +# endregion