Skip to content

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

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

Draft
ayushiahjolia wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

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

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
103 changes: 77 additions & 26 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 @@ -49,11 +49,13 @@
from opentelemetry.sdk.trace import Tracer as SdkTracer
from opentelemetry.trace import (
Link,
NonRecordingSpan,
Span,
SpanContext,
SpanKind,
StatusCode,
Tracer,
TraceFlags,
)

from aws_durable_execution_sdk_python_otel.context_extractors import (
Expand Down Expand Up @@ -118,6 +120,7 @@ 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: Context | None = None
self._workflow_span: Span | None = None
self._invocation_span: Span | None = None
Expand Down Expand Up @@ -166,6 +169,22 @@ def _pop_span(self, key: str) -> Span | None:
with self._lock:
return self._operation_spans.pop(key, None)

def _end_open_operation_spans(self) -> None:
"""End every operation span still open at invocation end (except invocation).

Spans are registered parent-first, so ending them in reverse keeps each
child contained within its parent. The invocation span is ended
separately by the caller.
"""
with self._lock:
keys = list(reversed(self._operation_spans))
for key in keys:
if key == _INVOCATION_KEY:
continue
popped = self._pop_span(key)
if popped is not None:
popped.end()

@staticmethod
def _attempt_key(
info: UserFunctionStartInfo | UserFunctionEndInfo,
Expand Down Expand Up @@ -310,6 +329,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
self._execution_trace_id = _to_otel_trace_id(
self._execution_arn, info.execution_start_time
)
self._execution_start_time = info.execution_start_time
self._extracted_context = self._context_extractor(info)

self._start_workflow_span(info)
Expand All @@ -327,21 +347,60 @@ 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. Using a non-recording span for that role means 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
workflow_span_context = SpanContext(
trace_id=self._execution_trace_id or 0,
span_id=derive_workflow_span_id(self._execution_arn),
is_remote=False,
trace_flags=TraceFlags(TraceFlags.SAMPLED),
)
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 trace and span IDs as the placeholder so the
exported root correlates with every operation span across all
invocations, and anchors the span at the execution start time.
"""
if not self._execution_arn:
return
# Empty context => root span with no parent.
with self._id_generator.use_ids(
trace_id=self._execution_trace_id,
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=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 _start_invocation_span(self, info: InvocationStartInfo) -> None:
self._invocation_span = self._tracer.start_span(
Expand All @@ -361,12 +420,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 @@ -389,23 +442,20 @@ 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()
# Operation spans still open here belong to operations that suspended
# (e.g. PENDING/RETRYING) rather than completed this invocation. End them
# so no recording span is abandoned; the authoritative span for an
# operation that resumes in a later invocation is created and ended by
# on_operation_end at that time.
self._end_open_operation_spans()

# 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 @@ -419,6 +469,7 @@ def _reset_state(self) -> None:
self._detach_remaining_contexts()
self._execution_arn = ""
self._execution_trace_id = None
self._execution_start_time = None
self._extracted_context = None
self._workflow_span = None
self._invocation_span = None
Expand Down
75 changes: 49 additions & 26 deletions ...-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,13 @@
from opentelemetry.sdk.trace import Tracer as SdkTracer
from opentelemetry.trace import (
Link,
NonRecordingSpan,
Span,
SpanContext,
SpanKind,
StatusCode,
Tracer,
TraceFlags,
)

from aws_durable_execution_sdk_python_otel.context_extractors import (
Expand Down Expand Up @@ -119,6 +121,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None:
# per invocation status:
self._execution_arn = ""
self._execution_trace_id: int | None = None
self._execution_start_time: datetime.datetime | None = None
self._extracted_context: Context | None = None
self._workflow_span: Span | None = None
# Maps operation ID (None for root) to the active span.
Expand Down Expand Up @@ -432,6 +435,7 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
self._execution_trace_id = _to_otel_trace_id(
self._execution_arn, info.execution_start_time
)
self._execution_start_time = info.execution_start_time
self._extracted_context = self._context_extractor(info)

self._start_workflow_span(info)
Expand All @@ -443,32 +447,61 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None:
)

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

The Workflow span is a parentless root keyed to a deterministic span ID
derived from the execution ARN, so every invocation of the same durable
execution contributes to one Workflow span. It is exported once, on a
terminal invocation. Operation and attempt spans link to it while
remaining parented to the invocation span. It is created unconditionally
-- InvocationOtelPlugin has no default/owned tracer-provider distinction,
so the span is emitted whether the provider is the ambient (ADOT/global)
one or an explicitly supplied one.
terminal invocation, by :meth:`_export_workflow_span`. During each
invocation the plugin only needs the deterministic SpanContext so
operation and attempt spans can link to it while remaining parented to
the invocation span; a non-recording placeholder fills that role so a
non-terminal invocation never abandons a recording span.
"""
if not self._execution_arn:
logger.warning("No execution ARN; skipping Workflow span creation")
return
workflow_span_context = SpanContext(
trace_id=self._execution_trace_id or 0,
span_id=derive_workflow_span_id(self._execution_arn),
is_remote=False,
trace_flags=TraceFlags(TraceFlags.SAMPLED),
)
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 trace and span IDs as the placeholder so the
exported root correlates with every operation span across all
invocations, and anchors the span at the execution start time.
"""
if not self._execution_arn:
return
# Empty context => root span with no parent.
with self._id_generator.use_ids(
trace_id=self._execution_trace_id,
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=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 on_invocation_end(self, info: InvocationEndInfo) -> None:
"""Called at the end of each invocation. Ends the invocation span and flushes."""
Expand Down Expand Up @@ -505,23 +538,12 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None:
# end the invocation span
self._end_span(None)

# The Workflow span (execution view) is exported only on a terminal
# status; on non-terminal statuses its reference is dropped without
# ending it (so it is not exported yet). SUCCEEDED -> OK, FAILED -> ERROR;
# RETRY/PENDING are non-terminal and leave it unexported.
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()
# The Workflow span (execution view) is a non-recording placeholder
# during the invocation, so only a terminal status materializes and ends
# the recording span. SUCCEEDED -> OK, FAILED -> ERROR; RETRY/PENDING are
# non-terminal and leave it unexported until a later terminal invocation.
if info.status in _TERMINAL_INVOCATION_STATUSES:
self._export_workflow_span(info)

self._reset_state()

Expand All @@ -534,6 +556,7 @@ def _reset_state(self) -> None:
self._detach_remaining_contexts()
self._execution_arn = ""
self._execution_trace_id = None
self._execution_start_time = None
self._extracted_context = None
self._workflow_span = None
with self._operation_spans_lock:
Expand Down
Loading
Loading