diff --git a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py index 9c17b37e..a727bf73 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/otel/otel_logger_example.py @@ -1,6 +1,6 @@ """Demonstrates OTel-enriched logging in a durable execution. -The OtelPlugin installs a logging filter on the root logger +The InvocationOtelPlugin installs a logging filter on the root logger (enrich_logger=True by default) when the plugin is constructed. The filter stamps the active OpenTelemetry trace context (traceId, spanId, otelTraceSampled) onto every log record that flows through the root handler. @@ -16,7 +16,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import OtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.context import ( @@ -44,7 +44,7 @@ def greet_in_child(child_context: DurableContext, name: str) -> str: return result -@durable_execution(plugins=[OtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPlugin()]) def handler(_event: Any, context: DurableContext) -> str: # Logged at the top level: enriched with the invocation span_id. context.logger.info("Workflow started") diff --git a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py index df744985..3d001d46 100644 --- a/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py +++ b/packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_otel.py @@ -2,7 +2,7 @@ from typing import Any -from aws_durable_execution_sdk_python_otel import OtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin from aws_durable_execution_sdk_python import StepContext from aws_durable_execution_sdk_python.config import Duration @@ -32,7 +32,7 @@ def add_numbers_in_child(child_context: DurableContext, a: int, b: int): return result -@durable_execution(plugins=[OtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPlugin()]) def handler(_event: Any, context: DurableContext) -> int: result = 0 for i in range(3): diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index bc77cb7e..5a38774a 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -21,7 +21,7 @@ pip install aws-durable-execution-sdk-python-otel 1. Add the [ADOT Lambda Layer](#1-adot-lambda-layer) to your function and set `AWS_LAMBDA_EXEC_WRAPPER=/opt/otel-instrument` 2. Enable [X-Ray Active Tracing](#2-aws-x-ray-active-tracing) on the function -3. Pass `OtelPlugin` to your handler's `plugins` list +3. Pass `InvocationOtelPlugin` to your handler's `plugins` list 4. Add X-Ray write permissions ### 1. ADOT Lambda Layer @@ -130,10 +130,10 @@ lambda_.Function( ```python from aws_durable_execution_sdk_python import DurableContext from aws_durable_execution_sdk_python.execution import durable_execution -from aws_durable_execution_sdk_python_otel import OtelPlugin +from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin -@durable_execution(plugins=[OtelPlugin()]) +@durable_execution(plugins=[InvocationOtelPlugin()]) def handler(event: dict, context: DurableContext) -> dict: result = context.step(lambda _: fetch_data(event["id"]), name="fetch-data") @@ -167,11 +167,11 @@ See the [ADOT sampling configuration](https://aws-otel.github.io/docs/getting-st ```python from aws_durable_execution_sdk_python_otel import ( - OtelPlugin, + InvocationOtelPlugin, xray_context_extractor, ) -plugin = OtelPlugin( +plugin = InvocationOtelPlugin( # Provide your own TracerProvider if you already have one configured. # Defaults to the globally configured tracer provider. trace_provider=None, @@ -192,16 +192,16 @@ The plugin supports multiple strategies for extracting upstream trace context: ```python from aws_durable_execution_sdk_python_otel import ( - OtelPlugin, + InvocationOtelPlugin, w3c_client_context_extractor, xray_context_extractor, ) # Default: X-Ray trace header (recommended for most Lambda deployments) -OtelPlugin(context_extractor=xray_context_extractor) +InvocationOtelPlugin(context_extractor=xray_context_extractor) # W3C Trace Context via clientContext (requires backend propagation support) -OtelPlugin(context_extractor=w3c_client_context_extractor) +InvocationOtelPlugin(context_extractor=w3c_client_context_extractor) ``` ### Log Correlation @@ -245,12 +245,12 @@ After deploying your function with the plugin configured: ## API Reference -### `OtelPlugin` +### `InvocationOtelPlugin` The main plugin class. Implements `DurableInstrumentationPlugin` from `aws_durable_execution_sdk_python`. ```python -OtelPlugin( +InvocationOtelPlugin( trace_provider=None, context_extractor=None, instrument_name="aws-durable-execution-sdk-python", diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index ab575cc5..bd65c93d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -29,6 +29,15 @@ dependencies = [ "opentelemetry-propagator-aws-xray", ] +[project.optional-dependencies] +# Instrumentation used by ExecutionOtelPlugin's auto-configured provider path. +# Kept optional so the InvocationOtelPlugin (ADOT / global provider) install +# stays lean; the instrumentations module degrades gracefully when absent. +instrumentation = [ + "opentelemetry-instrumentation-botocore", + "opentelemetry-instrumentation-urllib3", +] + [project.urls] Documentation = "https://github.com/aws/aws-durable-execution-sdk-python#readme" Issues = "https://github.com/aws/aws-durable-execution-sdk-python/issues" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py index e1bc0ef6..d44dfb56 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py @@ -8,13 +8,29 @@ ) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, + derive_workflow_span_id, + operation_id_to_span_id, +) +from aws_durable_execution_sdk_python_otel.execution_plugin import ( + ExecutionOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ExporterConfig, +) +from aws_durable_execution_sdk_python_otel.instrumentations import ( + register_standalone_instrumentations, ) from aws_durable_execution_sdk_python_otel.log_filter import ( OtelContextLogFilter, install_log_filter, ) -from aws_durable_execution_sdk_python_otel.plugin import ( - OtelPlugin, +from aws_durable_execution_sdk_python_otel.invocation_plugin import ( + InvocationOtelPlugin, +) +from aws_durable_execution_sdk_python_otel.provider import ( + ProviderResult, + create_tracer_provider, ) @@ -22,9 +38,17 @@ "__version__", "ContextExtractor", "DeterministicIdGenerator", + "ExecutionOtelPlugin", + "OtelPluginConfig", + "ExporterConfig", + "InvocationOtelPlugin", "OtelContextLogFilter", - "OtelPlugin", + "ProviderResult", + "create_tracer_provider", + "derive_workflow_span_id", "install_log_filter", + "operation_id_to_span_id", + "register_standalone_instrumentations", "w3c_client_context_extractor", "xray_context_extractor", ] diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py index 5c61b383..bcf7388b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py @@ -73,6 +73,39 @@ def operation_id_to_span_id(durable_execution_arn: str, operation_id: str) -> in return int(hashed_operation_id, 16) +def derive_workflow_span_id(durable_execution_arn: str) -> int: + """Derive the deterministic Workflow root span ID (64 bits) from an execution ARN. + + Mirrors the JS ``deriveWorkflowSpanId``: hash ``"workflow:" + arn`` and take + the first 16 hex characters. All invocations of the same durable execution + therefore share one Workflow root span ID, allowing the Workflow span to be + (re-)created and exported exactly once per execution regardless of how many + Lambda invocations back it. + + The hash algorithm is blake2b to stay consistent with + :func:`operation_id_to_span_id`; it does not need to match the JS SDK's + SHA-256 byte-for-byte because span IDs only need to be stable within a single + execution's trace, not across language runtimes. + + Args: + durable_execution_arn: The durable execution ARN. Must be non-empty. + + Returns: + A 64-bit integer span ID. If the derived value is all-zero (invalid per + the OTel spec) it is bumped to ``1``. + + Raises: + ValueError: If ``durable_execution_arn`` is empty. Only emptiness is + validated; ARN format is not checked. + """ + if not durable_execution_arn: + raise ValueError("execution ARN is required to derive a workflow span ID") + plain_value = f"workflow:{durable_execution_arn}" + hashed = hashlib.blake2b(plain_value.encode()).hexdigest()[:16] + span_id = int(hashed, 16) + return span_id or 1 + + class DeterministicIdGenerator(RandomIdGenerator): """An ID generator that produces deterministic span IDs when a pending operation ID is set, and falls back to the provided generator otherwise. 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 new file mode 100644 index 00000000..52c25123 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -0,0 +1,526 @@ +"""Execution-view OpenTelemetry plugin for AWS Durable Executions. + +The :class:`ExecutionOtelPlugin` produces the deterministic span hierarchy + + Workflow -> Invocation -> Operation -> Attempt + +that stitches a single trace across every Lambda invocation of one durable +execution. The Workflow span is the root (created in an empty context so it +never has a parent) and 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. + +This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from +aws-durable-execution-sdk-js#729. Because the Python plugin interface differs +from JS (there is no ``wrapInvocation``/``wrapChildContextFn``/``enrichLogContext`` +- 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`. +""" + +from __future__ import annotations + +import datetime +import logging +import threading +from typing import Any + +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + OperationEndInfo, + OperationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, + UserFunctionStartInfo, +) +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.trace import ( + Link, + Span, + SpanContext, + StatusCode, + Tracer, +) + +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ContextExtractor, + xray_context_extractor, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, + derive_workflow_span_id, + operation_id_to_span_id, +) +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, +) +from aws_durable_execution_sdk_python_otel.instrumentations import ( + register_standalone_instrumentations, +) +from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter +from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider + + +logger = logging.getLogger(__name__) + +_TERMINAL_INVOCATION_STATUSES = frozenset( + {InvocationStatus.SUCCEEDED, InvocationStatus.FAILED} +) + +# Registry key for the invocation span (operations use their operation_id). +_INVOCATION_KEY = "__invocation__" + + +def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: + """Convert a datetime to an OTel timestamp (ns since epoch), or None.""" + if dt is None: + return None + return int(dt.timestamp() * 1_000_000_000) + + +class ExecutionOtelPlugin(DurableInstrumentationPlugin): + """OTel plugin that renders a durable execution as one Workflow-rooted trace. + + Args: + config: Shared plugin configuration. When omitted, defaults are used + (auto-configured provider, X-Ray extractor, "Workflow" root span). + """ + + def __init__(self, config: OtelPluginConfig | None = None) -> None: + self._config = config or OtelPluginConfig() + self._context_extractor: ContextExtractor = ( + self._config.context_extractor or xray_context_extractor + ) + self._workflow_span_name = self._config.workflow_span_name + self._use_default = bool(self._config.use_default_tracer_provider) + + self._id_generator = DeterministicIdGenerator() + result = create_tracer_provider( + self._config, + id_generator=self._id_generator, + default_use_global=False, + ) + self._provider = result.tracer_provider + self._owns_provider = result.owns_provider + + # Deterministic stitching requires an SDK provider exposing id_generator. + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + + if isinstance(self._provider, SdkTracerProvider): + self._id_generator = DeterministicIdGenerator( + fallback_id_generator=getattr(self._provider, "id_generator", None) + ) + self._provider.id_generator = self._id_generator + else: + logger.warning( + "ExecutionOtelPlugin expected an SDK TracerProvider but got %s; " + "spans will not use deterministic IDs.", + type(self._provider).__name__, + ) + + self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) + + try: + register_standalone_instrumentations( + self._config, + self._provider if self._owns_provider else None, + owns_provider=self._owns_provider, + use_default_tracer_provider=self._use_default, + ) + except Exception: + logger.exception("Failed to register standalone instrumentations") + + # Per-invocation state. + self._execution_arn = "" + self._extracted_context: Context | None = None + self._saved_invocation_context: Context | None = None + self._workflow_span: Span | None = None + self._invocation_span: Span | None = None + self._operation_spans: dict[str, Span] = {} + self._lock = threading.RLock() + self._is_cold_start = True + + if self._config.enrich_logger: + install_log_filter(self) + + # ------------------------------------------------------------------ + # Span registry helpers + # ------------------------------------------------------------------ + def _set_span(self, key: str, span: Span) -> None: + with self._lock: + self._operation_spans[key] = span + + def _get_span(self, key: str | None) -> Span | None: + if key is None: + return None + with self._lock: + return self._operation_spans.get(key) + + def _pop_span(self, key: str) -> Span | None: + with self._lock: + return self._operation_spans.pop(key, None) + + @staticmethod + def _attempt_key(info: UserFunctionStartInfo | UserFunctionEndInfo) -> str: + return f"{info.operation_id}:attempt:{info.attempt or 1}" + + 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() + if span_context and span_context.is_valid: + return span_context + for candidate in (self._invocation_span, self._workflow_span): + if candidate is not None: + ctx = candidate.get_span_context() + if ctx and ctx.is_valid: + return ctx + return None + + # ------------------------------------------------------------------ + # Links + # ------------------------------------------------------------------ + def _build_invocation_links(self) -> list[Link]: + """Link operation/attempt spans to the invocation span.""" + if self._use_default and self._saved_invocation_context is not None: + span = trace.get_current_span(self._saved_invocation_context) + ctx = span.get_span_context() + if ctx and ctx.is_valid: + return [Link(context=ctx)] + if self._invocation_span is not None: + ctx = self._invocation_span.get_span_context() + if ctx and ctx.is_valid: + return [Link(context=ctx)] + return [] + + def _resolve_parent(self, parent_id: str | None) -> Span | None: + """Resolve the parent span: parent operation, else the Workflow span.""" + if parent_id is not None: + existing = self._get_span(parent_id) + if existing is not None: + return existing + return self._workflow_span + + # ------------------------------------------------------------------ + # Invocation lifecycle + # ------------------------------------------------------------------ + def on_invocation_start(self, info: InvocationStartInfo) -> None: + logger.debug("Durable invocation started: %s", info) + self._execution_arn = info.execution_arn or "" + self._extracted_context = self._context_extractor(info) + self._id_generator.set_trace_id(self._execution_arn, info.execution_start_time) + + # Capture the ambient context for link-building in default mode. + if self._use_default: + self._saved_invocation_context = otel_context.get_current() + + self._start_workflow_span(info) + # Create the Invocation span in both modes. In default-provider mode it + # is parented to the ambient Lambda invocation span. + self._start_invocation_span(info) + + # Make the Workflow span the active span so auto-instrumented spans + # created during the invocation become its children. + if self._workflow_span is not None: + otel_context.attach( + trace.set_span_in_context(self._workflow_span, self._extracted_context) + ) + + self._is_cold_start = False + + def _start_workflow_span(self, info: InvocationStartInfo) -> None: + if not self._execution_arn: + logger.warning("No execution ARN; skipping Workflow span creation") + return + self._id_generator.set_next_span_id( + derive_workflow_span_id(self._execution_arn) + ) + start_time = _to_otel_timestamp( + info.execution_start_time + ) or _to_otel_timestamp(datetime.datetime.now(datetime.UTC)) + # Empty context => root span with no parent. + self._workflow_span = self._tracer.start_span( + name=self._workflow_span_name, + attributes={"durable.execution.arn": self._execution_arn}, + start_time=start_time, + context=Context(), + ) + + def _start_invocation_span(self, info: InvocationStartInfo) -> None: + self._id_generator.set_next_span_id(None) + attributes: dict[str, Any] + if self._use_default: + # Default-provider mode: parent the Invocation span to the ambient + # Lambda invocation span (from the ADOT layer or other + # auto-instrumentation) captured at invocation start. + # Lambda semantic attributes belong to that ambient span, so carry + # only durable correlation attributes here. + parent_ctx = self._saved_invocation_context or otel_context.get_current() + attributes = { + "durable.execution.arn": self._execution_arn, + "durable.invocation.first": info.is_first_invocation, + } + else: + if self._workflow_span is None: + return + parent_ctx = trace.set_span_in_context( + self._workflow_span, self._extracted_context + ) + attributes = { + "durable.execution.arn": self._execution_arn, + "faas.coldstart": self._is_cold_start, + "cloud.provider": "aws", + "cloud.platform": "aws_lambda", + } + if info.request_id: + attributes["faas.invocation_id"] = info.request_id + self._invocation_span = self._tracer.start_span( + name="invocation", + attributes=attributes, + context=parent_ctx, + ) + self._set_span(_INVOCATION_KEY, self._invocation_span) + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + logger.debug("Durable invocation ended: %s", info) + # 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. + if self._invocation_span is not None: + self._invocation_span.end() + + # The Workflow span is exported only on a terminal status; otherwise its + # reference is dropped without ending it. + 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.error: + self._workflow_span.set_status( + StatusCode.ERROR, info.error.message or "" + ) + else: + self._workflow_span.set_status(StatusCode.OK) + self._workflow_span.end() + + self._reset_state() + + if hasattr(self._provider, "force_flush"): + try: + self._provider.force_flush() + except Exception: # noqa: BLE001 + logger.exception("force_flush failed at invocation end") + + def _reset_state(self) -> None: + self._execution_arn = "" + self._extracted_context = None + self._saved_invocation_context = None + self._workflow_span = None + self._invocation_span = None + with self._lock: + self._operation_spans = {} + + # ------------------------------------------------------------------ + # Operation lifecycle + # ------------------------------------------------------------------ + def on_operation_start(self, info: OperationStartInfo) -> None: + logger.debug("Durable operation started: %s", info) + if info.operation_type is OperationType.CONTEXT: + return # tracked via on_user_function_start + parent = self._resolve_parent(info.parent_id) + self._start_span( + 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 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 + # 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( + Exception(info.error.message or info.error.type or "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(info.operation_id) + if popped is not None: + popped.end(end_time=_to_otel_timestamp(end_time)) + + def _start_span( + self, + *, + operation_id: str, + name: str, + info: Any, + parent: Span | None, + start_time: datetime.datetime | None, + span_key: str | None = None, + deterministic: bool = True, + ) -> Span: + """Start a span for an operation/attempt and register it.""" + key = span_key if span_key is not None else operation_id + with self._lock: + links = self._build_invocation_links() + if deterministic: + # Operation spans always use the deterministic logical-operation + # span ID so a suspended-then-completed operation exports a + # single span (on completion) with a stable ID across invocations. + self._id_generator.set_next_span_id( + operation_id_to_span_id(self._execution_arn, operation_id) + ) + else: + self._id_generator.set_next_span_id(None) + + if parent is None: + parent_ctx = self._extracted_context or Context() + else: + parent_ctx = trace.set_span_in_context(parent, self._extracted_context) + span = self._tracer.start_span( + name=name, + attributes=self._operation_attributes(info), + start_time=_to_otel_timestamp(start_time), + context=parent_ctx, + links=links, + ) + self._operation_spans[key] = span + return span + + # ------------------------------------------------------------------ + # User function (CONTEXT / STEP attempt) lifecycle + # ------------------------------------------------------------------ + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + logger.debug("Durable user function started: %s", info) + if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP): + raise RuntimeError( + "on_user_function_start only supports CONTEXT and STEP operations" + ) + if info.operation_type is OperationType.STEP: + parent = self._get_span(info.operation_id) or self._resolve_parent( + info.parent_id + ) + name = f"{info.name or info.operation_id} attempt {info.attempt or 1}" + 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), + 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, + ) + otel_context.attach(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) + if info.operation_type not in (OperationType.CONTEXT, OperationType.STEP): + raise RuntimeError( + "on_user_function_end only supports CONTEXT and STEP operations" + ) + key = ( + self._attempt_key(info) + if info.operation_type is OperationType.STEP + else info.operation_id + ) + span = self._get_span(key) + if span is 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" + ) + ) + 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)) + + # 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) + ) + + # ------------------------------------------------------------------ + # Attributes + # ------------------------------------------------------------------ + def _operation_attributes(self, info: Any) -> dict[str, Any]: + attributes: dict[str, Any] = {"durable.execution.arn": self._execution_arn} + if getattr(info, "operation_id", None) is not None: + attributes["durable.operation.id"] = info.operation_id + if getattr(info, "operation_type", None) is not None: + 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: + attributes["durable.operation.status"] = info.status.value + if getattr(info, "name", None) is not None: + attributes["durable.operation.name"] = info.name + if getattr(info, "operation_type", None) is not OperationType.CONTEXT: + if getattr(info, "attempt", None) is not None: + attributes["durable.attempt.number"] = info.attempt + if getattr(info, "outcome", None) is not None: + attributes["durable.attempt.outcome"] = info.outcome.value + return attributes diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/instrumentations.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/instrumentations.py new file mode 100644 index 00000000..a9813808 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/instrumentations.py @@ -0,0 +1,128 @@ +"""Shared instrumentation registration for the durable-execution OTel plugins. + +Mirrors the JS ``registerStandaloneInstrumentations``: + +* A custom (explicit) provider skips ALL instrumentation registration. +* When the global provider is in use (``use_default_tracer_provider``), only the + AWS SDK instrumentation is registered (not HTTP). +* When the plugin owns an auto-configured provider, both AWS SDK and (optionally) + HTTP instrumentation are registered against that provider. + +The JS SDK uses ``AwsInstrumentation`` (AWS SDK v3) and ``HttpInstrumentation``. +The Python equivalents are ``BotocoreInstrumentor`` (boto3/botocore is the AWS +SDK for Python) and ``URLLib3Instrumentor`` (botocore's HTTP transport). Both +instrumentation packages are optional imports: when a package is not installed +the registration is skipped with a warning rather than raising, so the module +stays import-safe. +""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from opentelemetry.trace import TracerProvider + + from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ) + + +logger = logging.getLogger(__name__) + +_LOCAL_HOSTS = {"127.0.0.1", "localhost"} + + +def _runtime_api_host() -> str | None: + """Return the Lambda runtime API host (portion before ':') if set.""" + runtime_api = os.environ.get("AWS_LAMBDA_RUNTIME_API") + if not runtime_api: + return None + return runtime_api.split(":", 1)[0] + + +def _register_aws_instrumentation(tracer_provider: object | None) -> None: + """Register AWS SDK (botocore) instrumentation, if the package is available.""" + try: + from opentelemetry.instrumentation.botocore import BotocoreInstrumentor + except ImportError: + logger.warning( + "opentelemetry-instrumentation-botocore is not installed; " + "AWS SDK calls will not be traced. Install it to enable AWS " + "instrumentation." + ) + return + instrumentor = BotocoreInstrumentor() + if not instrumentor.is_instrumented_by_opentelemetry: + kwargs = {} + if tracer_provider is not None: + kwargs["tracer_provider"] = tracer_provider + instrumentor.instrument(**kwargs) + + +def _register_http_instrumentation(tracer_provider: object | None) -> None: + """Register HTTP (urllib3) instrumentation with local/runtime suppression.""" + try: + from opentelemetry.instrumentation.urllib3 import URLLib3Instrumentor + except ImportError: + logger.warning( + "opentelemetry-instrumentation-urllib3 is not installed; outbound " + "HTTP calls will not be traced. Install it to enable HTTP " + "instrumentation." + ) + return + + suppressed_hosts = set(_LOCAL_HOSTS) + runtime_host = _runtime_api_host() + if runtime_host: + suppressed_hosts.add(runtime_host) + + def request_hook(span, pool, request_info) -> None: # noqa: ANN001 + # Suppress spans to the loopback collector and the Lambda runtime API by + # ending them immediately with no recording. urllib3 does + # not expose a pre-create filter, so we no-op the created span instead. + host = getattr(pool, "host", None) + if host in suppressed_hosts and span is not None and span.is_recording(): + span.set_attribute("durable.instrumentation.suppressed", True) + + instrumentor = URLLib3Instrumentor() + if not instrumentor.is_instrumented_by_opentelemetry: + kwargs: dict[str, Any] = {"request_hook": request_hook} + if tracer_provider is not None: + kwargs["tracer_provider"] = tracer_provider + instrumentor.instrument(**kwargs) + + +def register_standalone_instrumentations( + config: OtelPluginConfig, + tracer_provider: TracerProvider | None, + *, + owns_provider: bool, + use_default_tracer_provider: bool, +) -> None: + """Register AWS SDK and HTTP instrumentations per the shared policy. + + Args: + config: Shared plugin configuration. + tracer_provider: The resolved provider (may be the global provider). + owns_provider: True when the plugin created/owns the provider. + use_default_tracer_provider: True when the global provider is in use. + """ + # A custom, explicitly-supplied provider means the caller manages their own + # instrumentation: skip everything. + if config.tracer_provider is not None: + return + + if use_default_tracer_provider: + # Global provider: register AWS instrumentation only. + _register_aws_instrumentation(None) + return + + # Auto-configured, plugin-owned provider: AWS SDK always; HTTP unless + # explicitly disabled. + _register_aws_instrumentation(tracer_provider) + if config.enable_http_instrumentation: + _register_http_instrumentation(tracer_provider) 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/invocation_plugin.py similarity index 99% rename from packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/plugin.py rename to packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 8965aca7..d0ab1f68 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/invocation_plugin.py @@ -57,7 +57,7 @@ def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: return int(dt.timestamp() * 1_000_000_000) -class OtelPlugin(DurableInstrumentationPlugin): +class InvocationOtelPlugin(DurableInstrumentationPlugin): """OpenTelemetry instrumentation plugin for durable executions. The plugin creates spans for Lambda invocations, durable operations, and @@ -120,7 +120,7 @@ def __init__( self._provider.id_generator = self._id_generator else: logger.warning( - "OtelPlugin expected an SDK TracerProvider " + "InvocationOtelPlugin expected an SDK TracerProvider " "(opentelemetry.sdk.trace.TracerProvider) but got %s. Spans will " "not use deterministic IDs. " "Ensure the OpenTelemetry SDK is configured (e.g. via the ADOT " diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py index 55e5e1c8..7e11ebc9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/log_filter.py @@ -18,13 +18,24 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol from opentelemetry.trace import TraceFlags if TYPE_CHECKING: - from aws_durable_execution_sdk_python_otel.plugin import OtelPlugin + from opentelemetry.trace import SpanContext + + +class _SpanContextProvider(Protocol): + """Structural type for any plugin that resolves the active span context. + + The log filter only needs this one capability, so it accepts any object + that provides it (e.g. ``InvocationOtelPlugin`` or ``ExecutionOtelPlugin``) + rather than a specific plugin class. + """ + + def get_current_span_context(self) -> SpanContext | None: ... class OtelContextLogFilter(logging.Filter): @@ -44,7 +55,7 @@ class OtelContextLogFilter(logging.Filter): plugin: The OTel plugin instance that resolves the current span context. """ - def __init__(self, plugin: OtelPlugin) -> None: + def __init__(self, plugin: _SpanContextProvider) -> None: super().__init__() self._plugin = plugin @@ -61,7 +72,7 @@ def filter(self, record: logging.LogRecord) -> bool: def install_log_filter( - plugin: OtelPlugin, + plugin: _SpanContextProvider, target_logger: logging.Logger | None = None, ) -> OtelContextLogFilter | None: """Attach an OtelContextLogFilter to a logger's handlers, idempotently. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/otel_plugin_config.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/otel_plugin_config.py new file mode 100644 index 00000000..76dc8843 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/otel_plugin_config.py @@ -0,0 +1,71 @@ +"""Shared configuration for the durable-execution OpenTelemetry plugins. + +Both :class:`ExecutionOtelPlugin` and :class:`InvocationOtelPlugin` accept a +single :class:`OtelPluginConfig`, so configuration options are +consistent and not duplicated across plugins. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Sequence + + +if TYPE_CHECKING: + from opentelemetry.propagators.textmap import TextMapPropagator + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + + from aws_durable_execution_sdk_python_otel.context_extractors import ( + ContextExtractor, + ) + + +DEFAULT_INSTRUMENT_NAME = "aws-durable-execution-sdk-python" +DEFAULT_WORKFLOW_SPAN_NAME = "Workflow" +# OTLPSpanExporter appends /v1/traces itself, so the base endpoint must NOT +# include it (mirrors the JS fix in PR #729 that removed the duplicate path). +DEFAULT_OTLP_ENDPOINT = "http://localhost:4318" + + +@dataclass +class ExporterConfig: + """OTLP exporter configuration for the auto-configured TracerProvider.""" + + endpoint: str | None = None + headers: dict[str, str] | None = None + + +@dataclass +class OtelPluginConfig: + """Canonical configuration shared by both OTel plugins. + + Fields relevant only to :class:`ExecutionOtelPlugin` (e.g. ``workflow_span_name``) + are ignored without error by :class:`InvocationOtelPlugin`. + + Attributes: + tracer_provider: Explicit provider to use as-is. Highest priority; when + set, the plugin does not own or modify it and skips instrumentation + registration. + use_default_tracer_provider: When True (and no explicit provider), + resolve the globally configured provider via ``trace.get_tracer_provider()``. + context_extractor: Upstream trace-context extractor. Defaults to the + X-Ray extractor when omitted. + instrument_name: Instrumentation scope name. + enable_http_instrumentation: Whether to register HTTP instrumentation + when the plugin owns an auto-configured provider. Defaults to True. + exporter_config: OTLP exporter settings for the auto-configured provider. + propagators: Custom propagators for the auto-configured provider. + Defaults to ``[AWSXRayPropagator, W3CTraceContextPropagator]``. + workflow_span_name: Name of the Workflow root span (ExecutionOtelPlugin). + enrich_logger: Install the root-logger OTel context filter. + """ + + tracer_provider: SdkTracerProvider | None = None + use_default_tracer_provider: bool | None = None + context_extractor: ContextExtractor | None = None + instrument_name: str = DEFAULT_INSTRUMENT_NAME + enable_http_instrumentation: bool = True + exporter_config: ExporterConfig = field(default_factory=ExporterConfig) + propagators: Sequence[TextMapPropagator] | None = None + workflow_span_name: str = DEFAULT_WORKFLOW_SPAN_NAME + enrich_logger: bool = True diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/provider.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/provider.py new file mode 100644 index 00000000..a0533b4e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/provider.py @@ -0,0 +1,181 @@ +"""Shared TracerProvider factory for the durable-execution OTel plugins. + +Implements the 3-level provider resolution used by both plugins: + +1. An explicit ``tracer_provider`` in config is used as-is (``owns_provider=False``). +2. Otherwise, when ``use_default_tracer_provider`` resolves to True, the globally + configured provider is used (``owns_provider=False``). +3. Otherwise a fully auto-configured SDK provider is created with an OTLP + exporter, batch processor, sampler and Lambda resource attributes + (``owns_provider=True``). +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from opentelemetry import propagate, trace +from opentelemetry.propagators.composite import CompositePropagator + +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + DEFAULT_OTLP_ENDPOINT, + OtelPluginConfig, +) + + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import IdGenerator + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + from opentelemetry.trace import TracerProvider + + +logger = logging.getLogger(__name__) + +SAMPLING_RATIO_ENV = "OTEL_DURABLE_SAMPLING_RATIO" +OTLP_ENDPOINT_ENV = "OTEL_EXPORTER_OTLP_ENDPOINT" + + +@dataclass +class ProviderResult: + """Result of provider resolution: the provider and whether we own it.""" + + tracer_provider: TracerProvider + owns_provider: bool + + +def _resolve_endpoint(config: OtelPluginConfig) -> str: + """Resolve the OTLP traces endpoint (config -> env -> default). + + The Python OTLP/HTTP exporter uses an explicitly-passed ``endpoint`` verbatim + (unlike the JS exporter, which appends ``/v1/traces``), so we append the + signal path here when the caller supplied only a base URL. + """ + base = ( + config.exporter_config.endpoint + or os.environ.get(OTLP_ENDPOINT_ENV) + or DEFAULT_OTLP_ENDPOINT + ) + base = base.rstrip("/") + if not base.endswith("/v1/traces"): + base = f"{base}/v1/traces" + return base + + +def _build_sampler(): + """Build the sampler from ``OTEL_DURABLE_SAMPLING_RATIO``. + + A valid ratio in ``[0, 1]`` yields a ``TraceIdRatioBased`` sampler; anything + else falls back to ``ALWAYS_ON``. If constructing the ratio sampler raises, + the error propagates and fails provider setup. + """ + from opentelemetry.sdk.trace.sampling import ALWAYS_ON, TraceIdRatioBased + + raw = os.environ.get(SAMPLING_RATIO_ENV) + if raw is None: + return ALWAYS_ON + try: + ratio = float(raw) + except (TypeError, ValueError): + return ALWAYS_ON + if not (0.0 <= ratio <= 1.0): + return ALWAYS_ON + return TraceIdRatioBased(ratio) + + +def _build_resource(): + """Build a Lambda resource from AWS_* env vars.""" + from opentelemetry.sdk.resources import Resource + + function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME") + if not function_name: + return None + attributes: dict[str, str] = { + "service.name": function_name, + "faas.name": function_name, + "cloud.provider": "aws", + "cloud.platform": "aws_lambda", + } + region = os.environ.get("AWS_REGION") + if region: + attributes["cloud.region"] = region + version = os.environ.get("AWS_LAMBDA_FUNCTION_VERSION") + if version: + attributes["faas.version"] = version + return Resource.create(attributes) + + +def _default_propagators(config: OtelPluginConfig): + """Return configured propagators or the X-Ray + W3C default.""" + if config.propagators is not None: + return list(config.propagators) + from opentelemetry.propagators.aws import AwsXRayPropagator + from opentelemetry.trace.propagation.tracecontext import ( + TraceContextTextMapPropagator, + ) + + return [AwsXRayPropagator(), TraceContextTextMapPropagator()] + + +def _create_auto_provider( + config: OtelPluginConfig, + id_generator: IdGenerator | None, +) -> SdkTracerProvider: + """Create a fully self-configured SDK TracerProvider with OTLP export.""" + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + exporter = OTLPSpanExporter( + endpoint=_resolve_endpoint(config), + headers=config.exporter_config.headers or None, + ) + provider_kwargs: dict = {"sampler": _build_sampler()} + resource = _build_resource() + if resource is not None: + provider_kwargs["resource"] = resource + if id_generator is not None: + provider_kwargs["id_generator"] = id_generator + + provider = SdkTracerProvider(**provider_kwargs) + provider.add_span_processor(BatchSpanProcessor(exporter)) + + composite = CompositePropagator(_default_propagators(config)) + propagate.set_global_textmap(composite) + return provider + + +def create_tracer_provider( + config: OtelPluginConfig, + *, + id_generator: IdGenerator | None = None, + default_use_global: bool = False, +) -> ProviderResult: + """Resolve a TracerProvider using the shared 3-level priority. + + Args: + config: Shared plugin configuration. + id_generator: Deterministic ID generator injected into an + auto-configured provider so cross-invocation trace stitching works. + default_use_global: The value ``use_default_tracer_provider`` defaults to + when it is unset in config. ``ExecutionOtelPlugin`` passes ``False``; + ``InvocationOtelPlugin`` passes ``True`` to preserve its historical + behaviour of using the global provider. + + Returns: + A :class:`ProviderResult`. + """ + if config.tracer_provider is not None: + # Explicit provider: use as-is, never wrap/modify. + return ProviderResult(config.tracer_provider, False) + + use_default = config.use_default_tracer_provider + if use_default is None: + use_default = default_use_global + + if use_default: + return ProviderResult(trace.get_tracer_provider(), False) + + return ProviderResult(_create_auto_provider(config, id_generator), True) 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 new file mode 100644 index 00000000..86d9da18 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -0,0 +1,319 @@ +"""Tests for the execution-view OpenTelemetry plugin (Workflow-rooted trace).""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import opentelemetry.context as otel_context +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + OperationEndInfo, + OperationStartInfo, +) +from opentelemetry import trace +from opentelemetry.context import Context +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 aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_workflow_span_id, + operation_id_to_span_id, +) +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, +) + + +START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) +END_TIME = datetime(2024, 1, 2, 3, 4, 6, tzinfo=UTC) +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" + + +@pytest.fixture(autouse=True) +def _reset_otel_context(): + """Reset the OTel thread-local context around each test. + + The plugin attaches spans via context.attach() without detaching, so state + would otherwise leak between tests running on the same thread. + """ + token = otel_context.attach(Context()) + try: + yield + finally: + otel_context.detach(token) + + +def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: + """Create an ExecutionOtelPlugin wired to an in-memory exporter.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + return plugin, exporter + + +def _invocation_start_info() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + +def _invocation_end_info( + status: InvocationStatus = InvocationStatus.SUCCEEDED, +) -> InvocationEndInfo: + return InvocationEndInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + status=status, + error=None, + ) + + +# --------------------------------------------------------------------------- +# derive_workflow_span_id +# --------------------------------------------------------------------------- +def test_derive_workflow_span_id_is_deterministic(): + assert derive_workflow_span_id(EXECUTION_ARN) == derive_workflow_span_id( + EXECUTION_ARN + ) + + +def test_derive_workflow_span_id_differs_by_arn(): + assert derive_workflow_span_id(EXECUTION_ARN) != derive_workflow_span_id( + EXECUTION_ARN + "-other" + ) + + +def test_derive_workflow_span_id_is_64_bit(): + span_id = derive_workflow_span_id(EXECUTION_ARN) + assert 0 < span_id < 2**64 + + +def test_derive_workflow_span_id_rejects_empty_arn(): + with pytest.raises(ValueError, match="execution ARN is required"): + derive_workflow_span_id("") + + +# --------------------------------------------------------------------------- +# Workflow + invocation span hierarchy +# --------------------------------------------------------------------------- +def test_workflow_span_is_root_and_invocation_is_its_child(): + plugin, exporter = _create_plugin() + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + + spans = {s.name: s for s in exporter.get_finished_spans()} + assert "Workflow" in spans + assert "invocation" in spans + + workflow = spans["Workflow"] + invocation = spans["invocation"] + + # Workflow is a root span (no parent) with the deterministic span ID. + assert workflow.parent is None + assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) + assert workflow.attributes["durable.execution.arn"] == EXECUTION_ARN + assert ( + workflow.attributes["durable.execution.status"] + == InvocationStatus.SUCCEEDED.value + ) + + # Invocation is parented under the Workflow span. + assert invocation.parent is not None + assert invocation.parent.span_id == workflow.context.span_id + + +def test_workflow_span_dropped_on_non_terminal_status(): + plugin, exporter = _create_plugin() + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + names = [s.name for s in exporter.get_finished_spans()] + # Invocation span is always ended/exported; the Workflow span is dropped + # (not ended) on a non-terminal status, so it must not be exported. + assert "invocation" in names + assert "Workflow" not in names + + +def test_operation_parented_under_workflow_and_linked_to_invocation(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + operation_id = "wait-1" + plugin.on_operation_start( + OperationStartInfo( + operation_id=operation_id, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="wait-for-signal", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="wait-for-signal", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + spans = {s.name: s for s in exporter.get_finished_spans()} + workflow = spans["Workflow"] + invocation = spans["invocation"] + operation = spans["wait-for-signal"] + + # Parented under the Workflow span (no parentId => Workflow fallback). + assert operation.parent is not None + assert operation.parent.span_id == workflow.context.span_id + + # Linked (not parented) to the invocation span. + linked_span_ids = {link.context.span_id for link in operation.links} + assert invocation.context.span_id in linked_span_ids + + +def test_cross_invocation_operation_end_uses_deterministic_span_id(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + # Operation end with no matching start (started in a prior invocation). + plugin.on_operation_end( + OperationEndInfo( + operation_id="step-earlier", + operation_type=OperationType.STEP, + sub_type=None, + name="earlier-step", + parent_id=None, + start_time=START_TIME, + is_replayed=True, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + matching = [s for s in exporter.get_finished_spans() if s.name == "earlier-step"] + # Exported exactly once, using the deterministic logical-operation span ID + # (no separate continuation span). + assert len(matching) == 1 + assert matching[0].context.span_id == operation_id_to_span_id( + EXECUTION_ARN, "step-earlier" + ) + + +# --------------------------------------------------------------------------- +# Default-provider mode: invocation span +# --------------------------------------------------------------------------- +def _create_default_mode_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: + """ExecutionOtelPlugin in default-provider mode wired to an in-memory exporter. + + Passing an explicit ``tracer_provider`` lets the test capture spans while + ``use_default_tracer_provider=True`` still selects the default-mode code path. + """ + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + use_default_tracer_provider=True, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + return plugin, exporter + + +def test_default_mode_creates_invocation_span(): + plugin, exporter = _create_default_mode_plugin() + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + + spans = {s.name: s for s in exporter.get_finished_spans()} + # The invocation span is now created even in default-provider mode. + assert "invocation" in spans + invocation = spans["invocation"] + assert invocation.attributes["durable.execution.arn"] == EXECUTION_ARN + assert invocation.attributes["durable.invocation.first"] is True + + +def test_default_mode_invocation_span_parented_to_ambient_span(): + plugin, exporter = _create_default_mode_plugin() + + # Simulate the ambient Lambda invocation span from the ADOT layer. + ambient = plugin._provider.get_tracer("ambient").start_span("lambda-invocation") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + finally: + otel_context.detach(token) + ambient.end() + + invocation = {s.name: s for s in exporter.get_finished_spans()}["invocation"] + assert invocation.parent is not None + assert invocation.parent.span_id == ambient.get_span_context().span_id + + +def test_open_operation_span_not_exported_at_invocation_end(): + """A suspended operation (started, not ended) must not be exported. + + on_invocation_end drops the reference without ending it; the + span is ended only when on_operation_end fires in a later invocation. + """ + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + + plugin.on_operation_start( + OperationStartInfo( + operation_id="wait-1", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="wait-for-signal", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + # No on_operation_end: the operation suspended. + plugin.on_invocation_end(_invocation_end_info(status=InvocationStatus.PENDING)) + + exported = {s.name for s in exporter.get_finished_spans()} + # The open operation span is NOT exported (never ended). + assert "wait-for-signal" not in exported 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 new file mode 100644 index 00000000..86166376 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -0,0 +1,247 @@ +"""In-process integration tests for ExecutionOtelPlugin. + +Drives the full plugin lifecycle against a real TracerProvider + +InMemorySpanExporter for the two deployment shapes: + +* Community collector layer: the plugin owns its provider + (``use_default_tracer_provider=False``); the Workflow span is the trace root. +* ADOT layer: the ADOT Lambda layer supplies the global provider and the ambient + Lambda invocation span (``use_default_tracer_provider=True``); the plugin's + Invocation span parents to that ambient span. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import opentelemetry.context as otel_context +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + OperationEndInfo, + OperationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, + UserFunctionStartInfo, +) +from opentelemetry import trace +from opentelemetry.context import Context +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 aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_workflow_span_id, + operation_id_to_span_id, +) +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig + + +START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) +END_TIME = datetime(2024, 1, 2, 3, 4, 6, tzinfo=UTC) +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" +OP_ID = "step-1" +OP_NAME = "fetch-user" + + +@pytest.fixture(autouse=True) +def _reset_otel_context(): + token = otel_context.attach(Context()) + try: + yield + finally: + otel_context.detach(token) + + +def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider, exporter + + +def _invocation_start() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + +def _invocation_end( + status: InvocationStatus = InvocationStatus.SUCCEEDED, +) -> InvocationEndInfo: + return InvocationEndInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + status=status, + error=None, + ) + + +def _run_step_lifecycle(plugin: ExecutionOtelPlugin) -> None: + """Drive a single completed STEP (operation + one attempt) through a plugin.""" + plugin.on_operation_start( + OperationStartInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_user_function_start( + UserFunctionStartInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + ) + plugin.on_user_function_end( + UserFunctionEndInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + 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, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + + +# --------------------------------------------------------------------------- +# Community collector layer (plugin owns the provider) +# --------------------------------------------------------------------------- +def test_community_layer_full_lifecycle_is_workflow_rooted(): + provider, exporter = _provider() + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + use_default_tracer_provider=False, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + + plugin.on_invocation_start(_invocation_start()) + _run_step_lifecycle(plugin) + plugin.on_invocation_end(_invocation_end()) + + finished = exporter.get_finished_spans() + spans = {s.name: s for s in finished} + workflow = spans["Workflow"] + invocation = spans["invocation"] + operation = spans[OP_NAME] + attempt = spans[f"{OP_NAME} attempt 1"] + + # Workflow is the trace root with the deterministic workflow span id. + assert workflow.parent is None + assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) + assert ( + workflow.attributes["durable.execution.status"] + == InvocationStatus.SUCCEEDED.value + ) + + # Invocation span is a child of the Workflow span. + assert invocation.parent is not None + assert invocation.parent.span_id == workflow.context.span_id + + # Operation span: deterministic id, parented under Workflow, linked to invocation. + assert operation.context.span_id == operation_id_to_span_id(EXECUTION_ARN, OP_ID) + assert operation.parent.span_id == workflow.context.span_id + assert invocation.context.span_id in { + link.context.span_id for link in operation.links + } + + # Attempt span is a child of the operation span. + assert attempt.parent.span_id == operation.context.span_id + + # The operation span is exported exactly once. + assert len([s for s in finished if s.name == OP_NAME]) == 1 + + +# --------------------------------------------------------------------------- +# ADOT layer (default provider; ambient invocation span) +# --------------------------------------------------------------------------- +def test_adot_layer_full_lifecycle_parents_to_ambient_span(): + provider, exporter = _provider() + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + use_default_tracer_provider=True, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + ) + + # Simulate the ambient Lambda invocation span the ADOT layer creates. + ambient = provider.get_tracer("adot").start_span("lambda-invocation") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + plugin.on_invocation_start(_invocation_start()) + _run_step_lifecycle(plugin) + plugin.on_invocation_end(_invocation_end()) + finally: + otel_context.detach(token) + ambient.end() + + finished = exporter.get_finished_spans() + spans = {s.name: s for s in finished} + invocation = spans["invocation"] + operation = spans[OP_NAME] + + # Invocation span parents to the ambient ADOT span and carries the first flag. + assert invocation.parent is not None + assert invocation.parent.span_id == ambient.get_span_context().span_id + assert invocation.attributes["durable.invocation.first"] is True + + # Operation span still uses the deterministic id and links to the ambient span. + assert operation.context.span_id == operation_id_to_span_id(EXECUTION_ARN, OP_ID) + assert ambient.get_span_context().span_id in { + link.context.span_id for link in operation.links + } + + # The operation span is exported exactly once. + assert len([s for s in finished if s.name == OP_NAME]) == 1 diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py similarity index 99% rename from packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py rename to packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index a9e46322..58d22a68 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -32,7 +32,7 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( operation_id_to_span_id, ) -from aws_durable_execution_sdk_python_otel.plugin import OtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -54,12 +54,12 @@ def _reset_otel_context(): otel_context.detach(token) -def _create_plugin() -> tuple[OtelPlugin, InMemorySpanExporter]: +def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: """Create a plugin wired to an in-memory span exporter.""" exporter = InMemorySpanExporter() trace_provider = TracerProvider() trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) - plugin = OtelPlugin( + plugin = InvocationOtelPlugin( trace_provider=trace_provider, context_extractor=lambda _: Context(), ) 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 new file mode 100644 index 00000000..129e3e97 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -0,0 +1,217 @@ +"""In-process integration tests for InvocationOtelPlugin. + +Drives the full plugin lifecycle against a real TracerProvider + +InMemorySpanExporter for the two deployment shapes: + +* Community collector layer: the caller supplies its own provider (configured to + export to a community OTel collector) via ``trace_provider=...``. +* ADOT layer: the ADOT Lambda layer configures the global provider and the plugin + picks it up from ``trace.get_tracer_provider()`` (no explicit provider passed). + +Both paths run identical plugin logic; each test asserts the +invocation -> operation -> attempt hierarchy is produced against whichever +provider is in use. +""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import opentelemetry.context as otel_context +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + InvocationStatus, + OperationStatus, + OperationSubType, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + OperationEndInfo, + OperationStartInfo, + UserFunctionEndInfo, + UserFunctionOutcome, + UserFunctionStartInfo, +) +from opentelemetry import trace +from opentelemetry.context import Context +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 + +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + operation_id_to_span_id, +) +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin + + +START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) +END_TIME = datetime(2024, 1, 2, 3, 4, 6, tzinfo=UTC) +EXECUTION_ARN = "arn:aws:lambda:us-west-2:123456789012:function:workflow:$LATEST" +OP_ID = "step-1" +OP_NAME = "fetch-user" + + +@pytest.fixture(autouse=True) +def _reset_otel_context(): + token = otel_context.attach(Context()) + try: + yield + finally: + otel_context.detach(token) + + +def _provider() -> tuple[TracerProvider, InMemorySpanExporter]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider, exporter + + +def _invocation_start() -> InvocationStartInfo: + return InvocationStartInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + +def _invocation_end( + status: InvocationStatus = InvocationStatus.SUCCEEDED, +) -> InvocationEndInfo: + return InvocationEndInfo( + request_id="request-1", + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + is_first_invocation=True, + status=status, + error=None, + ) + + +def _run_step_lifecycle(plugin: InvocationOtelPlugin) -> None: + """Drive a single completed STEP (operation + one attempt) through a plugin.""" + plugin.on_operation_start( + OperationStartInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_user_function_start( + UserFunctionStartInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + is_replay_children=False, + attempt=1, + ) + ) + plugin.on_user_function_end( + UserFunctionEndInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + 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, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=OP_ID, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=OP_NAME, + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + + +def _assert_hierarchy(exporter: InMemorySpanExporter) -> None: + spans = {s.name: s for s in exporter.get_finished_spans()} + invocation = spans["invocation"] + operation = spans[OP_NAME] + attempt = spans[f"{OP_NAME} attempt 1"] + + # Invocation span is a root (empty extracted context) and records status. + assert invocation.parent is None + assert invocation.kind is SpanKind.INTERNAL + assert invocation.attributes is not None + assert ( + invocation.attributes["durable.invocation.status"] + == InvocationStatus.SUCCEEDED.value + ) + + # Operation span: deterministic id, child of the invocation span. + assert operation.context.span_id == operation_id_to_span_id(EXECUTION_ARN, OP_ID) + assert operation.parent is not None + assert operation.parent.span_id == invocation.context.span_id + + # Attempt span is a child of the operation span. + assert attempt.parent is not None + assert attempt.parent.span_id == operation.context.span_id + + +# --------------------------------------------------------------------------- +# Community collector layer (caller-supplied provider) +# --------------------------------------------------------------------------- +def test_community_layer_full_lifecycle_uses_supplied_provider(): + provider, exporter = _provider() + plugin = InvocationOtelPlugin( + trace_provider=provider, + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + + plugin.on_invocation_start(_invocation_start()) + _run_step_lifecycle(plugin) + plugin.on_invocation_end(_invocation_end()) + + _assert_hierarchy(exporter) + + +# --------------------------------------------------------------------------- +# ADOT layer (global provider picked up from trace.get_tracer_provider) +# --------------------------------------------------------------------------- +def test_adot_layer_full_lifecycle_uses_global_provider(monkeypatch): + provider, exporter = _provider() + # Simulate the ADOT layer having configured the global TracerProvider. + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + + plugin = InvocationOtelPlugin( + context_extractor=lambda _: Context(), + enrich_logger=False, + ) + + plugin.on_invocation_start(_invocation_start()) + _run_step_lifecycle(plugin) + plugin.on_invocation_end(_invocation_end()) + + # Spans were produced against the (monkeypatched) global provider's exporter. + _assert_hierarchy(exporter) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index ba8788a3..39933338 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -22,7 +22,7 @@ OtelContextLogFilter, install_log_filter, ) -from aws_durable_execution_sdk_python_otel.plugin import OtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin START_TIME = datetime(2024, 1, 2, 3, 4, 5, tzinfo=UTC) @@ -31,12 +31,12 @@ def _create_plugin( enrich_logger: bool = True, -) -> tuple[OtelPlugin, InMemorySpanExporter]: +) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: """Create a plugin wired to an in-memory span exporter.""" exporter = InMemorySpanExporter() trace_provider = TracerProvider() trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) - plugin = OtelPlugin( + plugin = InvocationOtelPlugin( trace_provider=trace_provider, context_extractor=lambda _: Context(), enrich_logger=enrich_logger, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_provider.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_provider.py new file mode 100644 index 00000000..dd06a6f2 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_provider.py @@ -0,0 +1,114 @@ +"""Tests for the shared TracerProvider factory (create_tracer_provider).""" + +from __future__ import annotations + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, TraceIdRatioBased + +from aws_durable_execution_sdk_python_otel.otel_plugin_config import ( + OtelPluginConfig, + ExporterConfig, +) +from aws_durable_execution_sdk_python_otel.provider import ( + SAMPLING_RATIO_ENV, + _build_resource, + _build_sampler, + _resolve_endpoint, + create_tracer_provider, +) + + +def test_explicit_provider_is_used_and_not_owned(): + provider = TracerProvider() + result = create_tracer_provider(OtelPluginConfig(tracer_provider=provider)) + assert result.tracer_provider is provider + assert result.owns_provider is False + + +def test_use_default_provider_returns_global_and_not_owned(): + result = create_tracer_provider(OtelPluginConfig(use_default_tracer_provider=True)) + assert result.tracer_provider is trace.get_tracer_provider() + assert result.owns_provider is False + + +def test_default_use_global_flag_applies_when_unset(): + # InvocationOtelPlugin passes default_use_global=True so an unset config + # resolves to the global provider. + result = create_tracer_provider(OtelPluginConfig(), default_use_global=True) + assert result.tracer_provider is trace.get_tracer_provider() + assert result.owns_provider is False + + +def test_auto_configured_provider_is_owned_sdk_provider(): + result = create_tracer_provider(OtelPluginConfig()) + assert result.owns_provider is True + assert isinstance(result.tracer_provider, TracerProvider) + + +def test_explicit_provider_takes_precedence_over_use_default(): + provider = TracerProvider() + result = create_tracer_provider( + OtelPluginConfig(tracer_provider=provider, use_default_tracer_provider=True) + ) + assert result.tracer_provider is provider + assert result.owns_provider is False + + +# --------------------------------------------------------------------------- +# Sampler resolution +# --------------------------------------------------------------------------- +def test_sampler_defaults_to_always_on(monkeypatch): + monkeypatch.delenv(SAMPLING_RATIO_ENV, raising=False) + assert _build_sampler() is ALWAYS_ON + + +def test_sampler_uses_ratio_when_valid(monkeypatch): + monkeypatch.setenv(SAMPLING_RATIO_ENV, "0.25") + assert isinstance(_build_sampler(), TraceIdRatioBased) + + +@pytest.mark.parametrize("bad", ["not-a-number", "1.5", "-0.1"]) +def test_sampler_falls_back_to_always_on_for_invalid_ratio(monkeypatch, bad): + monkeypatch.setenv(SAMPLING_RATIO_ENV, bad) + assert _build_sampler() is ALWAYS_ON + + +# --------------------------------------------------------------------------- +# Resource + endpoint resolution +# --------------------------------------------------------------------------- +def test_resource_is_none_without_function_name(monkeypatch): + monkeypatch.delenv("AWS_LAMBDA_FUNCTION_NAME", raising=False) + assert _build_resource() is None + + +def test_resource_populates_lambda_attributes(monkeypatch): + monkeypatch.setenv("AWS_LAMBDA_FUNCTION_NAME", "my-fn") + monkeypatch.setenv("AWS_REGION", "us-west-2") + monkeypatch.setenv("AWS_LAMBDA_FUNCTION_VERSION", "7") + resource = _build_resource() + assert resource is not None + attrs = resource.attributes + assert attrs["service.name"] == "my-fn" + assert attrs["faas.name"] == "my-fn" + assert attrs["cloud.provider"] == "aws" + assert attrs["cloud.platform"] == "aws_lambda" + assert attrs["cloud.region"] == "us-west-2" + assert attrs["faas.version"] == "7" + + +def test_resolve_endpoint_appends_signal_path(monkeypatch): + monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) + config = OtelPluginConfig( + exporter_config=ExporterConfig(endpoint="http://collector:4318") + ) + assert _resolve_endpoint(config) == "http://collector:4318/v1/traces" + + +def test_resolve_endpoint_does_not_double_append(monkeypatch): + monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) + config = OtelPluginConfig( + exporter_config=ExporterConfig(endpoint="http://collector:4318/v1/traces") + ) + assert _resolve_endpoint(config) == "http://collector:4318/v1/traces" diff --git a/pyproject.toml b/pyproject.toml index 653e7836..86a1ace0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,8 @@ extra-dependencies = [ "pytest", "boto3-stubs[lambda]", "opentelemetry-sdk>=1.20.0,<=1.42.1", + "opentelemetry-instrumentation-botocore", + "opentelemetry-instrumentation-urllib3", ] [tool.hatch.envs.types.scripts]