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."""