From 7bf83c465c7e3ae3894205e9131a88c064afa055 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Mon, 17 Aug 2026 23:06:57 +0000 Subject: [PATCH] fix(otel): balance otel context attach and detach - Track every context.attach() token in the plugin and release it at the matching lifecycle end: user-function scopes in on_user_function_end, the invocation scope during invocation cleanup - Replace the re-attach of the enclosing span in on_user_function_end with a detach, so nested operations manage a balanced scope instead of stacking a new one - Release scopes still open at invocation end (a suspended user function or an aborted invocation) so nothing survives into a warm invocation - Skip tokens recorded on another thread, which OpenTelemetry cannot reset - Replace the autouse OTel context reset in the plugin tests with an assertion that each test leaves the context as it found it, and cover nested child contexts, sequential steps, failures, suspension, and warm invocation reuse Resolves #643 --- .../execution_plugin.py | 86 +++++- .../invocation_plugin.py | 76 ++++- .../tests/test_execution_plugin.py | 280 +++++++++++++++++- .../test_execution_plugin_integration.py | 13 +- .../tests/test_invocation_plugin.py | 217 ++++++++++++-- .../test_invocation_plugin_integration.py | 13 +- 6 files changed, 610 insertions(+), 75 deletions(-) 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 564ad905..9d328a44 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 @@ -17,6 +17,11 @@ - context is attached inside the synchronous ``on_user_function_*`` hooks and log correlation is handled by :mod:`log_filter`), the hook wiring mirrors the existing :class:`~aws_durable_execution_sdk_python_otel.invocation_plugin.InvocationOtelPlugin`. + +Every context the plugin attaches is tracked by its token and detached at the +matching lifecycle end: a user-function scope is released in +``on_user_function_end`` and the invocation scope in ``on_invocation_end``, so +the plugin never leaves an ended or suspended span current. """ from __future__ import annotations @@ -75,6 +80,9 @@ # Registry key for the invocation span (operations use their operation_id). _INVOCATION_KEY = "__invocation__" +# Token key for the invocation-level context scope attached at invocation start. +_INVOCATION_CONTEXT_KEY = "__invocation_context__" + def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: """Convert a datetime to an OTel timestamp (ns since epoch), or None.""" @@ -114,6 +122,11 @@ 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] = {} + # Tokens returned by context.attach(), keyed by the span registry key, + # paired with the thread that attached them. Every attach the plugin + # owns is released through _detach_context so the plugin never leaves a + # scope on the context stack. + self._context_tokens: dict[str, tuple[int, object]] = {} self._lock = threading.RLock() self._tracing_enabled = False @@ -157,6 +170,46 @@ 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}" + # ------------------------------------------------------------------ + # Context scope helpers + # ------------------------------------------------------------------ + def _attach_context(self, key: str, new_context: Context) -> None: + """Attach a context and remember its token under ``key``.""" + with self._lock: + self._context_tokens[key] = ( + threading.get_ident(), + otel_context.attach(new_context), + ) + + def _detach_context(self, key: str) -> None: + """Detach the context attached under ``key``, restoring its predecessor. + + A context token can only be reset on the thread that created it, so a + token recorded on another thread is dropped instead of detached (OTel + logs an error for a cross-thread reset). In practice the pairs always + line up: invocation hooks run on the Lambda handler thread and + user-function hooks run on the thread executing user code. + """ + with self._lock: + entry = self._context_tokens.pop(key, None) + if entry is None: + return + thread_ident, token = entry + if thread_ident == threading.get_ident(): + otel_context.detach(token) # type: ignore[arg-type] + + def _detach_remaining_contexts(self) -> None: + """Release scopes still open, newest first, so nothing outlives the plugin. + + Reached when a lifecycle end hook never fires -- for example a user + function that suspends, or a warm invocation that starts before the + previous one was cleaned up. + """ + with self._lock: + keys = list(reversed(self._context_tokens)) + for key in keys: + self._detach_context(key) + 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() @@ -225,10 +278,13 @@ 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. The token is + # released in _reset_state at invocation end, restoring the context that + # was active before the invocation started. if self._workflow_span is not None: - otel_context.attach( - trace.set_span_in_context(self._workflow_span, self._extracted_context) + self._attach_context( + _INVOCATION_CONTEXT_KEY, + trace.set_span_in_context(self._workflow_span, self._extracted_context), ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: @@ -324,6 +380,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: logger.exception("force_flush failed at invocation end") def _reset_state(self) -> None: + self._detach_remaining_contexts() self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None @@ -439,17 +496,19 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: info.parent_id ) name = f"{info.name or info.operation_id} attempt {info.attempt or 1}" + key = self._attempt_key(info) span = self._start_span( operation_id=info.operation_id, name=name, info=info, parent=parent, start_time=info.start_time, - span_key=self._attempt_key(info), + span_key=key, deterministic=False, ) else: # CONTEXT parent = self._resolve_parent(info.parent_id) + key = info.operation_id span = self._start_span( operation_id=info.operation_id, name=info.name or info.operation_id, @@ -457,7 +516,9 @@ 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._attach_context( + key, 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) @@ -500,16 +561,11 @@ 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) - ) + # Restore the enclosing context by releasing the scope this user + # function attached, so the parent operation (or, at the top level, the + # context that was active before the operation) becomes current again + # without stacking another scope. + self._detach_context(key) # ------------------------------------------------------------------ # 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 f0fe0b89..f702fc7a 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 @@ -123,6 +123,11 @@ 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] = {} + # Tokens returned by context.attach(), keyed by the span registry key, + # paired with the thread that attached them. Every attach the plugin + # owns is released through _detach_context so the plugin never leaves a + # scope on the context stack. + self._context_tokens: dict[str, tuple[int, object]] = {} self._operation_spans_lock = threading.RLock() self._tracing_enabled = False @@ -167,6 +172,46 @@ 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}" + # ------------------------------------------------------------------ + # Context scope helpers + # ------------------------------------------------------------------ + def _attach_context(self, key: str, new_context: Context) -> None: + """Attach a context and remember its token under ``key``.""" + with self._operation_spans_lock: + self._context_tokens[key] = ( + threading.get_ident(), + context.attach(new_context), + ) + + def _detach_context(self, key: str) -> None: + """Detach the context attached under ``key``, restoring its predecessor. + + A context token can only be reset on the thread that created it, so a + token recorded on another thread is dropped instead of detached (OTel + logs an error for a cross-thread reset). In practice the pairs always + line up: user-function hooks run on the thread executing user code, and + both the start and end hook for one attempt run on that same thread. + """ + with self._operation_spans_lock: + entry = self._context_tokens.pop(key, None) + if entry is None: + return + thread_ident, token = entry + if thread_ident == threading.get_ident(): + context.detach(token) # type: ignore[arg-type] + + def _detach_remaining_contexts(self) -> None: + """Release scopes still open, newest first, so nothing outlives the plugin. + + Reached when a lifecycle end hook never fires -- for example a user + function that suspends, or a warm invocation that starts before the + previous one was cleaned up. + """ + with self._operation_spans_lock: + keys = list(reversed(self._context_tokens)) + for key in keys: + self._detach_context(key) + def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. @@ -174,12 +219,15 @@ def get_current_span_context(self) -> SpanContext | None: 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). + on_user_function_start), and between the steps of a child context it + is the enclosing context span, restored when on_user_function_end + detaches the inner scope. 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. + resolve it. It also covers code between top-level operations, where + detaching the operation scope restores a context with no durable + span. Returns: A valid SpanContext, or None if no span is active. @@ -445,6 +493,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _reset_state(self) -> None: """Clear per-invocation state for warm Lambda environment reuse.""" + self._detach_remaining_contexts() self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None @@ -557,7 +606,9 @@ 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._attach_context( + span_key, 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. @@ -607,16 +658,13 @@ 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) - ) + # Restore the enclosing context by releasing the scope this user + # function attached. Code that runs after this operation (e.g. between + # steps in a child context) correlates to its enclosing operation + # again -- the parent context span for a nested operation, and the + # context active before the operation for a top-level one, where + # get_current_span_context falls back to the invocation span. + self._detach_context(span_key) 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 9ef7a765..f03cb337 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 from types import SimpleNamespace @@ -43,17 +44,18 @@ @pytest.fixture(autouse=True) -def _reset_otel_context(): - """Reset the OTel thread-local context around each test. +def _assert_otel_context_balanced(): + """Assert each test leaves the OTel thread-local context as it found it. - The plugin attaches spans via context.attach() without detaching, so state - would otherwise leak between tests running on the same thread. + The plugin pairs every context.attach() with a context.detach(), so no + global reset is needed to keep tests isolated. Asserting the invariant here + turns a plugin lifecycle leak into a test failure instead of hiding it. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + yield + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: @@ -380,6 +382,8 @@ 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 + plugin.on_invocation_end(_invocation_end_info()) + def test_step_attempt_span_omits_operation_status(): plugin, exporter = _create_plugin() @@ -429,6 +433,8 @@ def test_step_attempt_span_omits_operation_status(): ) assert "durable.operation.status" not in span.attributes + plugin.on_invocation_end(_invocation_end_info()) + # --------------------------------------------------------------------------- # Default-provider mode: invocation span @@ -541,3 +547,259 @@ def test_invocation_span_status_kind_and_attributes(status, expected_code): assert invocation.attributes["durable.invocation.status"] == status.value assert invocation.attributes["durable.invocation.first"] is True assert invocation.status.status_code is expected_code + + +# --------------------------------------------------------------------------- +# Context attach/detach balance across the plugin lifecycle +# --------------------------------------------------------------------------- +def _step_start_info( + operation_id: str, + parent_id: str | None = None, + attempt: int = 1, +) -> UserFunctionStartInfo: + return UserFunctionStartInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=attempt, + ) + + +def _step_end_info( + operation_id: str, + parent_id: str | None = None, + attempt: int = 1, + outcome: UserFunctionOutcome = UserFunctionOutcome.SUCCEEDED, +) -> UserFunctionEndInfo: + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=attempt, + outcome=outcome, + end_time=END_TIME, + error=None, + ) + + +def _context_start_info(operation_id: str) -> UserFunctionStartInfo: + return UserFunctionStartInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=operation_id, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + + +def _context_end_info(operation_id: str) -> UserFunctionEndInfo: + return UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, + name=operation_id, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + outcome=UserFunctionOutcome.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + + +def test_workflow_span_is_current_during_the_invocation(): + """Verify the Workflow span is the active span while the invocation runs.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + 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()) + + +@pytest.mark.parametrize( + "status", + [InvocationStatus.PENDING, InvocationStatus.SUCCEEDED, InvocationStatus.FAILED], +) +def test_invocation_end_restores_context_from_before_invocation_start(status): + """Verify invocation cleanup leaves no plugin span current. + + A non-terminal invocation used to leave the Workflow span attached, so work + after cleanup was parented to a span that had not ended. + """ + plugin, _ = _create_plugin() + before_context = otel_context.get_current() + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(status=status)) + + assert otel_context.get_current() == before_context + assert not trace.get_current_span().get_span_context().is_valid + assert plugin._context_tokens == {} + + +@pytest.mark.parametrize( + "outcome", [UserFunctionOutcome.SUCCEEDED, UserFunctionOutcome.FAILED] +) +def test_step_scope_is_released_at_user_function_end(outcome): + """Verify a finished step restores the context that enclosed it.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + enclosing_context = otel_context.get_current() + + plugin.on_user_function_start(_step_start_info("step-1")) + attempt_span = plugin._get_span("step-1:attempt:1") + assert attempt_span is not None + assert ( + trace.get_current_span().get_span_context().span_id + == attempt_span.get_span_context().span_id + ) + + plugin.on_user_function_end(_step_end_info("step-1", outcome=outcome)) + + assert otel_context.get_current() == enclosing_context + # Only the invocation-level scope remains open. + assert set(plugin._context_tokens) == {"__invocation_context__"} + + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._context_tokens == {} + + +def test_sequential_steps_do_not_accumulate_scopes(): + """Verify repeated steps unwind to the same enclosing context each time.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + enclosing_context = otel_context.get_current() + + for index in range(3): + operation_id = f"step-{index}" + plugin.on_user_function_start(_step_start_info(operation_id)) + plugin.on_user_function_end(_step_end_info(operation_id)) + assert otel_context.get_current() == enclosing_context + assert set(plugin._context_tokens) == {"__invocation_context__"} + + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._context_tokens == {} + + +def test_nested_scopes_are_released_without_accumulating(): + """Verify a child context and its inner step unwind to their entry contexts.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + context_id = "ctx-1" + plugin.on_user_function_start(_context_start_info(context_id)) + inside_context = otel_context.get_current() + context_span = plugin._get_span(context_id) + assert context_span is not None + + plugin.on_user_function_start(_step_start_info("ctx-1-step", parent_id=context_id)) + plugin.on_user_function_end(_step_end_info("ctx-1-step", parent_id=context_id)) + + # The inner step restored the child-context scope, not a copy of it. + assert otel_context.get_current() == inside_context + assert ( + trace.get_current_span().get_span_context().span_id + == context_span.get_span_context().span_id + ) + assert set(plugin._context_tokens) == {"__invocation_context__", context_id} + + plugin.on_user_function_end(_context_end_info(context_id)) + + assert otel_context.get_current() == before_context + assert set(plugin._context_tokens) == {"__invocation_context__"} + + plugin.on_invocation_end(_invocation_end_info()) + assert plugin._context_tokens == {} + + +def test_invocation_end_releases_scope_of_suspended_user_function(): + """Verify a user function that never ends does not leak its scope. + + A suspending user function raises before ``on_user_function_end`` runs, so + invocation cleanup is what releases the scope it attached. + """ + plugin, _ = _create_plugin() + before_context = otel_context.get_current() + plugin.on_invocation_start(_invocation_start_info()) + + plugin.on_user_function_start(_step_start_info("step-suspends")) + assert plugin._context_tokens + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + + +def test_warm_invocation_reuse_restores_ambient_span_each_time(): + """Verify repeated invocations leave the ambient Lambda span current.""" + plugin, _ = _create_plugin() + ambient_provider = TracerProvider() + ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + warm_context = otel_context.get_current() + for index in range(3): + plugin.on_invocation_start(_invocation_start_info()) + operation_id = f"step-{index}" + plugin.on_user_function_start(_step_start_info(operation_id)) + plugin.on_user_function_end(_step_end_info(operation_id)) + plugin.on_invocation_end(_invocation_end_info()) + + assert otel_context.get_current() == warm_context + assert ( + trace.get_current_span().get_span_context().span_id + == ambient.get_span_context().span_id + ) + assert plugin._context_tokens == {} + finally: + otel_context.detach(token) + ambient.end() + + +def test_detach_ignores_token_attached_on_another_thread(): + """Verify a scope attached on another thread is dropped, not reset here. + + A context token can only be reset on the thread that created it, so the + plugin drops foreign tokens instead of asking OpenTelemetry to fail. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit( + plugin.on_user_function_start, _step_start_info("step-1") + ).result() + + plugin._detach_context("step-1:attempt:1") + + assert "step-1:attempt:1" not in plugin._context_tokens + assert otel_context.get_current() == before_context + + plugin.on_invocation_end(_invocation_end_info()) 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 46dbed99..52000518 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 @@ -63,12 +63,13 @@ @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(): + """Assert each test leaves the OTel thread-local context as it found it.""" + before = otel_context.get_current() + yield + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) 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 3e0efd8a..34bd8bc8 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 @@ -46,17 +46,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(): + """Assert each test leaves the OTel thread-local context as it found it. - The plugin attaches spans via context.attach() without ever detaching, - so state would otherwise leak between tests running on the same thread. + The plugin pairs every context.attach() with a context.detach(), so no + global reset is needed to keep tests isolated. Asserting the invariant here + turns a plugin lifecycle leak into a test failure instead of hiding it. """ - token = otel_context.attach(Context()) - try: - yield - finally: - otel_context.detach(token) + before = otel_context.get_current() + yield + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: @@ -841,13 +842,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 enclosing the user function # ---------------------------------------------------------------------- -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 a completed step leaves the pre-step context current again.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) invocation_span_id = plugin._get_span(None).get_span_context().span_id + enclosing_context = otel_context.get_current() operation_id = "step-1" plugin.on_user_function_start(_user_function_start_info(operation_id)) @@ -861,15 +863,18 @@ 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 + # After the step, the enclosing context is restored and no scope is left + # behind. Log correlation resolves the invocation span from the registry. + assert otel_context.get_current() == enclosing_context + assert plugin._context_tokens == {} + assert plugin.get_current_span_context().span_id == invocation_span_id -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 enclosing 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 + enclosing_context = otel_context.get_current() operation_id = "step-fail" plugin.on_user_function_start(_user_function_start_info(operation_id)) @@ -877,21 +882,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() == enclosing_context + assert plugin._context_tokens == {} -def test_user_function_end_restores_invocation_span_across_multiple_steps(): - """Verify between-step context is the invocation span across many steps.""" +def test_user_function_end_restores_enclosing_context_across_multiple_steps(): + """Verify sequential steps do not accumulate context scopes.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) invocation_span_id = plugin._get_span(None).get_span_context().span_id + enclosing_context = 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 + # Between each step the context is back to where it started, and log + # correlation still resolves the invocation span. + assert otel_context.get_current() == enclosing_context + assert plugin._context_tokens == {} + assert plugin.get_current_span_context().span_id == invocation_span_id # ---------------------------------------------------------------------- @@ -928,6 +938,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 + # The step never ends here, so invocation cleanup releases its scope. + 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.""" @@ -986,18 +999,26 @@ def test_user_function_end_restores_parent_context_span_for_nested_step(): != plugin._get_span(None).get_span_context().span_id ) + # The child context never ends here, so invocation cleanup releases it. + 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_end_falls_back_to_invocation_for_correlation(): + """Verify a top-level step (parent_id=None) correlates to the invocation.""" plugin, _ = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) invocation_span_id = plugin._get_span(None).get_span_context().span_id + enclosing_context = 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 + # No durable span is attached at the top level, so the registry fallback + # supplies the invocation span for log correlation. + assert otel_context.get_current() == enclosing_context + assert not trace.get_current_span().get_span_context().is_valid + assert plugin.get_current_span_context().span_id == invocation_span_id def test_get_current_span_context_returns_context_span_between_nested_steps(): @@ -1029,6 +1050,9 @@ 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 + # The child context never ends here, so invocation cleanup releases it. + 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.""" @@ -1052,6 +1076,9 @@ 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 + # The child context never ends here, so invocation cleanup releases it. + plugin.on_invocation_end(_invocation_end_info()) + @pytest.mark.parametrize( ("status", "expected_code"), @@ -1152,3 +1179,143 @@ def test_workflow_span_name_is_configurable(): names = [s.name for s in exporter.get_finished_spans()] assert "MyExecution" in names assert "Workflow" not in names + + +# ---------------------------------------------------------------------- +# Context attach/detach balance across the plugin lifecycle +# ---------------------------------------------------------------------- +def test_child_context_end_restores_context_active_before_it(): + """Verify leaving a child context restores the context that enclosed it.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + enclosing_context = otel_context.get_current() + + context_id = "ctx-1" + plugin.on_user_function_start( + _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) + ) + assert otel_context.get_current() != enclosing_context + + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + + assert otel_context.get_current() == enclosing_context + assert plugin._context_tokens == {} + + +def test_nested_scopes_are_released_without_accumulating(): + """Verify a child context and its inner step unwind to their entry contexts.""" + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + context_id = "ctx-1" + plugin.on_user_function_start( + _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) + ) + inside_context = otel_context.get_current() + + inner_step_id = "ctx-1-step" + plugin.on_user_function_start( + _user_function_start_info(inner_step_id, parent_id=context_id) + ) + plugin.on_user_function_end( + _user_function_end_info(inner_step_id, parent_id=context_id) + ) + # The inner step restored the child-context scope, not a copy of it. + assert otel_context.get_current() == inside_context + assert set(plugin._context_tokens) == {context_id} + + plugin.on_user_function_end( + _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) + ) + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + + +def test_invocation_end_releases_scope_of_suspended_user_function(): + """Verify a user function that never ends does not leak its scope. + + A suspending user function raises before ``on_user_function_end`` runs, so + invocation cleanup is what releases the scope it attached. + """ + plugin, _ = _create_plugin() + before_context = otel_context.get_current() + plugin.on_invocation_start(_invocation_start_info()) + + plugin.on_user_function_start(_user_function_start_info("step-suspends")) + assert plugin._context_tokens + + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + assert otel_context.get_current() == before_context + assert plugin._context_tokens == {} + + +def test_ambient_span_is_current_again_after_full_lifecycle(): + """Verify an ambient (e.g. ADOT) span survives a full plugin lifecycle.""" + plugin, _ = _create_plugin() + ambient_provider = TracerProvider() + ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + plugin.on_invocation_start(_invocation_start_info()) + 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)) + plugin.on_invocation_end(_invocation_end_info()) + + assert ( + trace.get_current_span().get_span_context().span_id + == ambient.get_span_context().span_id + ) + finally: + otel_context.detach(token) + ambient.end() + + +def test_warm_invocation_reuse_does_not_accumulate_scopes(): + """Verify repeated invocations on one plugin instance stay balanced.""" + plugin, _ = _create_plugin() + ambient_provider = TracerProvider() + ambient = ambient_provider.get_tracer("ambient").start_span("AmbientLambda") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + warm_context = otel_context.get_current() + for index in range(3): + plugin.on_invocation_start(_invocation_start_info()) + 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)) + plugin.on_invocation_end(_invocation_end_info()) + + # Each invocation leaves the warm environment as it found it. + assert otel_context.get_current() == warm_context + assert plugin._context_tokens == {} + finally: + otel_context.detach(token) + ambient.end() + + +def test_detach_ignores_token_attached_on_another_thread(): + """Verify a scope attached on another thread is dropped, not reset here. + + A context token can only be reset on the thread that created it, so the + plugin drops foreign tokens instead of asking OpenTelemetry to fail. + """ + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + before_context = otel_context.get_current() + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit( + plugin.on_user_function_start, _user_function_start_info("step-1") + ).result() + + plugin._detach_context("step-1:attempt:1") + + assert plugin._context_tokens == {} + assert otel_context.get_current() == before_context + + plugin.on_invocation_end(_invocation_end_info()) 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 0e4bb0d5..b7f69152 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 @@ -67,12 +67,13 @@ @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(): + """Assert each test leaves the OTel thread-local context as it found it.""" + before = otel_context.get_current() + yield + assert otel_context.get_current() == before, ( + "test leaked OTel context state: an attach() was not detached" + ) def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: