fix(otel): balance otel context attach and detach - #650
Conversation
- Track every context.attach() token in the plugin and release it at the matching lifecycle end: user-function scopes in on_user_function_end, the invocation scope during invocation cleanup - Replace the re-attach of the enclosing span in on_user_function_end with a detach, so nested operations manage a balanced scope instead of stacking a new one - Release scopes still open at invocation end (a suspended user function or an aborted invocation) so nothing survives into a warm invocation - Skip tokens recorded on another thread, which OpenTelemetry cannot reset - Replace the autouse OTel context reset in the plugin tests with an assertion that each test leaves the context as it found it, and cover nested child contexts, sequential steps, failures, suspension, and warm invocation reuse Resolves #643
GitHub Actions
Inactive
GitHub Actions
Error
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| self._context_tokens[key] = ( | ||
| threading.get_ident(), | ||
| context.attach(new_context), | ||
| ) |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
| if thread_ident == threading.get_ident(): | ||
| otel_context.detach(token) # type: ignore[arg-type] |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
There was a problem hiding this comment.
tracked in #622 2 plugins should not be used at the same time
This comment has been minimized.
This comment has been minimized.
GitHub Actions
Inactive
GitHub Actions
Error
| def _attach_context(self, key: str, new_context: Context) -> None: | ||
| """Attach a context and remember its token under ``key``.""" | ||
| with self._lock: | ||
| self._context_tokens[key] = ( |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
| return | ||
| thread_ident, token = entry | ||
| if thread_ident == threading.get_ident(): | ||
| context.detach(token) # type: ignore[arg-type] |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment has been minimized.
This comment has been minimized.
GitHub Actions
Inactive
GitHub Actions
Error
| thread_ident, token = entry | ||
| if thread_ident == threading.get_ident(): | ||
| context.detach(token) # type: ignore[arg-type] |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
| with self._lock: | ||
| entry = self._context_tokens.pop(key, None) | ||
| if entry is None: | ||
| return | ||
| thread_ident, token = entry | ||
| if thread_ident == threading.get_ident(): | ||
| otel_context.detach(token) # type: ignore[arg-type] |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
There was a problem hiding this comment.
This sounds concerning.
This comment has been minimized.
This comment has been minimized.
GitHub Actions
Error
GitHub Actions
Inactive
| entry = self._context_tokens.pop(key, None) | ||
| if entry is None: | ||
| return | ||
| thread_ident, token = entry | ||
| if thread_ident == threading.get_ident(): | ||
| otel_context.detach(token) # type: ignore[arg-type] |
There was a problem hiding this comment.
Codex AI review
[P1] Suspended user functions never receive an end hook, and invocation cleanup runs on the handler thread. This therefore pops a branch-worker token without detaching it. A timed-suspended map/parallel branch can reuse that worker and key, restoring the stale suspended span after resumption and losing the original token. Add a same-thread suspend/abort hook or wrapper cleanup that detaches before the worker returns, apply it to both plugins, and test a real timed suspend/resume through the concurrent executor.
| if thread_ident == threading.get_ident(): | ||
| context.detach(token) # type: ignore[arg-type] |
There was a problem hiding this comment.
Codex AI review
[P2] Detaching here assumes this plugin owns the current topmost scope. PluginExecutor supports multiple plugins and invokes both start and end hooks in registration order, so configuring otel-invocation,otel-execution attaches A then B but detaches A then B. B's detach consequently restores A's ended span, corrupting subsequent log and span correlation. Dispatch paired end hooks in reverse registration order, or otherwise coordinate scope unwinding, and add a two-plugin lifecycle test.
Codex AI reviewTwo context-lifecycle regressions remain. Coverage does not exercise real worker reuse or multiple OTel plugins together. Reviewed commit |
2 participants
Problem
Both OTel plugins called
opentelemetry.context.attach()and discarded the returned token, so a workflow, invocation, operation, or attempt span stayed current after its scope ended.on_user_function_endthen "restored" the enclosing span by attaching again, so every operation pushed another unbalanced scope onto the thread's context stack. The test suites reset the global OTel context around every test to hide the leak.Resolves #643
Change
Each plugin now records the token for every context it attaches, keyed by the same registry key used for the span, and releases it at the matching lifecycle end:
on_user_function_enddetaches its own scope instead of attaching the enclosing span again. Detaching restores the exact enclosing context, which is what the re-attach was approximating.ExecutionOtelPlugin's invocation-start scope is released during invocation cleanup, so the context active before the invocation is restored — this is the case the issue's runtime probe observed.SuspendExecutionskipson_user_function_end) and an invocation that ends abnormally, so nothing survives into a warm invocation.No new classes; the helper trio lives on each plugin next to its existing span-registry helpers.
Behavior delta
Between top-level operations,
trace.get_current_span()on the worker thread is now the restored enclosing context rather than the plugin's Invocation span. Emitted log correlation is unchanged:get_current_span_context()already falls back to the invocation span in the plugin registry, which is the documented path for top-level handler code. Verified that the set of exported spans is byte-identical before and after the change.Tests
_reset_otel_contextfixtures are replaced with an autouse fixture that asserts each test leaves the OTel context as it found it, so a lifecycle leak fails the suite instead of being masked. That fixture immediately caught six pre-existing mid-lifecycle tests; each now ends its invocation, which exercises the new cleanup path.InvocationOtelPlugintests that asserted the old stacked behavior were rewritten to assert enclosing-context restoration whileget_current_span_context()still resolves the invocation span.Verification
hatch run dev-otel:test— 135 passedhatch run test:all— 3233 passed, 2 skippedhatch run types:check— cleanhatch fmt --check(otel package) — cleandurable_execution()(outer step + child context + inner step, two invocations) against both plugins. Before the change it reproduced the issue: context not restored, a plugin span still current after the handler returned. After: context restored, no span left current, no tokens open, zeroFailed to detach contextlogs, identical exported spans.