Skip to content

Commit 27f657b

Browse files
author
Alex Wang
committed
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.
1 parent 8ace0f5 commit 27f657b

8 files changed

Lines changed: 181 additions & 14 deletions

File tree

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
557557
raise RuntimeError(
558558
"on_user_function_end without matching on_user_function_start"
559559
)
560+
if info.outcome is UserFunctionOutcome.SUSPENDED:
561+
# The user function stopped so the execution can resume later. Leave
562+
# the span open and unexported, exactly as an operation that suspends
563+
# mid-invocation is treated: it is ended when the operation reaches a
564+
# terminal status, in a later invocation if necessary. Detaching the
565+
# scope above is all this hook owes.
566+
return
560567
if info.operation_type is OperationType.STEP:
561568
span.set_attributes(self._operation_attributes(info))
562569
if info.outcome is UserFunctionOutcome.FAILED:

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -662,6 +662,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
662662
"on_user_function_end called without matching on_user_function_start"
663663
)
664664

665+
if info.outcome is UserFunctionOutcome.SUSPENDED:
666+
# The user function stopped so the execution can resume later, so the
667+
# attempt did not conclude. Leave the span open rather than recording
668+
# an outcome on it; on_invocation_end closes whatever is still open.
669+
# Detaching the scope above is all this hook owes.
670+
return
671+
665672
if info.operation_type is OperationType.STEP:
666673
span.set_attributes(self._extract_attributes(info))
667674
if info.outcome is UserFunctionOutcome.FAILED:

packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from opentelemetry.sdk.trace import TracerProvider
3131
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
3232
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
33+
from opentelemetry.trace import StatusCode
3334

3435
from aws_durable_execution_sdk_python_otel import context_scope
3536
from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin
@@ -429,6 +430,25 @@ def _context_end(operation_id: str, parent_id: str | None) -> UserFunctionEndInf
429430
)
430431

431432

433+
def _step_suspended(operation_id: str) -> UserFunctionEndInfo:
434+
"""End info for a step whose user function suspended."""
435+
return UserFunctionEndInfo(
436+
operation_id=operation_id,
437+
operation_type=OperationType.STEP,
438+
sub_type=OperationSubType.STEP,
439+
name=operation_id,
440+
parent_id=None,
441+
start_time=START_TIME,
442+
end_time=END_TIME,
443+
is_replayed=False,
444+
status=OperationStatus.STARTED,
445+
is_replay_children=False,
446+
attempt=1,
447+
outcome=UserFunctionOutcome.SUSPENDED,
448+
error=None,
449+
)
450+
451+
432452
@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
433453
def test_invocation_end_unwinds_a_suspended_operation_scope(factory):
434454
"""A step that suspends never gets its end hook; invocation end cleans up.
@@ -747,6 +767,48 @@ def run_polls() -> None:
747767
plugin.on_invocation_end(_invocation_end())
748768

749769

770+
@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
771+
def test_suspended_outcome_detaches_scope_without_ending_the_span(factory):
772+
"""A suspended attempt releases its scope but is not exported as finished.
773+
774+
The core SDK fires on_user_function_end with SUSPENDED when a user function
775+
stops so the execution can resume later. The scope must come off -- that is
776+
the leak this hook exists to prevent -- but the attempt did not conclude, so
777+
the span must not be ended with an outcome here.
778+
"""
779+
plugin, exporter = factory()
780+
before = otel_context.get_current()
781+
plugin.on_invocation_start(_invocation_start())
782+
783+
plugin.on_user_function_start(_step_start("step-suspends"))
784+
assert context_scope.depth(plugin) == 1
785+
786+
plugin.on_user_function_end(_step_suspended("step-suspends"))
787+
788+
# Scope released, context restored.
789+
assert context_scope.depth(plugin) == 0
790+
assert otel_context.get_current() is before
791+
# Nothing exported for the attempt: it has not finished.
792+
assert [s.name for s in exporter.get_finished_spans()] == []
793+
794+
plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING))
795+
796+
797+
@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin])
798+
def test_suspended_outcome_is_not_recorded_as_an_error(factory):
799+
"""A suspension must not mark the attempt span ERROR."""
800+
plugin, exporter = factory()
801+
plugin.on_invocation_start(_invocation_start())
802+
plugin.on_user_function_start(_step_start("step-suspends"))
803+
804+
plugin.on_user_function_end(_step_suspended("step-suspends"))
805+
plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING))
806+
807+
for span in exporter.get_finished_spans():
808+
assert span.status.status_code is not StatusCode.ERROR
809+
assert span.attributes.get("durable.attempt.outcome") != "SUSPENDED"
810+
811+
750812
def test_two_plugins_on_one_thread_unwind_in_lifo_order():
751813
"""Both plugins ship as entry points and can be enabled together.
752814

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,13 @@ class OperationChangeInfo:
163163
class UserFunctionOutcome(Enum):
164164
SUCCEEDED = "SUCCEEDED"
165165
FAILED = "FAILED"
166+
# The user function did not finish: it suspended so the execution can be
167+
# resumed in a later invocation (e.g. a child context whose inner operation
168+
# is still pending). Reported as its own outcome rather than FAILED because
169+
# nothing went wrong -- plugins that count failures or set an error status
170+
# must not treat a suspension as one, and plugins holding per-operation
171+
# state need the hook to fire so they can release it.
172+
SUSPENDED = "SUSPENDED"
166173

167174
@classmethod
168175
def from_error(cls, error: ErrorObject | None) -> UserFunctionOutcome:
@@ -187,8 +194,20 @@ class UserFunctionEndInfo(OperationInfo):
187194

188195
@classmethod
189196
def from_start_info(
190-
cls, start_info: UserFunctionStartInfo, error: ErrorObject | None
197+
cls,
198+
start_info: UserFunctionStartInfo,
199+
error: ErrorObject | None,
200+
outcome: UserFunctionOutcome | None = None,
191201
) -> UserFunctionEndInfo:
202+
"""Build the end info for a user function that has stopped running.
203+
204+
Args:
205+
start_info: The info reported when the user function started.
206+
error: The failure, if the user function raised one.
207+
outcome: Overrides the outcome derived from ``error``. Used for
208+
suspension, which is neither a success nor a failure and carries
209+
no error.
210+
"""
192211
return UserFunctionEndInfo(
193212
operation_id=start_info.operation_id,
194213
operation_type=start_info.operation_type,
@@ -200,7 +219,9 @@ def from_start_info(
200219
status=start_info.status,
201220
is_replay_children=start_info.is_replay_children,
202221
attempt=start_info.attempt,
203-
outcome=UserFunctionOutcome.from_error(error),
222+
outcome=outcome
223+
if outcome is not None
224+
else UserFunctionOutcome.from_error(error),
204225
end_time=datetime.datetime.now(datetime.UTC),
205226
error=error,
206227
)
@@ -622,10 +643,15 @@ def on_user_function_start(
622643
self.execute_plugins(start_info, sync=True)
623644
return start_info
624645

625-
def on_user_function_end(self, start_info: UserFunctionStartInfo, error) -> None:
646+
def on_user_function_end(
647+
self,
648+
start_info: UserFunctionStartInfo,
649+
error,
650+
outcome: UserFunctionOutcome | None = None,
651+
) -> None:
626652
"""Execute any registered plugins for the operation when its user function finishes execution."""
627653
self.execute_plugins(
628-
UserFunctionEndInfo.from_start_info(start_info, error), sync=True
654+
UserFunctionEndInfo.from_start_info(start_info, error, outcome), sync=True
629655
)
630656

631657
def on_operation_action(

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
)
3838
from aws_durable_execution_sdk_python.plugin import (
3939
PluginExecutor,
40+
UserFunctionOutcome,
4041
)
4142
from aws_durable_execution_sdk_python.threading import CompletionEvent
4243

@@ -1170,6 +1171,16 @@ def wrapper(*args, **kwargs):
11701171
self._plugin_executor.on_user_function_end(start_info, None)
11711172
return result
11721173
except SuspendExecution:
1174+
# The user function did not finish -- it stopped so the execution
1175+
# can resume in a later invocation. The end hook still has to
1176+
# fire: it is the only signal a plugin gets that this operation's
1177+
# user code is no longer running, and without it any per-operation
1178+
# state a plugin opened in on_user_function_start (an OTel context
1179+
# scope, a timer, an open log group) is stranded. Reported as
1180+
# SUSPENDED with no error so plugins do not record a failure.
1181+
self._plugin_executor.on_user_function_end(
1182+
start_info, None, UserFunctionOutcome.SUSPENDED
1183+
)
11731184
raise
11741185
except Exception as e:
11751186
self._plugin_executor.on_user_function_end(

packages/aws-durable-execution-sdk-python/tests/execution_test.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2923,6 +2923,12 @@ def on_operation_attempt_start(self, info):
29232923
def on_operation_attempt_end(self, info):
29242924
self.calls.append(f"attempt_end:{info.operation_id}")
29252925

2926+
def on_user_function_start(self, info):
2927+
self.calls.append(f"user_function_start:{info.operation_id}")
2928+
2929+
def on_user_function_end(self, info):
2930+
self.calls.append(f"user_function_end:{info.operation_id}:{info.outcome.value}")
2931+
29262932

29272933
class _FailingPlugin(DurableInstrumentationPlugin):
29282934
"""Plugin that raises on every hook call."""
@@ -3182,6 +3188,46 @@ def test_handler(event: Any, context: DurableContext) -> dict:
31823188
assert len(execution_end_calls) == 0
31833189

31843190

3191+
def test_durable_execution_with_plugins_child_context_suspends():
3192+
"""A child context that suspends reports SUSPENDED, not FAILED.
3193+
3194+
This is the reachable suspension path: the child context's user function runs
3195+
inner durable operations, one of them is still pending, and SuspendExecution
3196+
propagates out of the user function. Plugins must see the end hook so they can
3197+
release whatever they opened at start, with an outcome that does not read as a
3198+
failure.
3199+
"""
3200+
mock_client = Mock(spec=DurableServiceClient)
3201+
mock_client.checkpoint.return_value = CheckpointOutput(
3202+
checkpoint_token="new_token", # noqa: S106
3203+
new_execution_state=CheckpointUpdatedExecutionState(),
3204+
)
3205+
3206+
plugin = _RecordingPlugin()
3207+
3208+
@durable_execution(plugins=[plugin])
3209+
def test_handler(event: Any, context: DurableContext) -> dict:
3210+
def child(ctx: DurableContext) -> dict:
3211+
raise SuspendExecution("inner operation still pending")
3212+
3213+
return context.run_in_child_context(child, name="child-1")
3214+
3215+
result = test_handler(
3216+
_make_invocation_input(mock_client),
3217+
_make_lambda_context(),
3218+
)
3219+
3220+
assert result["Status"] == InvocationStatus.PENDING.value
3221+
suspended = [
3222+
c
3223+
for c in plugin.calls
3224+
if c.startswith("user_function_end") and c.endswith(":SUSPENDED")
3225+
]
3226+
assert len(suspended) == 1, plugin.calls
3227+
# Never reported as a failure.
3228+
assert not [c for c in plugin.calls if c.endswith("user_function_end:FAILED")]
3229+
3230+
31853231
def test_durable_execution_with_plugins_retryable_error():
31863232
"""Test that plugins receive invocation end with RETRY status on retryable error."""
31873233
mock_client = Mock(spec=DurableServiceClient)

packages/aws-durable-execution-sdk-python/tests/plugin_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1788,7 +1788,7 @@ class TestUserFunctionOutcomeValues(unittest.TestCase):
17881788
def test_outcome_values(self):
17891789
self.assertEqual(
17901790
{o.value for o in UserFunctionOutcome},
1791-
{"SUCCEEDED", "FAILED"},
1791+
{"SUCCEEDED", "FAILED", "SUSPENDED"},
17921792
)
17931793

17941794

packages/aws-durable-execution-sdk-python/tests/state_test.py

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
OperationStartInfo,
4646
PluginExecutor,
4747
UserFunctionEndInfo,
48+
UserFunctionOutcome,
4849
)
4950
from aws_durable_execution_sdk_python.state import (
5051
CheckpointBatcherConfig,
@@ -4821,15 +4822,17 @@ def on_operation_end(self, info):
48214822
executor.shutdown(wait=True)
48224823

48234824

4824-
def test_wrap_user_function_suspend_does_not_fire_end_hook():
4825-
"""A user function that suspends does not fire the end hook.
4825+
def test_wrap_user_function_suspend_fires_end_hook_with_suspended_outcome():
4826+
"""A user function that suspends fires the end hook with SUSPENDED.
48264827
4827-
Regression: a timed suspend (TimedSuspendExecution) raised inside a wrapped
4828-
user function (e.g. a child context that waits) must not be surfaced to
4829-
plugins as a FAILED outcome. The suspend is normal durable control flow,
4830-
and the plugin observes it by absence (no end hook fires), with the
4831-
instrumentation plugin's own per-invocation span sweep closing any open
4832-
spans cleanly at invocation end.
4828+
A timed suspend (TimedSuspendExecution) raised inside a wrapped user function
4829+
(e.g. a child context that waits) is normal durable control flow, so it must
4830+
not be surfaced as a FAILED outcome. It must still fire the end hook: that is
4831+
the only signal a plugin gets that this operation's user code stopped
4832+
running, and per-operation state a plugin opened in on_user_function_start
4833+
cannot always be released at invocation end -- an OTel context token, for
4834+
one, is only detachable on the thread that attached it, which is not the
4835+
thread the invocation hooks run on.
48334836
"""
48344837
captured: list[UserFunctionEndInfo] = []
48354838

@@ -4858,7 +4861,12 @@ def suspends(_: object) -> None:
48584861
with pytest.raises(TimedSuspendExecution):
48594862
wrapped(None)
48604863

4861-
assert captured == []
4864+
assert len(captured) == 1
4865+
assert captured[0].outcome is UserFunctionOutcome.SUSPENDED
4866+
# A suspension is not a failure, so no error is reported.
4867+
assert captured[0].error is None
4868+
assert captured[0].operation_id == "op-1"
4869+
assert captured[0].attempt == 1
48624870

48634871

48644872
def test_plugin_executor_not_called_for_pending_operations():

0 commit comments

Comments
 (0)