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]: