From b59c13563819a239e6a7b19e9a3a4cec8b6d7c25 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 19 Aug 2026 21:16:52 +0000 Subject: [PATCH] feat(plugin): report user functions that do not complete wrap_user_function reported an outcome only when the user function returned or raised an ordinary Exception. SuspendExecution, TimedSuspendExecution, OrphanedChildException, BackgroundThreadError and SystemExit all subclass BaseException and bypassed both handlers, so a plugin that bound state to the user-function thread in on_user_function_start was never told to release it -- and never told on the thread that owns it. - Add on_user_function_incomplete, dispatched synchronously like the start and end hooks, with a UserFunctionIncompleteInfo payload - Fire it from a finally block in wrap_user_function, gated on whether an outcome was already reported, so exactly one of end/incomplete follows every start - Release the OTel context scope from the new hook in both plugins, leaving the span open because the operation may still resume - Document the hook and its threading contract in the core README Nested scopes now unwind in reverse order: the hook fires as the exception propagates outward, so an inner context is released before its enclosing one, and a resume no longer restores a scope captured for the suspended run. A suspended branch also releases its scope on the worker that ran it, which the plugins could not do from another thread. The two limitation tests added with the earlier re-entry mitigation now assert the corrected behaviour. Resolves #658 --- .../execution_plugin.py | 43 ++++- .../invocation_plugin.py | 45 +++++- .../tests/test_execution_plugin.py | 150 +++++++++++------- .../tests/test_invocation_plugin.py | 143 ++++++++++------- .../README.md | 19 +++ .../plugin.py | 70 ++++++++ .../aws_durable_execution_sdk_python/state.py | 12 ++ .../tests/state_test.py | 124 +++++++++++++++ 8 files changed, 477 insertions(+), 129 deletions(-) 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 27511809..1a399168 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 @@ -40,6 +40,7 @@ OperationStartInfo, OperationType, UserFunctionEndInfo, + UserFunctionIncompleteInfo, UserFunctionOutcome, UserFunctionStartInfo, ) @@ -167,9 +168,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 | UserFunctionIncompleteInfo, + ) -> str: return f"{info.operation_id}:attempt:{info.attempt or 1}" + @classmethod + def _user_function_key( + cls, + info: UserFunctionStartInfo | UserFunctionEndInfo | UserFunctionIncompleteInfo, + ) -> 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 # ------------------------------------------------------------------ @@ -544,11 +561,7 @@ 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( @@ -583,6 +596,24 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: # without stacking another scope. self._detach_context(key) + def on_user_function_incomplete(self, info: UserFunctionIncompleteInfo) -> None: + """Release the scope of a user function that reported no outcome. + + A suspended, orphaned or interrupted user function never reaches + ``on_user_function_end``, so its scope is released here instead, on the + thread that attached it. Nested scopes unwind in reverse order because + this fires as the exception propagates outward: the inner operation is + released before its enclosing one. + + The span is left open and registered. The operation is not finished -- + a suspended one resumes and ends later -- so ending it here would export + a span for work that is still in flight. + """ + logger.debug("Durable user function incomplete: %s", info) + if not self._tracing_enabled: + return + self._detach_context(self._user_function_key(info)) + # ------------------------------------------------------------------ # Attributes # ------------------------------------------------------------------ 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 9e31b0df..d953c003 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 @@ -16,6 +16,7 @@ OperationStartInfo, OperationType, UserFunctionEndInfo, + UserFunctionIncompleteInfo, UserFunctionOutcome, UserFunctionStartInfo, ) @@ -168,10 +169,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 | UserFunctionIncompleteInfo, + ) -> 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 | UserFunctionIncompleteInfo, + ) -> 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 # ------------------------------------------------------------------ @@ -642,12 +659,7 @@ 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( @@ -682,6 +694,25 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: # get_current_span_context falls back to the invocation span. self._detach_context(span_key) + def on_user_function_incomplete(self, info: UserFunctionIncompleteInfo) -> None: + """Release the scope of a user function that reported no outcome. + + A suspended, orphaned or interrupted user function never reaches + ``on_user_function_end``, so its scope is released here instead, on the + thread that attached it. Nested scopes unwind in reverse order because + this fires as the exception propagates outward: the inner operation is + released before its enclosing one. + + The span is left open and registered. The operation is not finished -- + a suspended one resumes and ends later -- so ending it here would export + a span for work that is still in flight. Spans still open at invocation + end are closed by ``on_invocation_end``. + """ + logger.debug("Durable user function incomplete: %s", info) + if not self._tracing_enabled: + return + self._detach_context(self._user_function_span_key(info)) + def _extract_attributes(self, info: Any) -> _SpanAttributes: """Extract durable execution fields as OpenTelemetry span attributes. 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 1fb2c9db..ab8e5492 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 @@ -22,6 +22,7 @@ OperationStartInfo, OperationType, UserFunctionEndInfo, + UserFunctionIncompleteInfo, UserFunctionOutcome, UserFunctionStartInfo, ) @@ -595,6 +596,42 @@ def _step_end_info( ) +def _step_incomplete_info( + operation_id: str, + parent_id: str | None = None, + attempt: int = 1, +) -> UserFunctionIncompleteInfo: + return UserFunctionIncompleteInfo( + 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, + ) + + +def _context_incomplete_info( + operation_id: str, parent_id: str | None = None +) -> UserFunctionIncompleteInfo: + return UserFunctionIncompleteInfo( + 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, + ) + + def _context_start_info( operation_id: str, parent_id: str | None = None ) -> UserFunctionStartInfo: @@ -864,75 +901,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_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_incomplete(_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 hook 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_incomplete( + _context_incomplete_info("ctx-inner", parent_id="ctx-outer") + ) + assert trace.get_current_span() is suspended_outer + + plugin.on_user_function_incomplete(_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")) @@ -942,27 +989,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 6e7e6742..d5749d92 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 @@ -23,6 +23,7 @@ OperationStartInfo, OperationType, UserFunctionEndInfo, + UserFunctionIncompleteInfo, UserFunctionOutcome, UserFunctionStartInfo, ) @@ -147,6 +148,27 @@ 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, +) -> UserFunctionIncompleteInfo: + """Create user function incomplete info (suspended, orphaned, interrupted).""" + return UserFunctionIncompleteInfo( + 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, + ) + + def test_extract_attributes_uses_structural_event_attributes(): plugin, _ = _create_plugin() @@ -1382,80 +1404,98 @@ 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_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_incomplete( + _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 hook 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_incomplete( + _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_incomplete( + _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( @@ -1469,10 +1509,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( @@ -1481,19 +1518,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/README.md b/packages/aws-durable-execution-sdk-python/README.md index bf7776f4..150831ee 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -73,6 +73,25 @@ example_audit = "example_audit:AUDIT_PLUGIN_PROVIDER" Set `plugin_api_version` to the literal API version the provider implements. Update it only after verifying the provider against that API version. +### User-function hooks and thread-bound state + +`on_user_function_start`, `on_user_function_end` and +`on_user_function_incomplete` run on the thread that executes the user +function, so a plugin can bind state to that thread. Exactly one of the latter +two follows every start: + +- `on_user_function_end` reports an outcome, whether the function returned or + raised an ordinary exception. +- `on_user_function_incomplete` reports no outcome. It fires when the operation + suspends, when a map or parallel branch is orphaned by its parent, when a + background checkpoint failure cuts the invocation short, or when the + interpreter is exiting. + +An incomplete user function is not a finished operation: a suspended one resumes +later, in this invocation or a subsequent one, and reports a real end then. +Release thread-bound state in `on_user_function_incomplete`, but do not record +the operation as complete. + Provider names must be unique across installed distributions. Missing, ambiguous, incompatible, or invalid providers raise `PluginLoadError` during handler initialization with the provider and distribution details. 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..911e3ce6 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 @@ -234,6 +234,44 @@ def from_start_info( ) +@dataclass(frozen=True) +class UserFunctionIncompleteInfo(OperationInfo): + """Reports a user function that ended without reporting an outcome. + + Emitted when the user function neither returns nor raises an ordinary + ``Exception`` -- it suspended, was orphaned by its parent, was cut short by a + background failure, or the interpreter is exiting. The operation is not + finished: a suspended operation resumes later, in this invocation or a + subsequent one, and reports a real end then. + + Plugins receive this on the thread that ran the user function, which is the + only thread that can release state bound to it (a thread-local scope, for + example). Release such state here, but do not treat the operation as + finished. + """ + + is_replay_children: ( + bool # True if user function is called to replay children (MAP/PARALLEL) + ) + + @classmethod + def from_start_info( + cls, start_info: UserFunctionStartInfo + ) -> UserFunctionIncompleteInfo: + return UserFunctionIncompleteInfo( + operation_id=start_info.operation_id, + operation_type=start_info.operation_type, + sub_type=start_info.sub_type, + name=start_info.name, + parent_id=start_info.parent_id, + start_time=start_info.start_time, + is_replayed=start_info.is_replayed, + status=start_info.status, + is_replay_children=start_info.is_replay_children, + attempt=start_info.attempt, + ) + + @dataclass(frozen=True) class InvocationInfo: request_id: str | None @@ -441,6 +479,25 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: """ pass + def on_user_function_incomplete(self, info: UserFunctionIncompleteInfo) -> None: + """Called when a user function ends without reporting an outcome. + + Fires instead of ``on_user_function_end`` when the user function + suspends, is orphaned by its parent, is cut short by a background + failure, or the interpreter is exiting. Exactly one of the two is called + for every ``on_user_function_start``. This is called within the thread + that runs the user provided function, so it is the only opportunity to + release state bound to that thread. + + The operation is not finished: a suspended operation resumes later and + reports a real end then. Release scoped state here, but do not record the + operation as complete. + + Args: + info: Information about the operation attempt. + """ + pass + @dataclass(frozen=True) class DurableInstrumentationPluginProvider: @@ -493,6 +550,8 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: plugin.on_user_function_start(info) case UserFunctionEndInfo(): plugin.on_user_function_end(info) + case UserFunctionIncompleteInfo(): + plugin.on_user_function_incomplete(info) case _: raise RuntimeError(f"Unknown info type: {type(info)}") except Exception: @@ -656,6 +715,17 @@ def on_user_function_end(self, start_info: UserFunctionStartInfo, error) -> None UserFunctionEndInfo.from_start_info(start_info, error), sync=True ) + def on_user_function_incomplete(self, start_info: UserFunctionStartInfo) -> None: + """Execute plugins for a user function that reported no outcome. + + Dispatched synchronously, like the start and end hooks, so plugins run on + the thread that executed the user function and can release state bound to + it. + """ + self.execute_plugins( + UserFunctionIncompleteInfo.from_start_info(start_info), sync=True + ) + def on_operation_action( self, update: OperationUpdate, 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..4e6bf46a 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 @@ -1165,9 +1165,15 @@ def wrapper(*args, **kwargs): start_info = self._plugin_executor.on_user_function_start( operation_identifier, is_replay_children, attempt ) + # Set once an outcome has been reported, so the finally below can + # tell an unreported exit from a reported one. Suspension is the + # common case, but OrphanedChildException, BackgroundThreadError and + # SystemExit all subclass BaseException and bypass the handlers too. + outcome_reported = False try: result = user_function(*args, **kwargs) self._plugin_executor.on_user_function_end(start_info, None) + outcome_reported = True return result except SuspendExecution: raise @@ -1175,6 +1181,12 @@ def wrapper(*args, **kwargs): self._plugin_executor.on_user_function_end( start_info, ErrorObject.from_exception(e) ) + outcome_reported = True raise + finally: + # Runs on the thread that executed the user function, which is + # the only thread that can release state bound to it. + if not outcome_reported: + self._plugin_executor.on_user_function_incomplete(start_info) return wrapper 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..e0e5560d 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 @@ -4333,6 +4334,9 @@ def on_user_function_start(self, info): def on_user_function_end(self, info): self.calls.append(f"user_function_end:{info.operation_id}") + def on_user_function_incomplete(self, info): + self.calls.append(f"user_function_incomplete:{info.operation_id}") + def test_execution_state_accepts_plugin_executor_parameter(): """Test that ExecutionState can be created with a plugin_executor parameter.""" @@ -5124,3 +5128,123 @@ 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_incomplete:step-1", + ] + + +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", + ] + + +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", + ] + + +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_incomplete(self, info) -> None: # noqa: ARG002 + 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