diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 4eb563aa..12201792 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -235,6 +235,23 @@ 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. +### Active Span Scope + +The plugin makes a span current only while your step or child-context function +runs, and detaches it when that function returns. Two consequences are worth +knowing: + +- Auto-instrumented calls (botocore, urllib3, and similar) made **inside** a step + or child context become children of that operation's span. +- Auto-instrumented calls made **outside** any operation -- for example directly + in the handler between two steps -- are not parented to the durable spans. In + an ADOT deployment they attach to the ambient Lambda invocation span instead. + Put such work in a step if you need it inside the durable trace. + +Log correlation is unaffected either way: the logging 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/context_scope.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py new file mode 100644 index 00000000..f36b238b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_scope.py @@ -0,0 +1,224 @@ +"""Balanced ``opentelemetry.context`` attach/detach bookkeeping for the plugins. + +The OpenTelemetry Context specification requires every ``context.attach()`` to +have a corresponding ``context.detach(token)``. Detaching is only possible with +the token that ``attach`` returned, so the token has to survive from the hook +that attached to the hook that pops it -- the plugin hooks are separate calls, +so the idiomatic ``with tracer.start_as_current_span(...)`` form is unavailable. + +Two properties of the runtime shape the design: + +* **Tokens are thread-confined.** The plugin hooks run on several threads: the + invocation hooks on the Lambda handler thread, the user-function hooks on the + ``dex-handler`` worker that runs user code, and on a branch worker for each + ``map``/``parallel`` branch. ``ContextVar.reset()`` only accepts a token + created in the same ``contextvars.Context``, so each thread keeps its own + stack and only ever detaches its own tokens. +* **Detach order matters.** Unlike OpenTelemetry Java's ``Scope.close()`` -- + which ignores a close that does not represent the current context -- + ``ContextVar.reset()`` unconditionally writes back the token's captured value. + Detaching out of order therefore *revives* a stale context instead of failing + safe. The stack is module level rather than per plugin instance so that two + plugins attaching on the same thread (both ship as separate entry points and + can be enabled together) still unwind in true LIFO order. + +Scopes are keyed by ``(owner, key)`` so a plugin instance can pop the exact +scope it pushed, while :func:`exit_scope` still unwinds anything stacked above +it. Nothing here raises: a plugin must never break an execution over +observability bookkeeping. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable + + +if TYPE_CHECKING: + from contextvars import Token + + from opentelemetry.context import Context + + +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class _Entry: + """One attached scope: who pushed it, under what key, and its token.""" + + owner_id: int + key: str + epoch: int + token: Token[Context] + + +class _ThreadState(threading.local): + """Per-thread LIFO stack of attached scopes.""" + + def __init__(self) -> None: + self.entries: list[_Entry] = [] + + +_state = _ThreadState() + + +def _detach(entry: _Entry) -> None: + """Detach one entry, swallowing any failure.""" + from opentelemetry import context as otel_context + + try: + otel_context.detach(entry.token) + except Exception: # noqa: BLE001 - observability must not break execution + logger.debug("Failed to detach OTel context scope %s", entry.key, exc_info=True) + + +def enter_scope( + owner: Any, + key: str, + context_factory: Callable[[], Context], + epoch: int = 0, +) -> None: + """Attach a context on this thread and remember how to restore it. + + Scopes left over from an earlier ``epoch``, or from an earlier entry under the + same ``key``, are unwound first: the SDK re-raises ``SuspendExecution`` without + calling ``on_user_function_end``, so a suspended operation leaves its scope + attached. + + ``context_factory`` is called *after* that cleanup, not before. The context to + attach is normally derived from what is current, so building it first would + copy values from a scope that is about to be detached -- baggage, suppression + flags -- and detaching afterwards cannot remove them from a context that has + already been constructed. + + Args: + owner: The plugin instance pushing the scope. + key: Registry key for the scope, unique per owner (operation or attempt). + context_factory: Builds the context to attach, called after cleanup. + epoch: The owner's invocation counter; scopes from older epochs are + discarded before the new scope is pushed. + """ + from opentelemetry import context as otel_context + + owner_id = id(owner) + _discard_stale(owner_id, epoch) + _discard_reentered(owner_id, key) + try: + token = otel_context.attach(context_factory()) + except Exception: # noqa: BLE001 + logger.debug("Failed to attach OTel context scope %s", key, exc_info=True) + return + _state.entries.append(_Entry(owner_id=owner_id, key=key, epoch=epoch, token=token)) + + +def exit_scope(owner: Any, key: str) -> None: + """Detach the scope ``owner`` pushed under ``key``, restoring what preceded it. + + Scopes stacked above the target are detached first so the underlying + ``ContextVar`` is always reset in LIFO order. A key this thread never pushed + is a no-op -- the scope belongs to another thread (or was already unwound), + and detaching someone else's token would corrupt the context. + """ + owner_id = id(owner) + index = _find_last(owner_id, key) + if index is None: + return + for entry in reversed(_state.entries[index:]): + _detach(entry) + del _state.entries[index:] + + +def unwind(owner: Any) -> None: + """Detach every scope ``owner`` still holds on this thread, newest first. + + Called at invocation end so the handler thread is left exactly as the plugin + found it. Scopes this owner pushed on *other* threads cannot be detached from + here; those threads are created per invocation and their context dies with + them. + """ + owner_id = id(owner) + index = _find_first(owner_id) + if index is None: + return + for entry in reversed(_state.entries[index:]): + _detach(entry) + del _state.entries[index:] + + +def depth(owner: Any | None = None) -> int: + """Return the number of scopes attached on this thread (for tests).""" + if owner is None: + return len(_state.entries) + owner_id = id(owner) + return sum(1 for entry in _state.entries if entry.owner_id == owner_id) + + +def _discard_reentered(owner_id: int, key: str) -> None: + """Unwind a scope this owner already holds under ``key`` on this thread. + + The epoch check only catches a *previous invocation's* leftovers. The same + operation key can also be entered twice inside one invocation, when a + suspended operation is re-entered after its branch is resubmitted, and its + first scope is still attached because the suspending path had no end hook to + pop it. Without this, the second enter would stack on the first and the + eventual end hook -- which pops one scope -- would leave the original + attached. + + A scope abandoned by a *different* operation on this thread cannot be + detected here. Physical nesting is not derivable from the hook payloads: + ``parent_id`` is checkpoint hierarchy, and a FLAT map/parallel branch + deliberately reports its inner operations' parent as the grandparent (see + ``DurableContext.is_virtual``), so a live branch scope would be + indistinguishable from an abandoned sibling. Closing that gap needs the SDK + to report the end of a suspended user function, which it does not do today. + """ + index = next( + ( + position + for position, entry in enumerate(_state.entries) + if entry.owner_id == owner_id and entry.key == key + ), + None, + ) + if index is None: + return + for entry in reversed(_state.entries[index:]): + _detach(entry) + del _state.entries[index:] + + +def _discard_stale(owner_id: int, epoch: int) -> None: + """Unwind this owner's scopes left over from a previous epoch.""" + index = next( + ( + position + for position, entry in enumerate(_state.entries) + if entry.owner_id == owner_id and entry.epoch != epoch + ), + None, + ) + if index is None: + return + for entry in reversed(_state.entries[index:]): + _detach(entry) + del _state.entries[index:] + + +def _find_last(owner_id: int, key: str) -> int | None: + """Index of this owner's most recent scope for ``key``, if any.""" + for position in range(len(_state.entries) - 1, -1, -1): + entry = _state.entries[position] + if entry.owner_id == owner_id and entry.key == key: + return position + return None + + +def _find_first(owner_id: int) -> int | None: + """Index of this owner's oldest scope, if any.""" + for position, entry in enumerate(_state.entries): + if entry.owner_id == owner_id: + return position + return None 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..6decc311 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 @@ -51,6 +51,7 @@ Tracer, ) +from aws_durable_execution_sdk_python_otel import context_scope from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, xray_context_extractor, @@ -143,6 +144,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._invocation_span: Span | None = None self._operation_spans: dict[str, Span] = {} self._lock = threading.RLock() + # Bumped every invocation. context_scope uses it to discard scopes a + # previous invocation left attached on a reused thread. + self._epoch = 0 if self._config.enrich_logger: install_log_filter(self) @@ -168,11 +172,57 @@ 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 pushed 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 _scope_base_context(self) -> Context: + """Return the context a new operation scope should be layered onto. + + For the outermost durable scope on a thread, the extracted upstream + context is the base: user code runs on a worker whose context starts + empty, so anything the context extractor supplied -- a remote parent, + baggage -- would be dropped by using the current context there, and + downstream propagation inside steps would lose it. + + For a nested scope the current context already carries that extracted + context transitively, via the enclosing scope, so it is used as the base + to preserve whatever ran in between (ambient spans, baggage added by user + code). + """ + if context_scope.depth(self) > 0: + return otel_context.get_current() + # An extractor may deliberately return an empty Context to isolate the + # operation from whatever is ambient. Context subclasses dict, so an empty + # one is falsy -- test for None, or that intent silently inverts into + # inheriting the worker's ambient baggage and suppression values. + if self._extracted_context is not None: + return self._extracted_context + return otel_context.get_current() + 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() - if span_context and span_context.is_valid: - return span_context + """Return the active span context for log correlation (see log_filter). + + The current span is used only while this plugin holds an operation scope + on this thread -- inside a step or child context, where the current span + is the one this plugin attached. Otherwise the registry is used, so a + record emitted between operations, or on the handler thread, correlates to + the durable Invocation span rather than to whatever else happens to be + current. That distinction matters in GLOBAL (ADOT) mode: the ambient + Lambda span is current on the handler thread and would otherwise be + reported in place of the durable span. + """ + if context_scope.depth(self) > 0: + span_context = trace.get_current_span().get_span_context() + if span_context and span_context.is_valid: + return span_context for candidate in (self._invocation_span, self._workflow_span): if candidate is not None: ctx = candidate.get_span_context() @@ -204,6 +254,7 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None: # ------------------------------------------------------------------ def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) + self._epoch += 1 self._execution_arn = info.execution_arn or "" self._extracted_context = self._context_extractor(info) self._id_generator.set_trace_id(self._execution_arn, info.execution_start_time) @@ -213,12 +264,16 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: # is parented to the ambient Lambda invocation span. self._start_invocation_span(info) - # Make the Workflow span the active span so auto-instrumented spans - # created during the invocation become its children. - if self._workflow_span is not None: - otel_context.attach( - trace.set_span_in_context(self._workflow_span, self._extracted_context) - ) + # No context is attached here. Nothing on this thread needs it: user code + # runs on a separate worker (ThreadPoolExecutor does not copy + # contextvars), so an attach here would never reach it, while the Lambda + # handler thread is reused across warm invocations -- an unpaired attach + # would leak an ended span into the next execution, whose context + # extractor and ambient-parent lookup would then adopt it and merge two + # executions into one trace. The Workflow and Invocation spans are used + # as explicit parents instead (see _resolve_parent), matching the Java + # plugins, which never make either span current. Log correlation for this + # thread resolves through get_current_span_context(). def _start_workflow_span(self, info: InvocationStartInfo) -> None: if not self._execution_arn: @@ -329,6 +384,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: logger.exception("force_flush failed at invocation end") def _reset_state(self) -> None: + # Detach anything this plugin still holds on this thread so the handler + # thread is left exactly as it was found. Scopes attached on the + # per-invocation worker threads cannot be detached from here (a token is + # only resettable in the context that created it); those threads are + # destroyed with the invocation, and any scope a suspended operation left + # behind is discarded by the epoch check on the next enter_scope. + context_scope.unwind(self) self._execution_arn = "" self._extracted_context = None self._workflow_span = None @@ -455,7 +517,18 @@ 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)) + # Attach on this worker thread so auto-instrumented calls made by the + # user function become children of this span. The span's own parent was + # chosen explicitly in _start_span; this only sets what is ambient while + # the user function runs. + context_scope.enter_scope( + self, + self._scope_key(info), + # Built after enter_scope's cleanup so it cannot inherit values + # from a scope that is about to be detached. + lambda: trace.set_span_in_context(span, self._scope_base_context()), + epoch=self._epoch, + ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: logger.debug("Durable user function ended: %s", info) @@ -463,6 +536,12 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end only supports CONTEXT and STEP operations" ) + # Detach first, on the same thread that attached, so the context this + # operation was entered from is restored exactly. Detaching (rather than + # attaching the enclosing span again) is what keeps the scopes balanced: + # a nested operation lands back on its parent's still-attached scope, and + # a top-level one lands back on the thread's ambient context. + context_scope.exit_scope(self, self._scope_key(info)) key = ( self._attempt_key(info) if info.operation_type is OperationType.STEP @@ -496,17 +575,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..890ac49b 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 @@ -34,6 +34,7 @@ Tracer, ) +from aws_durable_execution_sdk_python_otel import context_scope from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, xray_context_extractor, @@ -168,6 +169,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} self._operation_spans_lock = threading.RLock() + # Bumped every invocation. context_scope uses it to discard scopes a + # previous invocation left attached on a reused thread. + self._epoch = 0 if self._enrich_logger: # Install the root-logger filter so every log record is stamped with @@ -196,26 +200,65 @@ 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 pushed 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 _scope_base_context(self) -> Context: + """Return the context a new operation scope should be layered onto. + + For the outermost durable scope on a thread, the extracted upstream + context is the base: user code runs on a worker whose context starts + empty, so anything the context extractor supplied -- a remote parent, + baggage -- would be dropped by using the current context there, and + downstream propagation inside steps would lose it. + + For a nested scope the current context already carries that extracted + context transitively, via the enclosing scope, so it is used as the base + to preserve whatever ran in between (ambient spans, baggage added by user + code). + """ + if context_scope.depth(self) > 0: + return context.get_current() + # An extractor may deliberately return an empty Context to isolate the + # operation from whatever is ambient. Context subclasses dict, so an empty + # one is falsy -- test for None, or that intent silently inverts into + # inheriting the worker's ambient baggage and suppression values. + if self._extracted_context is not None: + return self._extracted_context + return context.get_current() + def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. Resolution order: - 1. The span attached to the OTel thread-local context. Inside a step - this is the active attempt span, and inside a child context this is - the active context span (attached in - on_user_function_start), and between operations it is the enclosing - operation span (restored in on_user_function_end). - 2. The invocation span from the plugin registry. This is the path used - for top-level handler code: the invocation span is never attached to - the worker thread's context, so the registry is the only way to - resolve it. + 1. The span attached to the OTel thread-local context, but only while this + plugin holds an operation scope on this thread. Inside a step that is + the active attempt span, and inside a child context the active context + span. After a nested operation ends, its scope is detached and the + enclosing child context span -- still attached -- becomes current + again. + 2. The invocation span from the plugin registry. This covers top-level + handler code (the invocation span is never attached to any thread's + context) and code between top-level operations. Gating step 1 on an + owned scope matters in GLOBAL (ADOT) mode: the ambient Lambda span is + current on the handler thread and would otherwise be reported in place + of the durable span. Returns: A valid SpanContext, or None if no span is active. """ - span_context = trace.get_current_span().get_span_context() - if span_context and span_context.is_valid: - return span_context + if context_scope.depth(self) > 0: + span_context = trace.get_current_span().get_span_context() + if span_context and span_context.is_valid: + return span_context invocation_span = self._get_span(None) if invocation_span: @@ -364,6 +407,7 @@ def _end_span( def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) + self._epoch += 1 self._execution_arn = info.execution_arn or "" self._extracted_context = self._context_extractor(info) self._id_generator.set_trace_id(self._execution_arn, info.execution_start_time) @@ -454,6 +498,13 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._workflow_span.end() # Clear all per-invocation state to prevent leaks across warm Lambda reuses + # Detach anything this plugin still holds on this thread so the handler + # thread is left exactly as it was found. Scopes attached on the + # per-invocation worker threads cannot be detached from here (a token is + # only resettable in the context that created it); those threads are + # destroyed with the invocation, and any scope a suspended operation left + # behind is discarded by the epoch check on the next enter_scope. + context_scope.unwind(self) self._execution_arn = "" self._extracted_context = None self._workflow_span = None @@ -562,7 +613,18 @@ 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)) + # Attach on this worker thread so auto-instrumented calls made by the + # user function become children of this span. The span's own parent was + # chosen explicitly in _start_span; this only sets what is ambient while + # the user function runs. + context_scope.enter_scope( + self, + self._scope_key(info), + # Built after enter_scope's cleanup so it cannot inherit values + # from a scope that is about to be detached. + lambda: trace.set_span_in_context(span, self._scope_base_context()), + epoch=self._epoch, + ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: """Called when a context or step operation finishes user code. @@ -578,6 +640,14 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: raise RuntimeError( "on_user_function_end should only be called for CONTEXT and STEP operations" ) + # Detach first, on the same thread that attached, so the context this + # operation was entered from is restored exactly. Detaching (rather than + # attaching the enclosing span again) is what keeps the scopes balanced: + # code after a nested operation lands back on its parent context span, + # still attached here, and a top-level operation lands back on the + # thread's ambient context. Between-operation log records resolve through + # get_current_span_context(), which falls back to the invocation span. + context_scope.exit_scope(self, self._scope_key(info)) # key = f"{info.operation_id}-{int(info.start_time.timestamp())}" span_key = ( self._attempt_span_key(info) @@ -610,16 +680,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_context_scope.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py new file mode 100644 index 00000000..bf787a0b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_scope.py @@ -0,0 +1,630 @@ +"""Tests for balanced OTel context scopes across the plugin lifecycles. + +Covers :mod:`context_scope` directly plus the lifecycle paths the plugins rely +on it for: a suspended operation whose end hook never runs, hooks that execute +on a worker thread, and both plugins attaching on the same thread. +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime + +import opentelemetry.context as otel_context +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, + UserFunctionStartInfo, +) +from opentelemetry import baggage, trace +from opentelemetry.context import Context +from opentelemetry.sdk.trace import TracerProvider +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 import context_scope +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ProviderSource, +) + + +START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) +END_TIME = datetime(2024, 1, 2, 3, 4, 6, tzinfo=UTC) +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" + + +@pytest.fixture(autouse=True) +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context scope attached.""" + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth + assert otel_context.get_current() is before + + +class _Owner: + """Stand-in for a plugin instance; context_scope only uses its identity.""" + + +def _span_context(name: str) -> Context: + """Return a context carrying a non-recording span named after ``name``.""" + return trace.set_span_in_context( + trace.NonRecordingSpan( + trace.SpanContext( + trace_id=abs(hash(name)) % (1 << 128) or 1, + span_id=abs(hash(name)) % (1 << 64) or 1, + is_remote=False, + trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED), + ) + ), + otel_context.get_current(), + ) + + +# --------------------------------------------------------------------------- +# context_scope helper +# --------------------------------------------------------------------------- +def test_enter_then_exit_restores_the_previous_context(): + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope(owner, "a", lambda: _span_context("a")) + assert otel_context.get_current() is not before + assert context_scope.depth(owner) == 1 + + context_scope.exit_scope(owner, "a") + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_nested_scopes_pop_in_lifo_order(): + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope(owner, "outer", lambda: _span_context("outer")) + outer = otel_context.get_current() + context_scope.enter_scope(owner, "inner", lambda: _span_context("inner")) + + context_scope.exit_scope(owner, "inner") + assert otel_context.get_current() is outer + context_scope.exit_scope(owner, "outer") + assert otel_context.get_current() is before + + +def test_exit_unwinds_scopes_stacked_above_the_target(): + """Popping an outer scope must also drop anything left above it. + + ``ContextVar.reset`` writes back the token's captured value unconditionally, + so resetting out of order would revive a stale context. Unwinding downwards + keeps the underlying variable strictly LIFO. + """ + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope(owner, "outer", lambda: _span_context("outer")) + context_scope.enter_scope(owner, "orphan", lambda: _span_context("orphan")) + + context_scope.exit_scope(owner, "outer") + + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_exit_with_unknown_key_is_a_noop(): + owner = _Owner() + before = otel_context.get_current() + + context_scope.exit_scope(owner, "never-pushed") + + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_reentering_the_same_key_replaces_the_previous_scope(): + """Re-entering a key inside one invocation must not stack a second scope. + + A suspended operation is re-entered when its branch is resubmitted, and its + first scope is still attached because the suspending path had no end hook to + pop it. The epoch is unchanged, so only the ancestry check catches this. + """ + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope(owner, "wfc-1", lambda: _span_context("poll-1"), epoch=1) + context_scope.enter_scope(owner, "wfc-1", lambda: _span_context("poll-2"), epoch=1) + + assert context_scope.depth(owner) == 1 + + context_scope.exit_scope(owner, "wfc-1") + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_reentry_guard_keeps_enclosing_scopes(): + """Re-entering a nested key must not disturb the scope it is nested in.""" + owner = _Owner() + context_scope.enter_scope(owner, "ctx", lambda: _span_context("ctx"), epoch=1) + enclosing = otel_context.get_current() + + context_scope.enter_scope(owner, "inner", lambda: _span_context("inner-1"), epoch=1) + context_scope.enter_scope(owner, "inner", lambda: _span_context("inner-2"), epoch=1) + + assert context_scope.depth(owner) == 2 + context_scope.exit_scope(owner, "inner") + assert otel_context.get_current() is enclosing + + context_scope.exit_scope(owner, "ctx") + + +def test_enter_discards_scopes_from_a_previous_epoch(): + """A scope a suspended operation left behind must not outlive its invocation. + + The SDK re-raises ``SuspendExecution`` without calling + ``on_user_function_end``, so the next invocation on a reused thread finds a + stale scope; the epoch check drops it before attaching. + """ + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope( + owner, "suspended", lambda: _span_context("suspended"), epoch=1 + ) + assert context_scope.depth(owner) == 1 + + context_scope.enter_scope(owner, "next", lambda: _span_context("next"), epoch=2) + assert context_scope.depth(owner) == 1 + + context_scope.exit_scope(owner, "next") + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_unwind_detaches_every_scope_for_one_owner(): + owner = _Owner() + before = otel_context.get_current() + + context_scope.enter_scope(owner, "a", lambda: _span_context("a")) + context_scope.enter_scope(owner, "b", lambda: _span_context("b")) + + context_scope.unwind(owner) + + assert otel_context.get_current() is before + assert context_scope.depth(owner) == 0 + + +def test_unwind_is_a_noop_for_an_owner_with_no_scopes(): + before = otel_context.get_current() + + context_scope.unwind(_Owner()) + + assert otel_context.get_current() is before + + +def test_scopes_are_confined_to_the_thread_that_attached_them(): + """A worker thread's scopes must not appear on, or be poppable from, another. + + Tokens are only resettable in the ``contextvars.Context`` that created them, + so the stack is per thread. + """ + owner = _Owner() + before = otel_context.get_current() + observed: dict[str, object] = {} + + def worker() -> None: + context_scope.enter_scope(owner, "worker", lambda: _span_context("worker")) + observed["worker_depth"] = context_scope.depth(owner) + observed["worker_span_valid"] = ( + trace.get_current_span().get_span_context().is_valid + ) + context_scope.exit_scope(owner, "worker") + observed["worker_depth_after"] = context_scope.depth(owner) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(worker).result() + + assert observed["worker_depth"] == 1 + assert observed["worker_span_valid"] is True + assert observed["worker_depth_after"] == 0 + # The calling thread never saw the worker's scope. + assert context_scope.depth(owner) == 0 + assert otel_context.get_current() is before + + +# --------------------------------------------------------------------------- +# Plugin lifecycle paths +# --------------------------------------------------------------------------- +def _execution_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return ( + ExecutionOtelPlugin( + OtelPluginConfig( + provider_source=ProviderSource.EXPLICIT, + tracer_provider=provider, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ), + exporter, + ) + + +def _invocation_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return ( + InvocationOtelPlugin( + OtelPluginConfig( + provider_source=ProviderSource.EXPLICIT, + tracer_provider=provider, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ), + exporter, + ) + + +def _invocation_start() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + +def _invocation_end( + status: InvocationStatus = InvocationStatus.SUCCEEDED, +) -> InvocationEndInfo: + return InvocationEndInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + status=status, + error=None, + ) + + +def _step_start(operation_id: str) -> UserFunctionStartInfo: + return UserFunctionStartInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + + +def _step_end(operation_id: str) -> UserFunctionEndInfo: + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + 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, + ) + + +def _context_start(operation_id: str, parent_id: str | None) -> UserFunctionStartInfo: + """Start info for a child-context (including virtual/FLAT branch) body.""" + return UserFunctionStartInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=None, + ) + + +def _context_end(operation_id: str, parent_id: str | None) -> UserFunctionEndInfo: + """End info for a child-context body.""" + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + end_time=END_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + is_replay_children=False, + attempt=None, + outcome=UserFunctionOutcome.SUCCEEDED, + error=None, + ) + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_invocation_end_unwinds_a_suspended_operation_scope(factory): + """A step that suspends never gets its end hook; invocation end cleans up. + + ``wrap_user_function`` re-raises ``SuspendExecution`` without calling + ``on_user_function_end``, so the scope is still attached when the invocation + winds down. + """ + plugin, _ = factory() + before = otel_context.get_current() + + plugin.on_invocation_start(_invocation_start()) + plugin.on_user_function_start(_step_start("step-suspends")) + assert context_scope.depth(plugin) == 1 + + plugin.on_invocation_end(_invocation_end(InvocationStatus.PENDING)) + + assert context_scope.depth(plugin) == 0 + assert otel_context.get_current() is before + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_user_function_hooks_on_a_worker_thread_leave_the_caller_alone(factory): + """Verify hooks running on the user-code thread do not touch the caller. + + User code runs on a worker the SDK owns, and ``ThreadPoolExecutor`` does not + copy contextvars, so the plugin's scope must stay on that thread. + """ + plugin, _ = factory() + plugin.on_invocation_start(_invocation_start()) + before = otel_context.get_current() + observed: dict[str, object] = {} + + def run_step() -> None: + plugin.on_user_function_start(_step_start("step-1")) + observed["inside"] = trace.get_current_span().get_span_context().is_valid + plugin.on_user_function_end(_step_end("step-1")) + observed["after"] = context_scope.depth(plugin) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_step).result() + + assert observed["inside"] is True + assert observed["after"] == 0 + assert otel_context.get_current() is before + assert trace.get_current_span().get_span_context().is_valid is False + + plugin.on_invocation_end(_invocation_end()) + + +def test_the_new_context_is_built_after_stale_scopes_are_dropped(): + """The attached context must not inherit a scope that is being detached. + + The context is normally derived from what is current, so building it before + cleanup would copy baggage or suppression values out of the stale scope, and + detaching afterwards cannot remove them from an already-built Context. + """ + owner = _Owner() + before = otel_context.get_current() + + # A stale scope under the same key, carrying baggage. + context_scope.enter_scope( + owner, + "op-1", + lambda: baggage.set_baggage( + "poll", "first", context=otel_context.get_current() + ), + epoch=1, + ) + assert baggage.get_baggage("poll") == "first" + + # Re-entering builds its context from whatever is current at that moment, + # which must already be the pre-stale context. + observed: dict[str, object] = {} + + def build() -> Context: + observed["seen_during_build"] = baggage.get_baggage("poll") + return otel_context.get_current() + + context_scope.enter_scope(owner, "op-1", build, epoch=1) + + assert observed["seen_during_build"] is None + assert baggage.get_baggage("poll") is None + + context_scope.exit_scope(owner, "op-1") + assert otel_context.get_current() is before + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_flat_branch_scope_survives_its_inner_operations(factory): + """A FLAT map/parallel branch scope must stay attached across inner steps. + + A virtual (FLAT) branch 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). Physical nesting is + therefore not derivable from parent_id, and treating an inner step as + root-level must not detach the live branch scope: work between two inner steps + would fall out of the durable trace. + """ + plugin, _ = factory() + plugin.on_invocation_start(_invocation_start()) + observed: dict[str, object] = {} + + def run_branch() -> None: + # The virtual branch body: a CONTEXT operation at the top level. + plugin.on_user_function_start( + _context_start("flat-branch", parent_id=None), + ) + branch_span_id = trace.get_current_span().get_span_context().span_id + + # Inner steps of a FLAT branch report parent_id=None, not the branch. + for index in range(2): + step_id = f"flat-branch-step-{index}" + plugin.on_user_function_start(_step_start(step_id)) + plugin.on_user_function_end(_step_end(step_id)) + # Between inner steps the branch scope is still current. + observed[f"between-{index}"] = ( + trace.get_current_span().get_span_context().span_id + ) + + observed["branch"] = branch_span_id + plugin.on_user_function_end( + _context_end("flat-branch", parent_id=None), + ) + observed["after_branch_depth"] = context_scope.depth(plugin) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_branch).result() + + assert observed["between-0"] == observed["branch"] + assert observed["between-1"] == observed["branch"] + assert observed["after_branch_depth"] == 0 + + plugin.on_invocation_end(_invocation_end()) + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_an_empty_extracted_context_isolates_the_operation(factory): + """An extractor returning an empty Context must not inherit ambient values. + + Context subclasses dict, so an empty one is falsy; treating it as "no context" + would silently invert the extractor's intent and leak the worker's ambient + baggage into the operation. + """ + plugin, _ = factory() + plugin._context_extractor = lambda _info: Context() + plugin.on_invocation_start(_invocation_start()) + observed: dict[str, object] = {} + + def run_step() -> None: + # Something ambient on the worker, as auto-instrumentation might leave. + token = otel_context.attach( + baggage.set_baggage("ambient", "leaked", context=otel_context.get_current()) + ) + try: + plugin.on_user_function_start(_step_start("step-1")) + observed["inside"] = baggage.get_baggage("ambient") + plugin.on_user_function_end(_step_end("step-1")) + finally: + otel_context.detach(token) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_step).result() + + assert observed["inside"] is None + + plugin.on_invocation_end(_invocation_end()) + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_extracted_context_values_reach_user_code_on_a_worker_thread(factory): + """Values from the context extractor must be current inside user code. + + The worker running user code starts with an empty context, so the outermost + durable scope has to be layered onto the extracted context -- otherwise + baggage and any other non-span values the extractor supplied are dropped and + downstream instrumentation inside the step cannot propagate them. + """ + plugin, _ = factory() + # An extractor that supplies baggage, as a propagator-based one would. + plugin._context_extractor = lambda _info: baggage.set_baggage( + "tenant", "acme", context=Context() + ) + plugin.on_invocation_start(_invocation_start()) + observed: dict[str, object] = {} + + def run_step() -> None: + plugin.on_user_function_start(_step_start("step-1")) + observed["inside"] = baggage.get_baggage("tenant") + # A nested scope keeps it too, since it layers onto the current context. + plugin.on_user_function_start(_step_start("step-2")) + observed["nested"] = baggage.get_baggage("tenant") + plugin.on_user_function_end(_step_end("step-2")) + plugin.on_user_function_end(_step_end("step-1")) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_step).result() + + assert observed["inside"] == "acme" + assert observed["nested"] == "acme" + + plugin.on_invocation_end(_invocation_end()) + + +@pytest.mark.parametrize("factory", [_execution_plugin, _invocation_plugin]) +def test_suspend_then_reenter_then_end_leaves_no_residue(factory): + """A suspended operation re-entered in the same invocation stays balanced. + + The first start has no matching end (the SDK re-raises SuspendExecution), so + the re-entry must replace that scope rather than stack on it. + """ + plugin, _ = factory() + before = otel_context.get_current() + plugin.on_invocation_start(_invocation_start()) + observed: dict[str, object] = {} + + def run_polls() -> None: + # Poll 1 suspends: start fires, end never does. + plugin.on_user_function_start(_step_start("wfc-1")) + # Poll 2 re-enters the same operation and completes. + plugin.on_user_function_start(_step_start("wfc-1")) + observed["depth_after_reentry"] = context_scope.depth(plugin) + plugin.on_user_function_end(_step_end("wfc-1")) + observed["depth_after_end"] = context_scope.depth(plugin) + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_polls).result() + + assert observed["depth_after_reentry"] == 1 + assert observed["depth_after_end"] == 0 + assert otel_context.get_current() is before + + plugin.on_invocation_end(_invocation_end()) + + +def test_two_plugins_on_one_thread_unwind_in_lifo_order(): + """Both plugins ship as entry points and can be enabled together. + + Hooks dispatch in registration order, so the second plugin's scope is + detached while the first plugin's end hook runs. The shared stack keeps the + underlying ContextVar LIFO instead of reviving the first plugin's scope. + """ + first, _ = _execution_plugin() + second, _ = _invocation_plugin() + before = otel_context.get_current() + + for plugin in (first, second): + plugin.on_invocation_start(_invocation_start()) + for plugin in (first, second): + plugin.on_user_function_start(_step_start("step-1")) + for plugin in (first, second): + plugin.on_user_function_end(_step_end("step-1")) + + assert context_scope.depth() == 0 + assert otel_context.get_current() is before + + for plugin in (first, second): + plugin.on_invocation_end(_invocation_end()) 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..39c42c44 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 @@ -28,6 +28,7 @@ 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 import context_scope from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, operation_id_to_span_id, @@ -45,17 +46,22 @@ @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 scope 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. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth, ( + "test left OTel context scopes attached: " + f"{context_scope.depth() - before_depth} extra" + ) + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: @@ -446,6 +452,77 @@ def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch): assert invocation.parent.span_id == ambient.get_span_context().span_id +# ---------------------------------------------------------------------- +# Warm invocation reuse: no context leaks from one execution into the next +# ---------------------------------------------------------------------- +def test_invocation_end_restores_the_pre_invocation_context(): + """Verify the plugin leaves the handler thread's context untouched. + + The invocation lifecycle attaches nothing on this thread, so the Workflow + span must not be current afterwards -- an ended span left current is what + previously bled into the next invocation on a warm container. + """ + plugin, _ = _create_plugin() + before = otel_context.get_current() + + plugin.on_invocation_start(_invocation_start_info()) + 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_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 + the previous invocation left its Workflow span attached, that span became the + parent and its trace ID won, merging two unrelated executions into one trace. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + # 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 = ExecutionOtelPlugin( + OtelPluginConfig( + provider_source=ProviderSource.GLOBAL, + enrich_logger=False, + ) + ) + + def run(arn: str) -> None: + info = InvocationStartInfo( + request_id="request-1", + execution_arn=arn, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + plugin.on_invocation_start(info) + 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") + + invocations = [s for s in exporter.get_finished_spans() if s.name == "Invocation"] + assert len(invocations) == 2 + first, second = invocations + assert first.context.trace_id != second.context.trace_id + # The second invocation is a root, not a child of the first execution's span. + assert second.parent is None + + 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..d760fed7 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 @@ -37,6 +37,7 @@ 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 import context_scope from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, operation_id_to_span_id, @@ -60,12 +61,17 @@ @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 scope attached.""" + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth, ( + "test left OTel context scopes attached" + ) + 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..64ebcc36 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 @@ -31,6 +31,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind, StatusCode +from aws_durable_execution_sdk_python_otel import context_scope from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, operation_id_to_span_id, @@ -48,17 +49,23 @@ @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 scope 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. A test that drives hooks + the SDK itself leaves unpaired (a suspension) is responsible for unwinding. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth, ( + "test left OTel context scopes attached: " + f"{context_scope.depth() - before_depth} extra" + ) + assert otel_context.get_current() is before, ( + "test did not restore the OTel context it started with" + ) def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -791,13 +798,14 @@ 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() + before_depth = context_scope.depth(plugin) operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) @@ -808,18 +816,23 @@ def test_user_function_end_restores_invocation_span(): trace.get_current_span().get_span_context().span_id == active_attempt_span.get_span_context().span_id ) + assert context_scope.depth(plugin) == before_depth + 1 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 + # The scope is detached, restoring the context byte for byte. The invocation + # span is never attached (matching the Java plugins), so between-step log + # correlation goes through get_current_span_context() instead. + assert otel_context.get_current() is before + assert context_scope.depth(plugin) == before_depth + 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 +840,27 @@ 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 + assert context_scope.depth(plugin) == 0 + plugin.on_invocation_end(_invocation_end_info()) -def test_user_function_end_restores_invocation_span_across_multiple_steps(): - """Verify between-step context is the invocation span across many steps.""" +def test_sequential_steps_do_not_accumulate_scopes(): + """Verify N sequential steps leave no residue: depth returns to 0 each time.""" 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)) + assert context_scope.depth(plugin) == 1 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 + # Between each step the context is exactly what preceded the step, and + # no scope has accumulated. + assert otel_context.get_current() is before + assert context_scope.depth(plugin) == 0 + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -878,6 +897,11 @@ 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 + # Close the lifecycle so the attempt scope is detached; the autouse fixture + # asserts no scope outlives the test. + 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.""" @@ -929,25 +953,37 @@ def test_user_function_end_restores_parent_context_span_for_nested_step(): ) # After the inner step, the enclosing child-context span is current again, - # NOT the invocation span. + # NOT the invocation span. The inner scope was detached, landing back on the + # context scope that is still attached. assert trace.get_current_span().get_span_context().span_id == context_span_id assert ( trace.get_current_span().get_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_user_function_end_falls_back_to_invocation_when_parent_missing(): - """Verify a top-level step (parent_id=None) restores the invocation span.""" +def test_top_level_step_restores_ambient_context(): + """Verify a top-level step (parent_id=None) restores the ambient context. + + There is no enclosing operation scope to fall back to, so the thread returns + to whatever was current before 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 +1015,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 +1043,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..e5344439 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 @@ -41,6 +41,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import SpanKind +from aws_durable_execution_sdk_python_otel import context_scope from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( operation_id_to_span_id, ) @@ -63,12 +64,17 @@ @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 scope attached.""" + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth, ( + "test left OTel context scopes attached" + ) + 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..dc9af8a9 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 @@ -6,18 +6,27 @@ from datetime import UTC, datetime from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, OperationStatus, OperationType, ) from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, InvocationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, UserFunctionStartInfo, ) +import opentelemetry.context as otel_context +from opentelemetry import trace +import pytest from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider 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 import context_scope +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, @@ -33,6 +42,32 @@ EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" +@pytest.fixture(autouse=True) +def _assert_otel_context_balanced(): + """Fail any test that leaves an OTel context scope attached.""" + before = otel_context.get_current() + before_depth = context_scope.depth() + yield + assert context_scope.depth() == before_depth + assert otel_context.get_current() is before + + +def _create_execution_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: + """Create an ExecutionOtelPlugin wired to an in-memory span exporter.""" + exporter = InMemorySpanExporter() + trace_provider = TracerProvider() + trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + provider_source=ProviderSource.EXPLICIT, + tracer_provider=trace_provider, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + return plugin, exporter + + def _create_plugin( enrich_logger: bool = True, ) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -77,6 +112,37 @@ def _user_function_start_info(operation_id: str) -> UserFunctionStartInfo: ) +def _invocation_end_info() -> InvocationEndInfo: + """Create standard invocation end info for tests.""" + return InvocationEndInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + status=InvocationStatus.SUCCEEDED, + error=None, + ) + + +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 +214,71 @@ 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)) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_execution_plugin_handler_thread_uses_the_invocation_span(): + """Handler-thread records correlate to the Invocation span, not Workflow. + + ExecutionOtelPlugin attaches nothing on the handler thread, so the filter + resolves through the plugin registry, which prefers the Invocation span. The + trace ID is shared with the Workflow span either way. + """ + plugin, _ = _create_execution_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + record = _make_record() + OtelContextLogFilter(plugin).filter(record) + + invocation_context = plugin._invocation_span.get_span_context() + workflow_context = plugin._workflow_span.get_span_context() + assert record.spanId == format(invocation_context.span_id, "016x") + assert record.spanId != format(workflow_context.span_id, "016x") + assert record.traceId == format(workflow_context.trace_id, "032x") + + plugin.on_invocation_end(_invocation_end_info()) + + +def test_ambient_lambda_span_does_not_displace_the_invocation_span(monkeypatch): + """An ambient ADOT span must not be reported in place of the durable span. + + In GLOBAL mode the Lambda invocation span from the ADOT layer stays current on + the handler thread. Log records emitted there must still carry the durable + Invocation span, so the filter only trusts the current span while the plugin + holds an operation scope on that thread. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + provider_source=ProviderSource.GLOBAL, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + + ambient = provider.get_tracer("ambient").start_span("lambda-invocation") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + plugin.on_invocation_start(_invocation_start_info()) + record = _make_record() + OtelContextLogFilter(plugin).filter(record) + + ambient_span_id = format(ambient.get_span_context().span_id, "016x") + invocation_span_id = format( + plugin._invocation_span.get_span_context().span_id, "016x" + ) + assert record.spanId == invocation_span_id + assert record.spanId != ambient_span_id + + plugin.on_invocation_end(_invocation_end_info()) + finally: + otel_context.detach(token) + ambient.end() + def test_install_log_filter_attaches_to_handlers(): """install_log_filter adds the filter to each handler on the target logger."""