Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -356,8 +356,6 @@ def on_operation_start(self, info: OperationStartInfo) -> None:

def on_operation_end(self, info: OperationEndInfo) -> None:
logger.debug("Durable operation ended: %s", info)
if info.operation_type is OperationType.CONTEXT:
return
span = self._get_span(info.operation_id)
if span is None:
# Cross-invocation stitching: operation started in a prior
Expand Down Expand Up @@ -477,25 +475,28 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
raise RuntimeError(
"on_user_function_end without matching on_user_function_start"
)
span.set_attributes(self._operation_attributes(info))
if info.outcome is UserFunctionOutcome.FAILED:
span.set_status(StatusCode.ERROR, info.error.message if info.error else "")
span.record_exception(
Exception(
(info.error.message or info.error.type)
if info.error
else "Unknown error"
if info.operation_type is OperationType.STEP:
span.set_attributes(self._operation_attributes(info))
if info.outcome is UserFunctionOutcome.FAILED:
span.set_status(
StatusCode.ERROR, info.error.message if info.error else ""
)
)
else:
span.set_status(StatusCode.OK)
span.record_exception(
Exception(
(info.error.message or info.error.type)
if info.error
else "Unknown error"
)
)
else:
span.set_status(StatusCode.OK)

end_time = info.end_time
if end_time is not None and end_time == info.start_time:
end_time += datetime.timedelta(microseconds=1)
popped = self._pop_span(key)
if popped is not None:
popped.end(end_time=_to_otel_timestamp(end_time))
end_time = info.end_time
if end_time is not None and end_time == info.start_time:
end_time += datetime.timedelta(microseconds=1)
popped = self._pop_span(key)
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 = (
Expand All @@ -519,7 +520,14 @@ def _operation_attributes(self, info: Any) -> dict[str, Any]:
attributes["durable.operation.type"] = info.operation_type.value
if getattr(info, "sub_type", None) is not None:
attributes["durable.operation.subtype"] = info.sub_type.value
if getattr(info, "status", None) is not None:
# STEP user-function spans represent attempts, not durable operations.
if (
not (
isinstance(info, (UserFunctionStartInfo, UserFunctionEndInfo))
and info.operation_type is OperationType.STEP
)
and getattr(info, "status", None) is not None
):
attributes["durable.operation.status"] = info.status.value
if getattr(info, "name", None) is not None:
attributes["durable.operation.name"] = info.name
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -470,9 +470,6 @@ def on_operation_end(self, info: OperationEndInfo) -> None:
operation span before being ended.
"""
logger.debug("Durable operation ended: %s", info)
if info.operation_type is OperationType.CONTEXT:
# Context operations are tracked using on_user_function_end.
return
span = self._get_span(info.operation_id)
if span is None:
# the span was not started in the current invocation, so we need to
Expand Down Expand Up @@ -548,9 +545,8 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None:
def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
"""Called when a context or step operation finishes user code.

This callback records the final attempt status, captures exceptions for
failed attempts, and ends the span that was attached in
``on_user_function_start``.
STEP attempt spans are finalized here. CONTEXT spans stay open until
``on_operation_end`` supplies the authoritative durable status.

Args:
info: Information about the operation attempt.
Expand All @@ -572,23 +568,26 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None:
"on_user_function_end called without matching on_user_function_start"
)

span.set_attributes(self._extract_attributes(info))
if info.outcome is UserFunctionOutcome.FAILED:
span.set_status(StatusCode.ERROR, info.error.message if info.error else "")
span.record_exception(
Exception(
(info.error.message or info.error.type)
if info.error
else "Unknown error"
if info.operation_type is OperationType.STEP:
span.set_attributes(self._extract_attributes(info))
if info.outcome is UserFunctionOutcome.FAILED:
span.set_status(
StatusCode.ERROR, info.error.message if info.error else ""
)
)
else:
span.set_status(StatusCode.OK)
span.record_exception(
Exception(
(info.error.message or info.error.type)
if info.error
else "Unknown error"
)
)
else:
span.set_status(StatusCode.OK)

end_timestamp = info.end_time
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)
end_timestamp = info.end_time
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
Expand Down Expand Up @@ -621,7 +620,15 @@ def _extract_attributes(self, info: Any) -> _SpanAttributes:
attributes["durable.operation.type"] = info.operation_type.value
if hasattr(info, "sub_type") and info.sub_type is not None:
attributes["durable.operation.subtype"] = info.sub_type.value
if hasattr(info, "status") and info.status is not None:
# STEP user-function spans represent attempts, not durable operations.
if (
not (
isinstance(info, (UserFunctionStartInfo, UserFunctionEndInfo))
and info.operation_type is OperationType.STEP
)
and hasattr(info, "status")
and info.status is not None
):
attributes["durable.operation.status"] = info.status.value
if hasattr(info, "name") and info.name is not None:
attributes["durable.operation.name"] = info.name
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import opentelemetry.context as otel_context
import pytest
from aws_durable_execution_sdk_python.lambda_service import (
ErrorObject,
InvocationStatus,
OperationStatus,
OperationSubType,
Expand All @@ -17,6 +18,9 @@
InvocationStartInfo,
OperationEndInfo,
OperationStartInfo,
UserFunctionEndInfo,
UserFunctionOutcome,
UserFunctionStartInfo,
)
from opentelemetry import trace
from opentelemetry.context import Context
Expand Down Expand Up @@ -235,6 +239,153 @@ def test_cross_invocation_operation_end_uses_deterministic_span_id():
)


@pytest.mark.parametrize(
("outcome", "terminal_status", "error", "expected_span_status"),
[
(
UserFunctionOutcome.SUCCEEDED,
OperationStatus.SUCCEEDED,
None,
trace.StatusCode.OK,
),
(
UserFunctionOutcome.SUCCEEDED,
OperationStatus.FAILED,
ErrorObject(
message="serialization failed",
type="SerializationError",
data=None,
stack_trace=None,
),
trace.StatusCode.ERROR,
),
],
)
def test_context_span_waits_for_terminal_operation_status(
outcome,
terminal_status,
error,
expected_span_status,
):
plugin, exporter = _create_plugin()
plugin.on_invocation_start(_invocation_start_info())
operation_id = "context-1"

plugin.on_user_function_start(
UserFunctionStartInfo(
operation_id=operation_id,
operation_type=OperationType.CONTEXT,
sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
name="book-trip",
parent_id=None,
start_time=START_TIME,
is_replayed=False,
status=OperationStatus.STARTED,
is_replay_children=False,
attempt=1,
)
)
active_span = plugin._get_span(operation_id)
assert active_span is not None
assert (
active_span.attributes["durable.operation.status"]
== OperationStatus.STARTED.value
)

plugin.on_user_function_end(
UserFunctionEndInfo(
operation_id=operation_id,
operation_type=OperationType.CONTEXT,
sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
name="book-trip",
parent_id=None,
start_time=START_TIME,
is_replayed=False,
status=OperationStatus.STARTED,
is_replay_children=False,
attempt=1,
outcome=outcome,
end_time=END_TIME,
error=None,
)
)

assert plugin._get_span(operation_id) is active_span
assert (
active_span.attributes["durable.operation.status"]
== OperationStatus.STARTED.value
)
assert not exporter.get_finished_spans()

plugin.on_operation_end(
OperationEndInfo(
operation_id=operation_id,
operation_type=OperationType.CONTEXT,
sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
name="book-trip",
parent_id=None,
start_time=START_TIME,
is_replayed=False,
status=terminal_status,
end_time=END_TIME,
error=error,
)
)

span = exporter.get_finished_spans()[0]
assert span.attributes["durable.operation.status"] == terminal_status.value
assert span.status.status_code is expected_span_status


def test_step_attempt_span_omits_operation_status():
plugin, exporter = _create_plugin()
plugin.on_invocation_start(_invocation_start_info())
operation_id = "step-1"

plugin.on_user_function_start(
UserFunctionStartInfo(
operation_id=operation_id,
operation_type=OperationType.STEP,
sub_type=OperationSubType.STEP,
name="fetch-user",
parent_id=None,
start_time=START_TIME,
is_replayed=False,
status=OperationStatus.STARTED,
is_replay_children=False,
attempt=1,
)
)
active_span = plugin._get_span("step-1:attempt:1")
assert active_span is not None
assert "durable.operation.status" not in active_span.attributes

plugin.on_user_function_end(
UserFunctionEndInfo(
operation_id=operation_id,
operation_type=OperationType.STEP,
sub_type=OperationSubType.STEP,
name="fetch-user",
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,
)
)

span = exporter.get_finished_spans()[0]
assert (
span.attributes["durable.attempt.outcome"]
== UserFunctionOutcome.SUCCEEDED.value
)
assert "durable.operation.status" not in span.attributes


# ---------------------------------------------------------------------------
# Default-provider mode: invocation span
# ---------------------------------------------------------------------------
Expand Down
Loading
Loading