diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py index 0159c721..8965aca7 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py @@ -7,7 +7,10 @@ import threading from typing import Any -from aws_durable_execution_sdk_python.lambda_service import OperationType +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationType, +) from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, InvocationEndInfo, @@ -25,6 +28,7 @@ Link, Span, SpanContext, + SpanKind, StatusCode, TraceFlags, Tracer, @@ -43,6 +47,8 @@ logger = logging.getLogger(__name__) +_SpanAttributes = dict[str, str | bool | int] + def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: """Convert a datetime to OTel timestamp (nanoseconds since epoch), or None.""" @@ -211,7 +217,7 @@ def _start_span( self, operation_id: str | None, name: str, - attributes: dict[str, str], + attributes: _SpanAttributes, start_time: datetime.datetime | None = None, parent_span: Span | None = None, existed: bool = False, @@ -282,6 +288,7 @@ def _start_span( ) span = self._tracer.start_span( name=name, + kind=SpanKind.INTERNAL, attributes=attributes, start_time=_to_otel_timestamp(start_time), context=parent_context, @@ -338,6 +345,18 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if operation_id: self._end_span(operation_id) + invocation_span = self._get_span(None) + if invocation_span: + invocation_span.set_attribute( + "durable.invocation.status", info.status.value + ) + if info.status is InvocationStatus.FAILED: + invocation_span.set_status( + StatusCode.ERROR, info.error.message if info.error else "" + ) + elif info.status is InvocationStatus.SUCCEEDED: + invocation_span.set_status(StatusCode.OK) + # end the invocation span self._end_span(None) @@ -512,7 +531,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: trace.set_span_in_context(parent_span, self._extracted_context) ) - def _extract_attributes(self, info: Any) -> dict[str, str]: + def _extract_attributes(self, info: Any) -> _SpanAttributes: """Extract durable execution fields as OpenTelemetry span attributes. Args: @@ -521,10 +540,12 @@ def _extract_attributes(self, info: Any) -> dict[str, str]: Returns: A dictionary of durable execution attributes suitable for a span. """ - attributes: dict[str, str] = { + attributes: _SpanAttributes = { "durable.execution.arn": self._execution_arn, } + if isinstance(info, InvocationStartInfo): + attributes["durable.invocation.first"] = info.is_first_invocation if hasattr(info, "operation_id") and info.operation_id is not None: attributes["durable.operation.id"] = info.operation_id if hasattr(info, "operation_type") and info.operation_type is not None: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py index 4bd650f6..a9e46322 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py @@ -27,6 +27,7 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( operation_id_to_span_id, @@ -65,24 +66,28 @@ def _create_plugin() -> tuple[OtelPlugin, InMemorySpanExporter]: return plugin, exporter -def _invocation_start_info() -> InvocationStartInfo: +def _invocation_start_info( + is_first_invocation: bool = True, +) -> InvocationStartInfo: """Create standard invocation start info for tests.""" return InvocationStartInfo( request_id="request-1", execution_arn=EXECUTION_ARN, execution_start_time=START_TIME, - is_first_invocation=True, + is_first_invocation=is_first_invocation, ) -def _invocation_end_info() -> InvocationEndInfo: +def _invocation_end_info( + status: InvocationStatus = InvocationStatus.SUCCEEDED, +) -> InvocationEndInfo: """Create standard invocation end info for tests.""" return InvocationEndInfo( request_id="request-1", execution_arn=EXECUTION_ARN, execution_start_time=START_TIME, is_first_invocation=True, - status=InvocationStatus.SUCCEEDED, + status=status, error=None, ) @@ -144,10 +149,55 @@ def test_invocation_start_and_end_emit_invocation_span(): spans = exporter.get_finished_spans() assert [span.name for span in spans] == ["invocation"] + assert spans[0].kind is SpanKind.INTERNAL assert spans[0].attributes["durable.execution.arn"] == EXECUTION_ARN + assert spans[0].attributes["durable.invocation.first"] is True + assert ( + spans[0].attributes["durable.invocation.status"] + == InvocationStatus.SUCCEEDED.value + ) assert plugin._get_span(None) is None +def test_invocation_span_records_subsequent_invocation(): + """Invocation spans preserve a false first-invocation attribute.""" + plugin, exporter = _create_plugin() + + plugin.on_invocation_start(_invocation_start_info(is_first_invocation=False)) + plugin.on_invocation_end(_invocation_end_info()) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["durable.invocation.first"] is False + + +@pytest.mark.parametrize( + ("invocation_status", "expected_span_status"), + [ + (InvocationStatus.PENDING, StatusCode.UNSET), + (InvocationStatus.RETRY, StatusCode.UNSET), + (InvocationStatus.SUCCEEDED, StatusCode.OK), + (InvocationStatus.FAILED, StatusCode.ERROR), + ], +) +def test_invocation_span_status_reflects_execution_status( + invocation_status: InvocationStatus, + expected_span_status: StatusCode, +): + """Only terminal invocation spans receive a success or failure status.""" + plugin, exporter = _create_plugin() + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(invocation_status)) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + attributes = spans[0].attributes + assert attributes is not None + assert attributes["durable.invocation.status"] == invocation_status.value + assert spans[0].status.status_code is expected_span_status + + def test_operation_callbacks_emit_child_span_with_deterministic_span_id(): """Verify non-user-function operations are traced beneath the invocation.""" plugin, exporter = _create_plugin() @@ -193,6 +243,7 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id(): plugin.on_invocation_end(_invocation_end_info()) spans_by_name = {span.name: span for span in exporter.get_finished_spans()} + assert all(span.kind is SpanKind.INTERNAL for span in spans_by_name.values()) wait_span = spans_by_name["wait-for-signal"] invocation_span = spans_by_name["invocation"] assert wait_span.context.span_id == operation_id_to_span_id( @@ -444,6 +495,8 @@ def test_step_operation_span_parents_attempt_span(): spans_by_name = {span.name: span for span in exporter.get_finished_spans()} finished_step_span = spans_by_name["fetch-user"] attempt_span = spans_by_name["fetch-user attempt 1"] + assert finished_step_span.kind is SpanKind.INTERNAL + assert attempt_span.kind is SpanKind.INTERNAL assert attempt_span.parent.span_id == finished_step_span.context.span_id assert ( finished_step_span.attributes["durable.operation.status"]