From 580da37e6c906c929c0e69245eccf7a9316c4a2d Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Fri, 14 Aug 2026 21:40:06 +0000 Subject: [PATCH] fix(plugin): report SUSPENDED user functions wrap_user_function re-raised SuspendExecution without calling on_user_function_end, so a user function that stopped so the execution could resume later never reported its end. Plugins were expected to "observe it by absence" and clean up during their own invocation-end sweep. That contract cannot be honoured for state that is thread-confined. The OTel plugins attach an opentelemetry.context token in on_user_function_start, and a token is only detachable in the contextvars.Context that created it -- the user-code worker thread, not the handler thread the invocation hooks run on. A suspended operation therefore stranded its context scope with no hook able to release it. The same applies to any plugin holding per-operation state: a timer, an open log group, a span. The Java SDK already fires the end hook here. BaseDurableOperation.runUserFunction catches Throwable -- which covers SuspendExecutionException -- and its javadoc gives the same reason: onUserFunctionEnd fires for failures and suspensions alike so plugins can clean up the attempt rather than leak state. Changes: - Add UserFunctionOutcome.SUSPENDED. Suspension is its own outcome rather than reusing FAILED: nothing went wrong, and plugins that count failures or set an error status must not treat it as one. Java models this as succeeded=false plus the suspend exception as the error, which reads as a failure to exactly those consumers. - Allow an explicit outcome on UserFunctionEndInfo.from_start_info and PluginExecutor.on_user_function_end, so the suspension path reports SUSPENDED with error=None instead of deriving the outcome from an absent error. - Fire the hook from wrap_user_function's SuspendExecution branch and re-raise unchanged, so durable control flow is untouched. - Teach both OTel plugins to treat SUSPENDED as "release the scope, leave the span open": the attempt has not concluded, so it must not be ended with an outcome here. It is ended when the operation reaches a terminal status, matching how an operation that suspends mid-invocation is already handled. test_wrap_user_function_suspend_does_not_fire_end_hook pinned the old behaviour and is inverted accordingly. Adds an end-to-end test driving a real child context that suspends, and OTel tests asserting a suspended attempt releases its scope, exports nothing, and is never marked ERROR. Note for reviewers: this makes Python the first of the three SDKs with a third user-function outcome. JS has no hook on this path at all, and Java reports suspension through the existing boolean. A follow-up should decide whether JS and Java adopt SUSPENDED. --- .../execution_plugin.py | 7 +++ .../invocation_plugin.py | 7 +++ .../tests/test_context_scope.py | 62 +++++++++++++++++++ .../plugin.py | 34 ++++++++-- .../aws_durable_execution_sdk_python/state.py | 11 ++++ .../tests/execution_test.py | 46 ++++++++++++++ .../tests/plugin_test.py | 2 +- .../tests/state_test.py | 26 +++++--- 8 files changed, 181 insertions(+), 14 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 6decc311..f413ab04 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 @@ -552,6 +552,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end without matching on_user_function_start" ) + if info.outcome is UserFunctionOutcome.SUSPENDED: + # The user function stopped so the execution can resume later. Leave + # the span open and unexported, exactly as an operation that suspends + # mid-invocation is treated: it is ended when the operation reaches a + # terminal status, in a later invocation if necessary. Detaching the + # scope above is all this hook owes. + return if info.operation_type is OperationType.STEP: span.set_attributes(self._operation_attributes(info)) if info.outcome is UserFunctionOutcome.FAILED: 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 890ac49b..b3ce2d92 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 @@ -660,6 +660,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: "on_user_function_end called without matching on_user_function_start" ) + if info.outcome is UserFunctionOutcome.SUSPENDED: + # The user function stopped so the execution can resume later, so the + # attempt did not conclude. Leave the span open rather than recording + # an outcome on it; on_invocation_end closes whatever is still open. + # Detaching the scope above is all this hook owes. + return + if info.operation_type is OperationType.STEP: span.set_attributes(self._extract_attributes(info)) if info.outcome is UserFunctionOutcome.FAILED: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py index bf787a0b..557c60aa 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py @@ -30,6 +30,7 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode from aws_durable_execution_sdk_python_otel import context_scope from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin @@ -370,6 +371,25 @@ def _context_end(operation_id: str, parent_id: str | None) -> UserFunctionEndInf ) +def _step_suspended(operation_id: str) -> UserFunctionEndInfo: + """End info for a step whose user function suspended.""" + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + parent_id=None, + start_time=START_TIME, + end_time=END_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + outcome=UserFunctionOutcome.SUSPENDED, + error=None, + ) + + @pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) def test_invocation_end_unwinds_a_suspended_operation_scope(factory): """A step that suspends never gets its end hook; invocation end cleans up. @@ -605,6 +625,48 @@ def run_polls() -> None: plugin.on_invocation_end(_invocation_end()) +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_suspended_outcome_detaches_scope_without_ending_the_span(factory): + """A suspended attempt releases its scope but is not exported as finished. + + The core SDK fires on_user_function_end with SUSPENDED when a user function + stops so the execution can resume later. The scope must come off -- that is + the leak this hook exists to prevent -- but the attempt did not conclude, so + the span must not be ended with an outcome here. + """ + plugin, exporter = factory() + before = otel_context.get_current() + plugin.on_invocation_start(_invocation_start()) + + plugin.on_user_function_start(_step_start("step-suspends")) + assert context_scope.depth(plugin) == 1 + + plugin.on_user_function_end(_step_suspended("step-suspends")) + + # Scope released, context restored. + assert context_scope.depth(plugin) == 0 + assert otel_context.get_current() is before + # Nothing exported for the attempt: it has not finished. + assert [s.name for s in exporter.get_finished_spans()] == [] + + plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING)) + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_suspended_outcome_is_not_recorded_as_an_error(factory): + """A suspension must not mark the attempt span ERROR.""" + plugin, exporter = factory() + plugin.on_invocation_start(_invocation_start()) + plugin.on_user_function_start(_step_start("step-suspends")) + + plugin.on_user_function_end(_step_suspended("step-suspends")) + plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING)) + + for span in exporter.get_finished_spans(): + assert span.status.status_code is not StatusCode.ERROR + assert span.attributes.get("durable.attempt.outcome") != "SUSPENDED" + + def test_two_plugins_on_one_thread_unwind_in_lifo_order(): """Both plugins ship as entry points and can be enabled together. 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 ffc5a65f..5eed2cd8 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 @@ -163,6 +163,13 @@ class OperationChangeInfo: class UserFunctionOutcome(Enum): SUCCEEDED = "SUCCEEDED" FAILED = "FAILED" + # The user function did not finish: it suspended so the execution can be + # resumed in a later invocation (e.g. a child context whose inner operation + # is still pending). Reported as its own outcome rather than FAILED because + # nothing went wrong -- plugins that count failures or set an error status + # must not treat a suspension as one, and plugins holding per-operation + # state need the hook to fire so they can release it. + SUSPENDED = "SUSPENDED" @classmethod def from_error(cls, error: ErrorObject | None) -> UserFunctionOutcome: @@ -187,8 +194,20 @@ 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: + """Build the end info for a user function that has stopped running. + + Args: + start_info: The info reported when the user function started. + error: The failure, if the user function raised one. + outcome: Overrides the outcome derived from ``error``. Used for + suspension, which is neither a success nor a failure and carries + no error. + """ return UserFunctionEndInfo( operation_id=start_info.operation_id, operation_type=start_info.operation_type, @@ -200,7 +219,9 @@ 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, ) @@ -622,10 +643,15 @@ 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: + def on_user_function_end( + self, + start_info: UserFunctionStartInfo, + error, + outcome: UserFunctionOutcome | None = None, + ) -> None: """Execute any registered plugins for the operation when its user function finishes execution.""" self.execute_plugins( - UserFunctionEndInfo.from_start_info(start_info, error), sync=True + UserFunctionEndInfo.from_start_info(start_info, error, 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..d8dc1792 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 @@ -37,6 +37,7 @@ ) from aws_durable_execution_sdk_python.plugin import ( PluginExecutor, + UserFunctionOutcome, ) from aws_durable_execution_sdk_python.threading import CompletionEvent @@ -1170,6 +1171,16 @@ def wrapper(*args, **kwargs): self._plugin_executor.on_user_function_end(start_info, None) return result except SuspendExecution: + # The user function did not finish -- it stopped so the execution + # can resume in a later invocation. The end hook still has to + # fire: it is the only signal a plugin gets that this operation's + # user code is no longer running, and without it any per-operation + # state a plugin opened in on_user_function_start (an OTel context + # scope, a timer, an open log group) is stranded. Reported as + # SUSPENDED with no error so plugins do not record a failure. + self._plugin_executor.on_user_function_end( + start_info, None, UserFunctionOutcome.SUSPENDED + ) raise except Exception as e: self._plugin_executor.on_user_function_end( diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index ee04ce30..0ea0dc3b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -2923,6 +2923,12 @@ def on_operation_attempt_start(self, info): def on_operation_attempt_end(self, info): self.calls.append(f"attempt_end:{info.operation_id}") + def on_user_function_start(self, info): + self.calls.append(f"user_function_start:{info.operation_id}") + + def on_user_function_end(self, info): + self.calls.append(f"user_function_end:{info.operation_id}:{info.outcome.value}") + class _FailingPlugin(DurableInstrumentationPlugin): """Plugin that raises on every hook call.""" @@ -3182,6 +3188,46 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert len(execution_end_calls) == 0 +def test_durable_execution_with_plugins_child_context_suspends(): + """A child context that suspends reports SUSPENDED, not FAILED. + + This is the reachable suspension path: the child context's user function runs + inner durable operations, one of them is still pending, and SuspendExecution + propagates out of the user function. Plugins must see the end hook so they can + release whatever they opened at start, with an outcome that does not read as a + failure. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + plugin = _RecordingPlugin() + + @durable_execution(plugins=[plugin]) + def test_handler(event: Any, context: DurableContext) -> dict: + def child(ctx: DurableContext) -> dict: + raise SuspendExecution("inner operation still pending") + + return context.run_in_child_context(child, name="child-1") + + result = test_handler( + _make_invocation_input(mock_client), + _make_lambda_context(), + ) + + assert result["Status"] == InvocationStatus.PENDING.value + suspended = [ + c + for c in plugin.calls + if c.startswith("user_function_end") and c.endswith(":SUSPENDED") + ] + assert len(suspended) == 1, plugin.calls + # Never reported as a failure. + assert not [c for c in plugin.calls if c.endswith("user_function_end:FAILED")] + + def test_durable_execution_with_plugins_retryable_error(): """Test that plugins receive invocation end with RETRY status on retryable error.""" mock_client = Mock(spec=DurableServiceClient) 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 c33c370a..229de757 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1788,7 +1788,7 @@ class TestUserFunctionOutcomeValues(unittest.TestCase): def test_outcome_values(self): self.assertEqual( {o.value for o in UserFunctionOutcome}, - {"SUCCEEDED", "FAILED"}, + {"SUCCEEDED", "FAILED", "SUSPENDED"}, ) 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..d071f8a6 100644 --- a/packages/aws-durable-execution-sdk-python/tests/state_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/state_test.py @@ -45,6 +45,7 @@ OperationStartInfo, PluginExecutor, UserFunctionEndInfo, + UserFunctionOutcome, ) from aws_durable_execution_sdk_python.state import ( CheckpointBatcherConfig, @@ -4821,15 +4822,17 @@ 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_end_hook_with_suspended_outcome(): + """A user function that suspends fires the end hook with SUSPENDED. - 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. + A timed suspend (TimedSuspendExecution) raised inside a wrapped user function + (e.g. a child context that waits) is normal durable control flow, so it must + not be surfaced as a FAILED outcome. It must still fire the end hook: that is + the only signal a plugin gets that this operation's user code stopped + running, and per-operation state a plugin opened in on_user_function_start + cannot always be released at invocation end -- an OTel context token, for + one, is only detachable on the thread that attached it, which is not the + thread the invocation hooks run on. """ captured: list[UserFunctionEndInfo] = [] @@ -4858,7 +4861,12 @@ def suspends(_: object) -> None: with pytest.raises(TimedSuspendExecution): wrapped(None) - assert captured == [] + assert len(captured) == 1 + assert captured[0].outcome is UserFunctionOutcome.SUSPENDED + # A suspension is not a failure, so no error is reported. + assert captured[0].error is None + assert captured[0].operation_id == "op-1" + assert captured[0].attempt == 1 def test_plugin_executor_not_called_for_pending_operations():