Skip to content

fix(plugin): fire on_user_function_end with SUSPENDED on suspension - #648

Closed
wangyb-A wants to merge 6 commits into
fix/otel-balance-context-scopesfrom
fix/user-function-end-on-suspension
Closed

fix(plugin): fire on_user_function_end with SUSPENDED on suspension#648
wangyb-A wants to merge 6 commits into
fix/otel-balance-context-scopesfrom
fix/user-function-end-on-suspension

Conversation

@wangyb-A

Copy link
Copy Markdown
Contributor

Stacked on #647. Base is fix/otel-balance-context-scopes; review that first. It touches the same on_user_function_end functions, and landing the core change without the OTel side would export suspended attempt spans as if they had completed successfully.

Problem

wrap_user_function re-raised SuspendExecution without calling on_user_function_end (state.py:1171):

except SuspendExecution:
    raise
except Exception as e:
    self._plugin_executor.on_user_function_end(start_info, ErrorObject.from_exception(e))
    raise

This was deliberate — test_wrap_user_function_suspend_does_not_fire_end_hook pinned it, on the rationale that "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."

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. So a suspended operation stranded its context scope with no hook able to release it: the invocation-end sweep runs on the wrong thread. The same applies to any plugin holding per-operation state — a timer, an open log group, a span.

Suspension is also not a rare path. It is the normal outcome of any child context whose inner operation is still pending, which is the reachable case here: SuspendExecution propagates out of the child context's user function, and the hook silently never fires.

Java already fires the hook here. BaseDurableOperation.runUserFunction catches Throwable — which covers SuspendExecutionException — and the javadoc gives the same reason this PR does:

onUserFunctionEnd fires for failures and suspensions alike so plugins (e.g. OTel) can end/clean up the attempt rather than leak state.

Changes

  • UserFunctionOutcome.SUSPENDED. Reported as its own outcome rather than reusing FAILED: nothing went wrong, and plugins that count failures or set an error span 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.
  • Explicit outcome override on UserFunctionEndInfo.from_start_info and PluginExecutor.on_user_function_end, so the suspension path can report SUSPENDED with error=None rather than deriving the outcome from an absent error.
  • Fire the hook from wrap_user_function's SuspendExecution branch, then re-raise unchanged — durable control flow is untouched.
  • Both OTel plugins 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.

Note SuspendExecution derives from BaseException, not Exception, so the existing except Exception clause never caught it — the gap was structural, not a missing branch. This PR keeps the catch narrow (SuspendExecution only); whether wrap_user_function should also fire the hook for other BaseExceptions such as KeyboardInterrupt, as Java's catch (Throwable) effectively does, is left as a separate question.

API impact

Adding an enum member is source-compatible, but plugin authors who exhaustively branch on UserFunctionOutcome will now see a third value. Anyone writing if outcome is FAILED: ... else: <treat as success> will classify a suspension as success, which is the intended reading for span status but may not be for metrics.

This also makes Python the first of the three SDKs with a third outcome. JS has no hook on this path at all — its attempt hooks cover steps and waitForCondition only, not child contexts, and its plugins scope context through wrapOperationAttemptFn/wrapChildContextFn, which restore on throw automatically. Java reports suspension through the existing boolean. Worth deciding whether JS and Java adopt SUSPENDED for parity; happy to split this into a discussion first if the team would rather agree the shape before the code lands.

Testing

3229 pass across the monorepo; hatch fmt --check and hatch run types:check clean.

  • test_wrap_user_function_suspend_does_not_fire_end_hook is inverted into test_wrap_user_function_suspend_fires_end_hook_with_suspended_outcome, asserting one end hook with SUSPENDED and error is None.
  • New end-to-end test drives a real run_in_child_context that suspends and asserts exactly one SUSPENDED end hook and no FAILED one, with the invocation still returning PENDING.
  • New OTel tests assert a suspended attempt releases its scope, restores the prior context, exports nothing, and is never marked ERROR (parametrised over both plugins).
  • The UserFunctionOutcome value-set test is updated.

Both halves mutation-tested: removing the hook call fails the two core tests; ignoring SUSPENDED in the plugins fails the four OTel tests.

Alex Wang added 2 commits August 14, 2026 22:11
Both OTel plugins called opentelemetry.context.attach() without
keeping the returned token, and "restored" the enclosing span by
attaching another context rather than detaching. Every operation
pushed two context layers and popped none, leaving an ended span
current after its scope had finished.

The worst consequence is not in the issue. The invocation-start
attach ran on the Lambda handler thread, which is reused across
warm invocations, so the next execution's context extractor and
GLOBAL-mode ambient-parent lookup adopted the previous
execution's ended Workflow span. A valid parent overrides the
deterministic ID generator's trace ID, so two unrelated durable
executions merged into a single trace.

Two runtime properties shaped the fix. The hooks run on several
threads -- the invocation hooks on the Lambda handler thread, the
user-function hooks on the worker that runs user code, and on a
branch worker per map/parallel branch -- and ContextVar.reset()
only accepts a token created in the same contextvars.Context.
And unlike Java's ScopeImpl.close(), which ignores a close that
does not represent the current context, ContextVar.reset() writes
back its captured value unconditionally, so detaching out of
order revives a stale context instead of failing safe.

Changes:

- Add context_scope, a thread-confined LIFO stack of attach
  tokens. Detaches unwind downwards so the ContextVar is always
  reset in order. The stack is module level so the two plugins,
  which ship as separate entry points and can be enabled
  together, still unwind in true LIFO order. An epoch check
  discards scopes a suspended operation left behind, since the
  SDK re-raises SuspendExecution without calling
  on_user_function_end.
- Pair every user-function attach with a detach on the same
  thread, replacing the re-attach that previously stood in for
  restoring the enclosing context.
- Unwind any remaining scopes at invocation end.
- Drop the invocation-start attach. User code runs on a separate
  worker and ThreadPoolExecutor does not copy contextvars, so
  that attach never reached the code it was meant to parent; it
  only leaked. The Workflow and Invocation spans are used as
  explicit parents instead, matching the Java plugins, which
  never make either span current.

Ambient spans emitted outside any operation are no longer
parented to the Invocation span; the README documents this, and
log correlation is unchanged because the logging filter resolves
through the plugin's span registry.

Tests no longer reset the OTel context to isolate themselves; an
autouse fixture asserts instead that every test leaves the
context exactly as it found it. Adds coverage for nested
contexts, sequential steps, failures, suspension, worker-thread
hooks, both plugins on one thread, and warm invocation reuse
keeping two executions in separate traces.

Fixes #643
Three review findings, all real:

Same-invocation re-entry. The epoch check only caught a previous
invocation's leftovers, but the same operation key can be entered
twice inside one invocation: a suspended operation is re-entered
when its branch is resubmitted, and its first scope is still
attached because the suspending path has no end hook. The second
enter stacked on the first and the end hook popped one, leaving a
stale layer per re-entry. enter_scope now unwinds an existing
(owner, key) even when the epoch matches.

Extracted context values. Basing every scope on the current
context dropped baggage and other non-span values supplied by the
context extractor, because the worker running user code starts
with an empty context. The outermost scope on a thread is now
layered onto the extracted context; nested scopes keep using the
current one, which already carries it transitively.

Ambient span vs durable span. In GLOBAL mode the ADOT Lambda span
stays current on the handler thread, so get_current_span_context
returned it instead of the Invocation span, contradicting what the
previous commit documented. The current span is now trusted only
while this plugin holds a scope on this thread; otherwise the
registry answers. The earlier test missed this by using an
explicit provider with no ambient span.

Adds the tests each finding asked for: same-key re-entry at the
helper level and through the plugin hooks on a worker thread,
baggage surviving into user code and into a nested scope, and a
GLOBAL-mode ambient Lambda span not displacing the Invocation
span in log records.
@wangyb-A
wangyb-A force-pushed the fix/otel-balance-context-scopes branch from ceb4551 to c79cfda Compare August 14, 2026 22:12
@wangyb-A
wangyb-A force-pushed the fix/user-function-end-on-suspension branch from 3fe5c79 to d868e27 Compare August 14, 2026 22:19
Two more review findings, both real.

Branch workers have no branch affinity. If branch A suspends
without an end hook and its pool worker next runs branch B, A's
scope has the same epoch and a different key, so neither the epoch
check nor the same-key check cleared it. B nested inside A and, on
exit, detached back into it -- later records on that worker
correlated to the wrong branch, and one layer accumulated per
suspended branch.

Replaces the same-key guard with an ancestry check: a scope may
stay attached only while the operation it belongs to is still
running on this thread, so anything above the new scope's parent is
stale, and when the parent is absent -- a root-level operation, or
one whose parent ran elsewhere -- nothing held here can enclose it.
The plugins pass the enclosing operation as parent_key. The normal
nesting path is a no-op, which matters because detaching
necessarily discards entries above the cut, including a second
plugin's.

An empty extracted context was being discarded. Context subclasses
dict, so an empty one is falsy and `extracted or current` silently
inverted the intent of an extractor that returns an empty context
to isolate the operation, inheriting the worker's ambient baggage
and suppression values instead. Now tested with `is not None`.

Adds the tests both findings asked for: a sibling scope dropped
rather than nested into, a nested scope whose parent never ran on
this thread, a plugin-level branch-A-suspends-then-branch-B case
pinned to one worker, and an empty extracted context isolating the
operation from ambient baggage. Two existing helper tests nested
without declaring a parent, which now reads as a root operation;
they pass parent_key like the plugins do.
@wangyb-A
wangyb-A force-pushed the fix/user-function-end-on-suspension branch from d868e27 to e1e03c7 Compare August 14, 2026 23:03
Reverts the ancestry check from the previous commit and fixes the
order in which the attached context is built. Both were review
findings; the first one was a regression I introduced.

parent_id is checkpoint hierarchy, not the Python call stack. A
virtual (FLAT) map/parallel branch deliberately reports its inner
operations' parent as the grandparent -- None for a top-level
branch -- while the branch's own context scope is still running
(see DurableContext.is_virtual and create_child_context, where
child_parent_id is the *parent's* parent when is_virtual). Treating
such an inner step as root-level therefore detached the live branch
scope at the first inner step, and work between two inner steps
fell out of the durable trace. That is worse than the abandoned
sibling scope the check was meant to catch, so the narrower
same-key guard is restored.

The gap that leaves -- a scope abandoned by a *different* operation
on the same branch-pool worker -- cannot be closed from the hook
payloads, because a live FLAT branch scope is indistinguishable
from an abandoned sibling. It needs the SDK to report the end of a
suspended user function, which is tracked separately; the docstring
says so rather than implying the helper handles it.

Second finding: the context to attach was built by the caller
before enter_scope ran its cleanup, so it copied baggage and
suppression values out of the very scope about to be detached, and
detaching afterwards could not remove them from an already-built
Context. enter_scope now takes a factory and calls it after
cleanup.

Adds a FLAT-branch test asserting the branch scope stays current
across two inner steps that report no parent, and a test that the
factory observes the post-cleanup context. Drops the three tests
that asserted the reverted rule.
@wangyb-A
wangyb-A force-pushed the fix/user-function-end-on-suspension branch from e1e03c7 to 580da37 Compare August 14, 2026 23:50
Alex Wang added 2 commits August 17, 2026 17:35
Both plugins decided "is the current span mine?" from the
thread-local token count. That is the wrong place to keep the
answer. When user code propagates an operation's context to
another thread -- asyncio.to_thread, contextvars.copy_context, an
instrumented executor -- that thread has the operation's span
current but holds no tokens, so the count read zero: log records
fell back to the invocation span, losing the attempt span, and a
scope entered there was treated as outermost and rebased onto the
extracted context, discarding what had been propagated.

Ownership now lives in the attached context itself, so it travels
with any copy of it. enter_scope stamps the entering plugin's id
into the context it attaches, and owns_current reads it back. The
ambient Lambda span an ADOT layer makes current on the handler
thread carries no marker, so it still loses to the registry --
the reason the check exists at all.

Ids accumulate rather than overwrite, so with both plugins enabled
each recognises its own scope instead of only the innermost one;
previously the depth check made a plugin trust the *other*
plugin's span.

Adds tests for a propagated context resolving to its own span in
both plugins, a nested scope built on a propagated context keeping
its baggage, an unowned ambient context not being claimed, and two
owners each recognising their own scope.
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant