fix(plugin): fire on_user_function_end with SUSPENDED on suspension - #648
Closed
wangyb-A wants to merge 6 commits into
Closed
fix(plugin): fire on_user_function_end with SUSPENDED on suspension#648wangyb-A wants to merge 6 commits into
wangyb-A wants to merge 6 commits into
Conversation
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
force-pushed
the
fix/otel-balance-context-scopes
branch
from
August 14, 2026 22:12
ceb4551 to
c79cfda
Compare
wangyb-A
force-pushed
the
fix/user-function-end-on-suspension
branch
from
August 14, 2026 22:19
3fe5c79 to
d868e27
Compare
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
force-pushed
the
fix/user-function-end-on-suspension
branch
from
August 14, 2026 23:03
d868e27 to
e1e03c7
Compare
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
force-pushed
the
fix/user-function-end-on-suspension
branch
from
August 14, 2026 23:50
e1e03c7 to
580da37
Compare
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.
wangyb-A
force-pushed
the
fix/user-function-end-on-suspension
branch
from
August 17, 2026 17:39
580da37 to
27f657b
Compare
wangyb-A
force-pushed
the
fix/otel-balance-context-scopes
branch
from
August 17, 2026 22:12
8ace0f5 to
4e59e90
Compare
This was referenced Aug 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
wrap_user_functionre-raisedSuspendExecutionwithout callingon_user_function_end(state.py:1171):This was deliberate —
test_wrap_user_function_suspend_does_not_fire_end_hookpinned 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.contexttoken inon_user_function_start, and a token is only detachable in thecontextvars.Contextthat 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:
SuspendExecutionpropagates out of the child context's user function, and the hook silently never fires.Java already fires the hook here.
BaseDurableOperation.runUserFunctioncatchesThrowable— which coversSuspendExecutionException— and the javadoc gives the same reason this PR does:Changes
UserFunctionOutcome.SUSPENDED. Reported as its own outcome rather than reusingFAILED: nothing went wrong, and plugins that count failures or set an error span status must not treat it as one. Java models this assucceeded=falseplus the suspend exception as the error, which reads as a failure to exactly those consumers.UserFunctionEndInfo.from_start_infoandPluginExecutor.on_user_function_end, so the suspension path can reportSUSPENDEDwitherror=Nonerather than deriving the outcome from an absent error.wrap_user_function'sSuspendExecutionbranch, then re-raise unchanged — durable control flow is untouched.SUSPENDEDas "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
SuspendExecutionderives fromBaseException, notException, so the existingexcept Exceptionclause never caught it — the gap was structural, not a missing branch. This PR keeps the catch narrow (SuspendExecutiononly); whetherwrap_user_functionshould also fire the hook for otherBaseExceptions such asKeyboardInterrupt, as Java'scatch (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
UserFunctionOutcomewill now see a third value. Anyone writingif 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
waitForConditiononly, not child contexts, and its plugins scope context throughwrapOperationAttemptFn/wrapChildContextFn, which restore on throw automatically. Java reports suspension through the existing boolean. Worth deciding whether JS and Java adoptSUSPENDEDfor 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 --checkandhatch run types:checkclean.test_wrap_user_function_suspend_does_not_fire_end_hookis inverted intotest_wrap_user_function_suspend_fires_end_hook_with_suspended_outcome, asserting one end hook withSUSPENDEDanderror is None.run_in_child_contextthat suspends and asserts exactly oneSUSPENDEDend hook and noFAILEDone, with the invocation still returningPENDING.ERROR(parametrised over both plugins).UserFunctionOutcomevalue-set test is updated.Both halves mutation-tested: removing the hook call fails the two core tests; ignoring
SUSPENDEDin the plugins fails the four OTel tests.