From 4e59e90afd45853affd05590dc4dda3a040c6388 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Mon, 17 Aug 2026 22:12:00 +0000 Subject: [PATCH] fix(otel): balance context attach and detach Both 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: the invocation-start attach runs 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. Each attach now keeps its token, keyed the same way as the span registry, and the hook that pairs with it pops the scope: - on_invocation_start attaches; on_invocation_end detaches. - on_user_function_start attaches; on_user_function_end detaches, replacing the re-attach that stood in for restoring the enclosing context. - Anything still held when the invocation ends is swept, since the SDK re-raises SuspendExecution without calling on_user_function_end. A scope is only detached while it is still the current one. That mirrors OpenTelemetry Java's ScopeImpl.close(), which ignores a close that does not represent the current context, and it matters more in Python: ContextVar.reset writes back its captured value unconditionally, so an out-of-order or wrong-thread detach would revive a stale context instead of failing safe. A skipped detach keeps its entry so the owning thread can still undo it. 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, and the tests that drove hooks without completing the lifecycle now complete it. Adds coverage for warm invocation reuse keeping two executions in separate traces, a suspended operation's scope being swept, sequential steps not accumulating layers, exact restore on success and failure, nested child contexts, worker-thread confinement, and the identity guard skipping both out-of-order and cross-thread detaches. --- .../README.md | 6 + .../execution_plugin.py | 107 ++++++- .../invocation_plugin.py | 94 ++++++- .../tests/test_execution_plugin.py | 261 +++++++++++++++++- .../test_execution_plugin_integration.py | 18 +- .../tests/test_invocation_plugin.py | 89 ++++-- .../test_invocation_plugin_integration.py | 18 +- .../tests/test_log_filter.py | 23 ++ 8 files changed, 536 insertions(+), 80 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 4eb563aa..71cf07b1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -235,6 +235,12 @@ context onto every emitted log record using these attributes: These attributes are only set when a valid span context is active, so any log formatter or schema must treat the fields as optional. +Between two operations the plugin holds no span current: the scope a step or child +context attached is detached when that function returns. Log correlation is +unaffected -- the filter resolves the trace context from the plugin's own span +registry, so records emitted between operations still carry the invocation's +`traceId` and `spanId`. + ## Verification After deploying your function with the plugin configured: 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 bc829db6..03f5984d 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 @@ -23,6 +23,7 @@ import datetime import logging import threading +from contextvars import Token from typing import Any from aws_durable_execution_sdk_python.lambda_service import ( @@ -142,11 +143,70 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._workflow_span: Span | None = None self._invocation_span: Span | None = None self._operation_spans: dict[str, Span] = {} + # Contexts this plugin has attached, keyed the same way as the span + # registry, so each attach can be undone by the hook that pairs with it. + self._scopes: dict[str, tuple[Token[Context], Context]] = {} self._lock = threading.RLock() if self._config.enrich_logger: install_log_filter(self) + # ------------------------------------------------------------------ + # Context scopes + # ------------------------------------------------------------------ + def _enter_scope(self, key: str, context: Context) -> None: + """Attach ``context`` and remember what is needed to restore it. + + ``otel_context.attach`` returns a token that is the only way to undo it, + and the hook that attaches is not the hook that pops, so the token has to + be kept. The context is kept alongside it for the identity check in + :meth:`_exit_scope`. + """ + with self._lock: + self._scopes[key] = (otel_context.attach(context), context) + + def _exit_scope(self, key: str) -> None: + """Restore the context that preceded the scope attached under ``key``. + + Only detaches when the scope being popped is still the current one. This + mirrors OpenTelemetry Java's ``ScopeImpl.close()``, which ignores a close + that does not represent the current context, and it matters more here: + ``ContextVar.reset`` writes back its captured value unconditionally, so an + out-of-order or wrong-thread detach would *revive* a stale context instead + of failing safe. Skipping leaves the layer attached, which is inert. + """ + with self._lock: + entry = self._scopes.get(key) + if entry is None: + return + token, context = entry + if otel_context.get_current() is not context: + # Not ours to pop right now: another scope is stacked above it, or + # this is not the thread that attached it. The entry is left in + # place so it can still be undone later -- discarding the token + # here would strand the context permanently. + logger.debug("Skipping out-of-scope OTel context detach for %s", key) + return + del self._scopes[key] + try: + otel_context.detach(token) + except Exception: # noqa: BLE001 - observability must not break execution + logger.debug("Failed to detach OTel context for %s", key, exc_info=True) + + def _exit_all_scopes(self) -> None: + """Pop every scope this plugin still holds, newest first. + + Reached when a hook that would have popped a scope never ran: the SDK + re-raises ``SuspendExecution`` without calling ``on_user_function_end``. + Scopes attached on another thread fail the identity check and are dropped + without detaching; those threads are per-invocation and their context dies + with them. + """ + with self._lock: + keys = list(reversed(self._scopes)) + for key in keys: + self._exit_scope(key) + # ------------------------------------------------------------------ # Span registry helpers # ------------------------------------------------------------------ @@ -168,6 +228,17 @@ def _pop_span(self, key: str) -> Span | None: def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: return f"{info.operation_id}:attempt:{info.attempt or 1}" + @classmethod + def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: + """Return the context-scope key for a user-function hook pair. + + Mirrors the span registry key so the scope attached by + ``on_user_function_start`` is the one ``on_user_function_end`` pops. + """ + if info.operation_type is OperationType.STEP: + return cls._attempt_key(info) + return info.operation_id + def get_current_span_context(self) -> SpanContext | None: """Return the active span context for log correlation (see log_filter).""" span_context = trace.get_current_span().get_span_context() @@ -214,10 +285,16 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._start_invocation_span(info) # Make the Workflow span the active span so auto-instrumented spans - # created during the invocation become its children. + # created during the invocation become its children. Paired with the + # _exit_scope in on_invocation_end: this thread is the Lambda handler + # thread, which is reused across warm invocations, so leaving it attached + # let the next execution's context extractor and ambient-parent lookup + # adopt this execution's ended Workflow span -- merging two unrelated + # executions into one trace. if self._workflow_span is not None: - otel_context.attach( - trace.set_span_in_context(self._workflow_span, self._extracted_context) + self._enter_scope( + _INVOCATION_KEY, + trace.set_span_in_context(self._workflow_span, self._extracted_context), ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: @@ -329,6 +406,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: logger.exception("force_flush failed at invocation end") def _reset_state(self) -> None: + # Undo the invocation scope, and anything a suspended operation left + # behind, so the handler thread is returned to the state it was found in. + self._exit_all_scopes() self._execution_arn = "" self._extracted_context = None self._workflow_span = None @@ -455,7 +535,10 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: parent=parent, start_time=info.start_time, ) - otel_context.attach(trace.set_span_in_context(span, self._extracted_context)) + self._enter_scope( + self._scope_key(info), + trace.set_span_in_context(span, self._extracted_context), + ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: logger.debug("Durable user function ended: %s", info) @@ -463,6 +546,11 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end only supports CONTEXT and STEP operations" ) + # Pop the scope this operation attached, restoring exactly what preceded + # it. Detaching rather than attaching the enclosing span again is what + # keeps this balanced: the previous code pushed a second context here, so + # every operation added a layer and removed none. + self._exit_scope(self._scope_key(info)) key = ( self._attempt_key(info) if info.operation_type is OperationType.STEP @@ -496,17 +584,6 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: if popped is not None: popped.end(end_time=_to_otel_timestamp(end_time)) - # Restore the enclosing span as active (parent op, else invocation/workflow). - enclosing = ( - self._get_span(info.parent_id) - or self._invocation_span - or self._workflow_span - ) - if enclosing is not None: - otel_context.attach( - trace.set_span_in_context(enclosing, self._extracted_context) - ) - # ------------------------------------------------------------------ # Attributes # ------------------------------------------------------------------ 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 4f42ca32..d6cf65e9 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 @@ -5,6 +5,7 @@ import datetime import logging import threading +from contextvars import Token from typing import Any from aws_durable_execution_sdk_python.lambda_service import ( @@ -167,6 +168,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._workflow_span: Span | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} + # Contexts this plugin has attached, keyed the same way as the span + # registry, so each attach can be undone by the hook that pairs with it. + self._scopes: dict[str, tuple[Token[Context], Context]] = {} self._operation_spans_lock = threading.RLock() if self._enrich_logger: @@ -176,6 +180,59 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # plugin is constructed), so the handlers are available here. install_log_filter(self) + def _enter_scope(self, key: str, context_to_attach: Context) -> None: + """Attach a context and remember what is needed to restore it. + + ``context.attach`` returns a token that is the only way to undo it, and + the hook that attaches is not the hook that pops, so the token has to be + kept. The context is kept alongside it for the identity check in + :meth:`_exit_scope`. + """ + with self._operation_spans_lock: + self._scopes[key] = (context.attach(context_to_attach), context_to_attach) + + def _exit_scope(self, key: str) -> None: + """Restore the context that preceded the scope attached under ``key``. + + Only detaches when the scope being popped is still the current one. This + mirrors OpenTelemetry Java's ``ScopeImpl.close()``, which ignores a close + that does not represent the current context, and it matters more here: + ``ContextVar.reset`` writes back its captured value unconditionally, so an + out-of-order or wrong-thread detach would *revive* a stale context instead + of failing safe. Skipping leaves the layer attached, which is inert. + """ + with self._operation_spans_lock: + entry = self._scopes.get(key) + if entry is None: + return + token, attached = entry + if context.get_current() is not attached: + # Not ours to pop right now: another scope is stacked above it, or + # this is not the thread that attached it. The entry is left in + # place so it can still be undone later -- discarding the token + # here would strand the context permanently. + logger.debug("Skipping out-of-scope OTel context detach for %s", key) + return + del self._scopes[key] + try: + context.detach(token) + except Exception: # noqa: BLE001 - observability must not break execution + logger.debug("Failed to detach OTel context for %s", key, exc_info=True) + + def _exit_all_scopes(self) -> None: + """Pop every scope this plugin still holds, newest first. + + Reached when a hook that would have popped a scope never ran: the SDK + re-raises ``SuspendExecution`` without calling ``on_user_function_end``. + Scopes attached on another thread fail the identity check and are dropped + without detaching; those threads are per-invocation and their context dies + with them. + """ + with self._operation_spans_lock: + keys = list(reversed(self._scopes)) + for key in keys: + self._exit_scope(key) + def _set_span(self, operation_id: str | None, span: Span) -> None: """Register the active span for an operation ID.""" with self._operation_spans_lock: @@ -196,6 +253,17 @@ def _attempt_span_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: """Return the registry key for a STEP attempt span.""" return f"{info.operation_id}:attempt:{info.attempt or 1}" + @classmethod + def _scope_key(cls, info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: + """Return the context-scope key for a user-function hook pair. + + Mirrors the span registry key so the scope attached by + ``on_user_function_start`` is the one ``on_user_function_end`` pops. + """ + if info.operation_type is OperationType.STEP: + return cls._attempt_span_key(info) + return info.operation_id + def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. @@ -453,6 +521,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._workflow_span.set_status(StatusCode.OK) self._workflow_span.end() + # Undo anything a suspended operation left attached, so no scope outlives + # the invocation that created it. + self._exit_all_scopes() + # Clear all per-invocation state to prevent leaks across warm Lambda reuses self._execution_arn = "" self._extracted_context = None @@ -562,7 +634,10 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: span_key=span_key, deterministic_span_id=info.operation_type is not OperationType.STEP, ) - context.attach(trace.set_span_in_context(span, self._extracted_context)) + self._enter_scope( + self._scope_key(info), + trace.set_span_in_context(span, self._extracted_context), + ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: """Called when a context or step operation finishes user code. @@ -578,6 +653,13 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end should only be called for CONTEXT and STEP operations" ) + # Pop the scope this operation attached, restoring exactly what preceded + # it. Detaching rather than attaching the enclosing span again is what + # keeps this balanced: the previous code pushed a second context here, so + # every operation added a layer and removed none. Between operations the + # log filter resolves through the span registry (see + # get_current_span_context), so correlation is unaffected. + self._exit_scope(self._scope_key(info)) # key = f"{info.operation_id}-{int(info.start_time.timestamp())}" span_key = ( self._attempt_span_key(info) @@ -610,16 +692,6 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: if end_timestamp is not None and end_timestamp == info.start_time: end_timestamp += datetime.timedelta(microseconds=1) self._end_span(span_key, end_timestamp) - # Restore the enclosing operation span as current so code that runs - # after this operation (e.g. between steps in a child context) - # correlates to its enclosing operation, not the operation that just - # ended. For a top-level operation (parent_id is None) this is the - # invocation span; for a nested operation it is the parent context span. - parent_span = self._get_span(info.parent_id) or self._get_span(None) - if parent_span: - context.attach( - trace.set_span_in_context(parent_span, self._extracted_context) - ) def _extract_attributes(self, info: Any) -> _SpanAttributes: """Extract durable execution fields as OpenTelemetry span attributes. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 75488662..10fcd24e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime import opentelemetry.context as otel_context @@ -28,11 +29,17 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from aws_durable_execution_sdk_python_otel.context_extractors import ( + xray_context_extractor, +) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, operation_id_to_span_id, ) -from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + _INVOCATION_KEY, + ExecutionOtelPlugin, +) from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( OtelPluginConfig, ProviderSource, @@ -45,17 +52,18 @@ @pytest.fixture(autouse=True) -def _reset_otel_context(): - """Reset the OTel thread-local context around each test. +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context attached. - The plugin attaches spans via context.attach() without detaching, so state - would otherwise leak between tests running on the same thread. + The plugins must detach every context they attach, so no reset is needed to + isolate tests -- instead this asserts the invariant. Resetting here would hide + exactly the leak this suite exists to catch. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + yield + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: @@ -338,6 +346,10 @@ def test_context_span_waits_for_terminal_operation_status( assert span.attributes["durable.operation.status"] == terminal_status.value assert span.status.status_code is expected_span_status + # Close the invocation so its scope is detached; the autouse + # fixture asserts no context outlives the test. + plugin.on_invocation_end(_invocation_end_info()) + def test_step_attempt_span_omits_operation_status(): plugin, exporter = _create_plugin() @@ -387,10 +399,15 @@ def test_step_attempt_span_omits_operation_status(): ) assert "durable.operation.status" not in span.attributes + # --------------------------------------------------------------------------- + # Default-provider mode: invocation span + # --------------------------------------------------------------------------- + + # Close the invocation so its scope is detached; the autouse + # fixture asserts no context outlives the test. + plugin.on_invocation_end(_invocation_end_info()) + -# --------------------------------------------------------------------------- -# Default-provider mode: invocation span -# --------------------------------------------------------------------------- def _create_default_mode_plugin( monkeypatch, ) -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: @@ -446,6 +463,224 @@ def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch): assert invocation.parent.span_id == ambient.get_span_context().span_id +# ---------------------------------------------------------------------- +# Invocation scope is paired, so nothing leaks into the next invocation +# ---------------------------------------------------------------------- +def test_invocation_end_restores_the_pre_invocation_context(): + """Verify the invocation scope is detached when the invocation ends. + + The Lambda handler thread is reused across warm invocations, so an unpaired + attach here left an ended Workflow span current for the next execution. + """ + plugin, _ = _create_plugin() + before = otel_context.get_current() + + plugin.on_invocation_start(_invocation_start_info()) + # The Workflow span is current while the invocation runs. + assert ( + trace.get_current_span().get_span_context().span_id + == plugin._workflow_span.get_span_context().span_id + ) + + plugin.on_invocation_end(_invocation_end_info()) + + assert otel_context.get_current() is before + assert trace.get_current_span().get_span_context().is_valid is False + + +def test_a_suspended_operation_scope_is_swept_at_invocation_end(): + """Verify a scope whose end hook never ran does not outlive the invocation. + + The SDK re-raises SuspendExecution without calling on_user_function_end, so + the operation's scope is still attached when the invocation winds down. + """ + plugin, _ = _create_plugin() + before = otel_context.get_current() + plugin.on_invocation_start(_invocation_start_info()) + + plugin.on_user_function_start( + UserFunctionStartInfo( + operation_id="step-suspends", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="step-suspends", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + ) + assert len(plugin._scopes) == 2 # invocation + operation + + plugin.on_invocation_end(_invocation_end_info(InvocationStatus.PENDING)) + + assert plugin._scopes == {} + assert otel_context.get_current() is before + + +def test_warm_reuse_does_not_share_a_trace_between_executions(monkeypatch): + """Verify a reused plugin instance keeps two executions in separate traces. + + In GLOBAL mode the Invocation span is parented to whatever is ambient. When a + previous invocation left its Workflow span attached, that span became the + parent and its trace ID won, merging two unrelated executions into one trace. + """ + plugin, _ = _create_default_mode_plugin(monkeypatch) + # The default X-Ray extractor falls back to the ambient context when no trace + # header is present, so it observes any leak too. + monkeypatch.delenv("_X_AMZN_TRACE_ID", raising=False) + plugin._context_extractor = xray_context_extractor + + traces: list[int] = [] + + def run(arn: str) -> None: + plugin.on_invocation_start( + InvocationStartInfo( + request_id="request-1", + execution_arn=arn, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + ) + traces.append(plugin._invocation_span.get_span_context().trace_id) + assert plugin._invocation_span.parent is None, ( + "the Invocation span adopted a parent from a previous invocation" + ) + plugin.on_invocation_end( + InvocationEndInfo( + request_id="request-1", + execution_arn=arn, + execution_start_time=START_TIME, + is_first_invocation=True, + status=InvocationStatus.SUCCEEDED, + error=None, + ) + ) + + run(EXECUTION_ARN) + run(EXECUTION_ARN + "-second") + + assert traces[0] != traces[1], "two executions were merged into one trace" + + +# ---------------------------------------------------------------------- +# The identity guard: a detach that is not the current scope is skipped +# ---------------------------------------------------------------------- +def test_out_of_order_exit_is_skipped_rather_than_reviving_a_stale_context(): + """Verify a mismatched detach leaves the context alone. + + ``ContextVar.reset`` writes back its captured value unconditionally, so + detaching a scope that is no longer current would *revive* what preceded it, + silently making a stale span current again. The guard makes that a no-op, + matching OpenTelemetry Java's ``ScopeImpl.close()``. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + outer = otel_context.get_current() + + # A second scope stacked on top of the invocation scope. + plugin._enter_scope("inner", trace.set_span_in_context(plugin._invocation_span)) + inner = otel_context.get_current() + assert inner is not outer + + # Popping the *outer* scope while the inner one is current must not revive + # the pre-invocation context. The entry is kept, so it can still be undone + # once it is current again. + plugin._exit_scope(_INVOCATION_KEY) + assert otel_context.get_current() is inner + assert _INVOCATION_KEY in plugin._scopes + + plugin._exit_scope("inner") + assert otel_context.get_current() is outer + # And the retained outer scope is undone at invocation end. + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._scopes == {} + + +def test_exit_from_another_thread_is_skipped(): + """Verify a scope is never detached from a thread that did not attach it. + + A token can only be reset in the context that created it, so a cross-thread + detach would corrupt the calling thread. The identity check rejects it because + the other thread's current context is not the attached one. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + attached = otel_context.get_current() + observed: dict[str, object] = {} + + def worker() -> None: + # This thread never attached anything, so its context does not match. + plugin._exit_scope(_INVOCATION_KEY) + observed["worker_valid"] = trace.get_current_span().get_span_context().is_valid + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(worker).result() + + assert observed["worker_valid"] is False + # The attaching thread is untouched, and the scope was not consumed by the + # failed attempt -- so the thread that owns it can still undo it. + assert otel_context.get_current() is attached + assert _INVOCATION_KEY in plugin._scopes + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._scopes == {} + + +def test_a_worker_thread_scope_does_not_disturb_the_caller(): + """Verify user-function scopes stay on the thread that runs user code. + + User code runs on a worker the SDK owns, and ThreadPoolExecutor does not copy + contextvars, so the plugin's scope must not reach the handler thread. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before = otel_context.get_current() + observed: dict[str, object] = {} + + def run_step() -> None: + plugin.on_user_function_start( + UserFunctionStartInfo( + operation_id="step-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="step-1", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + ) + observed["inside"] = trace.get_current_span().get_span_context().is_valid + plugin.on_user_function_end( + UserFunctionEndInfo( + operation_id="step-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="step-1", + parent_id=None, + start_time=START_TIME, + end_time=END_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + is_replay_children=False, + attempt=1, + outcome=UserFunctionOutcome.SUCCEEDED, + error=None, + ) + ) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_step).result() + + assert observed["inside"] is True + assert otel_context.get_current() is before + plugin.on_invocation_end(_invocation_end_info()) + + def test_open_operation_span_not_exported_at_invocation_end(): """A suspended operation (started, not ended) must not be exported. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index 98f450c5..e9ed2cd1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -60,12 +60,18 @@ @pytest.fixture(autouse=True) -def _reset_otel_context(): - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context attached. + + The plugins must detach every context they attach, so no reset is needed to + isolate tests -- instead this asserts the invariant. Resetting here would hide + exactly the leak this suite exists to catch. + """ + before = otel_context.get_current() + yield + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index d6cf063a..bd00642f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -48,17 +48,18 @@ @pytest.fixture(autouse=True) -def _reset_otel_context(): - """Reset the OTel thread-local context before and after each test. +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context attached. - The plugin attaches spans via context.attach() without ever detaching, - so state would otherwise leak between tests running on the same thread. + The plugins must detach every context they attach, so no reset is needed to + isolate tests -- instead this asserts the invariant. Resetting here would hide + exactly the leak this suite exists to catch. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + yield + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -791,17 +792,16 @@ def update_span(index: int) -> None: # ---------------------------------------------------------------------- -# on_user_function_end restores the invocation span to the context +# on_user_function_end restores the context the operation was entered from # ---------------------------------------------------------------------- -def test_user_function_end_restores_invocation_span(): - """Verify the invocation span is current again after a step completes.""" +def test_user_function_end_restores_enclosing_context(): + """Verify the exact pre-step context is restored after a step completes.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - invocation_span_id = plugin._get_span(None).get_span_context().span_id + before = otel_context.get_current() operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) - # Inside the step, the current span is the attempt span. active_attempt_span = plugin._get_span("step-1:attempt:1") assert active_attempt_span is not None assert ( @@ -811,15 +811,17 @@ def test_user_function_end_restores_invocation_span(): plugin.on_user_function_end(_user_function_end_info(operation_id)) - # After the step, the invocation span is restored. - assert trace.get_current_span().get_span_context().span_id == invocation_span_id + # Detached, not re-attached: the context is restored byte for byte. Between + # operations the log filter resolves through the registry instead. + assert otel_context.get_current() is before + plugin.on_invocation_end(_invocation_end_info()) -def test_user_function_end_restores_invocation_span_on_failure(): - """Verify the invocation span is restored even when the step fails.""" +def test_user_function_end_restores_enclosing_context_on_failure(): + """Verify the context is restored even when the step fails.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - invocation_span_id = plugin._get_span(None).get_span_context().span_id + before = otel_context.get_current() operation_id = "step-fail" plugin.on_user_function_start(_user_function_start_info(operation_id)) @@ -827,21 +829,26 @@ def test_user_function_end_restores_invocation_span_on_failure(): _user_function_end_info(operation_id, outcome=UserFunctionOutcome.FAILED) ) - assert trace.get_current_span().get_span_context().span_id == invocation_span_id + assert otel_context.get_current() is before + plugin.on_invocation_end(_invocation_end_info()) + +def test_sequential_steps_do_not_accumulate_context_layers(): + """Verify N sequential steps leave no residue. -def test_user_function_end_restores_invocation_span_across_multiple_steps(): - """Verify between-step context is the invocation span across many steps.""" + The previous code attached the enclosing span again on every end hook, so a + handler running many steps grew one context layer per step. + """ plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - invocation_span_id = plugin._get_span(None).get_span_context().span_id + before = otel_context.get_current() for index in range(3): operation_id = f"step-{index}" plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # Between each step, the invocation span is the current span. - assert trace.get_current_span().get_span_context().span_id == invocation_span_id + assert otel_context.get_current() is before + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -878,6 +885,9 @@ def test_get_current_span_context_returns_operation_span_inside_step(): assert active_attempt_span is not None assert span_context.span_id == active_attempt_span.get_span_context().span_id + plugin.on_user_function_end(_user_function_end_info(operation_id)) + plugin.on_invocation_end(_invocation_end_info()) + def test_get_current_span_context_returns_invocation_span_between_steps(): """Verify between-step code resolves back to the invocation span context.""" @@ -936,18 +946,29 @@ def test_user_function_end_restores_parent_context_span_for_nested_step(): != plugin._get_span(None).get_span_context().span_id ) + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_top_level_step_restores_the_ambient_context(): + """Verify a top-level step (parent_id=None) restores the ambient context. -def test_user_function_end_falls_back_to_invocation_when_parent_missing(): - """Verify a top-level step (parent_id=None) restores the invocation span.""" + There is no enclosing operation scope to land on, so the thread returns to + whatever preceded the step. Log correlation for that window is covered by + get_current_span_context()'s invocation-span fallback. + """ plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) - invocation_span_id = plugin._get_span(None).get_span_context().span_id + before = otel_context.get_current() operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - assert trace.get_current_span().get_span_context().span_id == invocation_span_id + assert otel_context.get_current() is before + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_context_span_between_nested_steps(): @@ -979,6 +1000,11 @@ def test_get_current_span_context_returns_context_span_between_nested_steps(): assert span_context.span_id == context_span.get_span_context().span_id assert span_context.span_id != plugin._get_span(None).get_span_context().span_id + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + def test_nested_steps_restore_context_span_across_multiple_iterations(): """Verify each inner step restores the child-context span between iterations.""" @@ -1002,6 +1028,11 @@ def test_nested_steps_restore_context_span_across_multiple_iterations(): # Between each inner step, the child-context span is current. assert trace.get_current_span().get_span_context().span_id == context_span_id + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + @pytest.mark.parametrize( ("status", "expected_code"), diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py index c1c6d4b1..a1292433 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -63,12 +63,18 @@ @pytest.fixture(autouse=True) -def _reset_otel_context(): - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context attached. + + The plugins must detach every context they attach, so no reset is needed to + isolate tests -- instead this asserts the invariant. Resetting here would hide + exactly the leak this suite exists to catch. + """ + before = otel_context.get_current() + yield + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index f419f9ad..f166a485 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -11,6 +11,8 @@ ) from aws_durable_execution_sdk_python.plugin import ( InvocationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, UserFunctionStartInfo, ) from opentelemetry.context import Context @@ -77,6 +79,25 @@ def _user_function_start_info(operation_id: str) -> UserFunctionStartInfo: ) +def _user_function_end_info(operation_id: str) -> UserFunctionEndInfo: + """Create standard user function end info for tests.""" + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=None, + name="fetch-user", + parent_id=None, + start_time=START_TIME, + end_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + is_replay_children=False, + attempt=1, + outcome=UserFunctionOutcome.SUCCEEDED, + error=None, + ) + + def _make_record() -> logging.LogRecord: """Create a bare LogRecord for filtering.""" return logging.LogRecord( @@ -148,6 +169,8 @@ def test_filter_uses_attempt_span_inside_user_function(): expected_span_id = format(attempt_span.get_span_context().span_id, "016x") assert record.spanId == expected_span_id + plugin.on_user_function_end(_user_function_end_info(operation_id)) + def test_install_log_filter_attaches_to_handlers(): """install_log_filter adds the filter to each handler on the target logger."""