Skip to content

fix(otel): end recording spans on non-terminal #696

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ayushiahjolia wants to merge 1 commit into from
Open

fix(otel): end recording spans on non-terminal #696

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
209 changes: 151 additions & 58 deletions ...e-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@
ancestor: a propagated backend parent when present, otherwise a deterministic
synthetic root. The Workflow span is exported exactly once, when the execution
reaches a terminal status. Operations are parented under the Workflow span (or
their parent operation) and *linked* to the current Invocation span.
their parent operation) and *linked* to the current Invocation span. Each
operation is likewise exported exactly once, on its deterministic span ID, when
it reaches a terminal status; while it spans invocations it is held as a
non-recording placeholder so no recording span is abandoned.

This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from
aws-durable-execution-sdk-js#729. Because the Python plugin interface differs
Expand Down Expand Up @@ -56,6 +59,7 @@
SpanKind,
StatusCode,
Tracer,
TraceState,
)

from aws_durable_execution_sdk_python_otel.context_extractors import (
Expand Down Expand Up @@ -133,12 +137,16 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
# Per-invocation state.
self._execution_arn = ""
self._execution_trace_id: int | None = None
self._execution_start_time: datetime.datetime | None = None
self._extracted_context: ExtractedContext | None = None
self._execution_trace_context: ExecutionTraceContext | None = None
self._sampling_intent: DurableSamplingIntent | None = None
self._workflow_span: Span | None = None
self._invocation_span: Span | None = None
self._operation_spans: dict[str, Span] = {}
# Operations whose span was already exported this invocation, so a
# repeated on_operation_end does not export it twice.
self._ended_operation_ids: set[str] = set()
# 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
Expand Down Expand Up @@ -292,6 +300,41 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None:
return existing
return self._workflow_span

def _resolved_trace_state(self) -> TraceState:
"""Return the resolved sampling trace state, else the ancestor state.

The sampling result preserves a same-trace ambient ``tracestate`` that
the empty ancestor state would drop.
"""
intent = self._sampling_intent
if intent is not None and intent.result.trace_state is not None:
return intent.result.trace_state
if self._execution_trace_context is not None:
return self._execution_trace_context.execution_ancestor.trace_state
return TraceState()

def _operation_span_context(self, operation_id: str) -> SpanContext | None:
"""Return the deterministic SpanContext for a logical operation."""
execution_trace_context = self._execution_trace_context
if execution_trace_context is None:
return None
return SpanContext(
trace_id=execution_trace_context.trace_id,
span_id=operation_id_to_span_id(self._execution_arn, operation_id),
is_remote=False,
trace_flags=execution_trace_context.trace_flags,
trace_state=self._resolved_trace_state(),
)

def _register_operation_placeholder(self, operation_id: str) -> Span | None:
"""Register a non-recording placeholder holding the operation context."""
span_context = self._operation_span_context(operation_id)
if span_context is None:
return None
placeholder = NonRecordingSpan(span_context)
self._set_span(operation_id, placeholder)
return placeholder

def _invocation_parent_context(self) -> Context:
"""Return same-trace ambient context, else execution ancestor context."""
execution_trace_context = self._execution_trace_context
Expand Down Expand Up @@ -346,6 +389,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
)
self._tracing_enabled = False
return
self._execution_start_time = info.execution_start_time
self._extracted_context = _ensure_extracted_context(
self._context_extractor(info)
)
Expand Down Expand Up @@ -393,11 +437,40 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
)

def _start_workflow_span(self, info: InvocationStartInfo) -> None:
"""Install a non-recording placeholder for the execution-scoped Workflow span.

The Workflow span spans the whole durable execution and is exported once,
on the terminal invocation. During every invocation the plugin only needs
its deterministic SpanContext -- to parent operation spans, to keep the
Workflow current so auto-instrumented spans join the execution trace, and
for log correlation. A non-recording placeholder fills that role so a
non-terminal invocation never abandons a recording span. The recording
span is created and ended once by :meth:`_export_workflow_span`.
"""
if not self._execution_arn:
logger.warning("No execution ARN; skipping Workflow span creation")
return
if self._execution_trace_context is None:
return
workflow_span_context = SpanContext(
trace_id=self._execution_trace_context.trace_id,
span_id=derive_workflow_span_id(self._execution_arn),
is_remote=False,
trace_flags=self._execution_trace_context.trace_flags,
trace_state=self._resolved_trace_state(),
)
self._workflow_span = NonRecordingSpan(workflow_span_context)

def _export_workflow_span(self, info: InvocationEndInfo) -> None:
"""Create and end the recording Workflow span once, on a terminal status.

Uses the same deterministic span ID as the placeholder and the shared
execution ancestor as its parent, so the exported Workflow span stays on
the execution trace and correlates with every operation span across all
invocations. Anchored at the execution start time.
"""
if not self._execution_arn or self._execution_trace_context is None:
return
parent_context = self._with_sampling(
trace.set_span_in_context(
NonRecordingSpan(self._execution_trace_context.execution_ancestor),
Expand All @@ -408,13 +481,44 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None:
trace_id=None,
span_id=derive_workflow_span_id(self._execution_arn),
):
self._workflow_span = self._tracer.start_span(
workflow_span = self._tracer.start_span(
name=self._workflow_span_name,
kind=SpanKind.INTERNAL,
attributes={"durable.execution.arn": self._execution_arn},
start_time=_to_otel_timestamp(info.execution_start_time),
attributes={
"durable.execution.arn": self._execution_arn,
"durable.execution.status": (
info.status.value if info.status else ""
),
},
start_time=_to_otel_timestamp(self._execution_start_time),
context=parent_context,
)
if info.status is InvocationStatus.FAILED:
workflow_span.set_status(
StatusCode.ERROR, info.error.message if info.error else ""
)
elif info.status is InvocationStatus.SUCCEEDED:
workflow_span.set_status(StatusCode.OK)
workflow_span.end()

def _end_open_recording_spans(self) -> None:
"""End recording user-function spans left open by a suspended operation.

Operation placeholders are non-recording and export their span from
on_operation_end, so they are skipped. Reverse order keeps each child
contained within its parent; the invocation span is ended by the caller.
"""
with self._lock:
keys = list(reversed(self._operation_spans))
Comment on lines +511 to +512

This comment was marked as outdated.

Sign in to view

Comment on lines +511 to +512

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex AI review · Finding arf_v1_dv5gcy7ci2c3j7gxu7vdwsejwy

[P2] Close attempts overwritten by in-process re-entry. This sweep can see only the latest span stored for each key. If an incomplete STEP re-enters the same attempt, _start_span() replaces the first recording span; ending the replacement then leaves the original unreachable and unended. End or retain the existing attempt before replacement, and extend the re-entry test to assert the original span is no longer recording.

for key in keys:
if key == _INVOCATION_KEY:
continue
span = self._get_span(key)
if span is None or not span.is_recording():
continue
popped = self._pop_span(key)
if popped is not None:
popped.end()
Comment thread
ayushiahjolia marked this conversation as resolved.
Comment thread
ayushiahjolia marked this conversation as resolved.
Comment thread
zhongkechen marked this conversation as resolved.
Comment thread
ayushiahjolia marked this conversation as resolved.
Comment on lines +504 to +521

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude AI review · Finding arf_v1_d3v4nkts2wbw4vfdae6gv4poj6

_end_open_recording_spans() force-closes any recording span still open at invocation end (the realistic case is a STEP attempt whose outcome was INCOMPLETE, e.g. an in-process retry that didn't finish before the invocation suspended) by calling popped.end() with no status or attribute set. The exported span therefore has StatusCode.UNSET and none of the durable.attempt.* attributes that a normal completed attempt gets from on_user_function_end, so in a trace viewer it is indistinguishable from a short, uneventful successful attempt — hiding the fact that it was actually truncated because the invocation ended mid-attempt.

Consider setting an explicit marker before ending, e.g. popped.set_attribute("durable.span.truncated_at_invocation_boundary", True) (and/or leaving status UNSET is fine, but the attribute gives downstream consumers a way to distinguish this from a genuinely completed, short attempt).


def _start_invocation_span(self, info: InvocationStartInfo) -> None:
self._invocation_span = self._tracer.start_span(
Expand All @@ -434,12 +538,6 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
self._reset_state()
return

# Operation spans still open here belong to operations that suspended
# (e.g. PENDING/RETRYING) rather than completed this invocation. They are
# ended only by on_operation_end; drop the references without ending them
# so they are not exported as if completed. _reset_state
# clears the span map below.

# End the invocation span regardless of terminal status. Record the
# invocation status and map it to a span status:
# SUCCEEDED/PENDING -> OK (this invocation did its work, whether it
Expand All @@ -462,23 +560,16 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
)
self._invocation_span.end()

# The Workflow span (execution view) is exported only on a terminal
# status; otherwise its reference is dropped without ending it. Its span
# status reflects the execution outcome: SUCCEEDED -> OK, FAILED -> ERROR
# (RETRY/PENDING are non-terminal and never reach here -> UNSET).
if self._workflow_span is not None:
if info.status in _TERMINAL_INVOCATION_STATUSES:
self._workflow_span.set_attribute(
"durable.execution.status",
info.status.value if info.status else "",
)
if info.status is InvocationStatus.FAILED:
self._workflow_span.set_status(
StatusCode.ERROR, info.error.message if info.error else ""
)
elif info.status is InvocationStatus.SUCCEEDED:
self._workflow_span.set_status(StatusCode.OK)
self._workflow_span.end()
# End recording user-function spans left open by a suspended operation.
self._end_open_recording_spans()
Comment thread
ayushiahjolia marked this conversation as resolved.

# The Workflow span (execution view) is a non-recording placeholder
# during the invocation, so only a terminal status materializes and ends
# the recording span. Its span status reflects the execution outcome:
# SUCCEEDED -> OK, FAILED -> ERROR (RETRY/PENDING are non-terminal and
# leave the Workflow span unexported until a later terminal invocation).
if info.status in _TERMINAL_INVOCATION_STATUSES:
self._export_workflow_span(info)

self._reset_state()

Expand All @@ -495,10 +586,12 @@ def _reset_state(self) -> None:
self._extracted_context = None
self._execution_trace_context = None
self._sampling_intent = None
self._execution_start_time = None
self._workflow_span = None
self._invocation_span = None
with self._lock:
self._operation_spans = {}
self._ended_operation_ids = set()
self._tracing_enabled = False

# ------------------------------------------------------------------
Expand All @@ -510,34 +603,31 @@ def on_operation_start(self, info: OperationStartInfo) -> None:
return
if info.operation_type is OperationType.CONTEXT:
return # tracked via on_user_function_start
# Hold a non-recording placeholder while the operation is open; its
# recording span is exported once on terminal on_operation_end.
self._register_operation_placeholder(info.operation_id)

def on_operation_end(self, info: OperationEndInfo) -> None:
logger.debug("Durable operation ended: %s", info)
if not self._tracing_enabled:
return
# Export the span only on the first end for this operation.
with self._lock:
if info.operation_id in self._ended_operation_ids:
return
self._ended_operation_ids.add(info.operation_id)
Comment on lines +615 to +618

This comment was marked as outdated.

Sign in to view

Comment on lines +615 to +618

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex AI review · Finding arf_v1_2pna33ayjyhsqx2ehym77e3n7t

[P1] Do not re-export terminal contexts during replay. This set is reset after every invocation, so ReplayChildren contexts calling on_child_context_end(..., is_replayed=True) export the same deterministic span again on each replay. Collectors may overwrite or reject the original span. Skip replayed terminal callbacks or otherwise make export uniqueness execution-wide, and add a multi-invocation ReplayChildren test asserting one export.

# An open operation is held as a non-recording placeholder; drop it and
# create the single recording span for the operation now.
self._pop_span(info.operation_id)
parent = self._resolve_parent(info.parent_id)
self._start_span(
span = self._start_span(

This comment was marked as outdated.

Sign in to view

@github-actions github-actions Bot Sep 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex AI review · Finding arf_v1_5yc45ob4ohc33i7vrvqz7pezbu

[P2] Preserve a start timestamp for checkpointless contexts. on_child_context_end() captures end_time but supplies start_time=None, so this span starts at materialization time and is then ended using an earlier timestamp. Virtual and ReplayChildren context spans can consequently have negative durations. Store the user-function start time alongside the placeholder or include it in the end hook, and test both paths.

operation_id=info.operation_id,
name=info.name or info.operation_id,
info=info,
parent=parent,
start_time=info.start_time,
)

def on_operation_end(self, info: OperationEndInfo) -> None:
logger.debug("Durable operation ended: %s", info)
if not self._tracing_enabled:
return
span = self._get_span(info.operation_id)
if span is None:
# Cross-invocation stitching: operation started in a prior
# invocation. Create + immediately end a linked span.
parent = self._resolve_parent(info.parent_id)
span = self._start_span(
operation_id=info.operation_id,
name=info.name or info.operation_id,
info=info,
parent=parent,
start_time=info.start_time,
)
else:
span.set_attributes(self._operation_attributes(info))

if info.error:
span.set_status(StatusCode.ERROR, info.error.message or "")
span.record_exception(
Expand All @@ -564,7 +654,11 @@ def _start_span(
span_key: str | None = None,
deterministic: bool = True,
) -> Span:
"""Start a span for an operation/attempt and register it."""
"""Start a recording span for an operation/attempt and register it.

Operation spans use the deterministic operation span ID; attempt spans
pass ``deterministic=False`` for a fresh ID beneath the operation span.
"""
key = span_key if span_key is not None else operation_id
with self._lock:
links = self._build_invocation_links()
Expand Down Expand Up @@ -603,6 +697,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
"on_user_function_start only supports CONTEXT and STEP operations"
)
key = self._user_function_key(info)
span: Span | None
if info.operation_type is OperationType.STEP:
parent = self._get_span(info.operation_id) or self._resolve_parent(
info.parent_id
Expand All @@ -618,17 +713,15 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
deterministic=False,
)
else: # CONTEXT
parent = self._resolve_parent(info.parent_id)
span = self._start_span(
operation_id=info.operation_id,
name=info.name or info.operation_id,
info=info,
parent=parent,
start_time=info.start_time,
# A child context can suspend before completing, so hold a
# non-recording placeholder while it runs; on_operation_end
# materializes its single recording span. This keeps a suspended
# context from being exported early and re-exported on replay.
span = self._register_operation_placeholder(info.operation_id)
if span is not None:
self._attach_context(
key, trace.set_span_in_context(span, otel_context.get_current())
)
self._attach_context(
key, trace.set_span_in_context(span, otel_context.get_current())
)

def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
logger.debug("Durable user function ended: %s", info)
Expand Down
Loading
Loading