diff --git a/pyproject.toml b/pyproject.toml index d6397a7b6..035a59122 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ openai-agents = ["openai-agents>=0.19.2,<0.20", "mcp>=1.9.4, <2"] google-adk = ["google-adk>=2.2.0,<3", "mcp>=1.24,<2"] langgraph = ["langgraph>=1.1.0"] langsmith = ["langsmith>=0.7.34,<0.9"] +datadog = ["ddtrace>=2.9,<4"] deepagents = [ "deepagents>=0.6.12,<0.7; python_version >= '3.11'", "langchain>=1.3.11,<2; python_version >= '3.11'", @@ -100,6 +101,7 @@ dev = [ "strands-agents>=1.39.0", "strands-agents-tools>=0.5.2", "mcp>=1.9.4,<2", + "ddtrace>=2.9,<4", ] [tool.poe.tasks] diff --git a/temporalio/contrib/datadog/README.md b/temporalio/contrib/datadog/README.md new file mode 100644 index 000000000..fc4d675d1 --- /dev/null +++ b/temporalio/contrib/datadog/README.md @@ -0,0 +1,132 @@ +# Datadog Tracing Interceptor for Temporal Python + +## Background + +### Why this is more complex than Go + +The Python Temporal SDK does not provide a common tracing interface to +implement. Instead, it exposes interceptor base classes that users extend to +inject behaviour at each operation boundary. This is idiomatic Python but +makes integration more involved than the Go equivalent. + +The more significant challenge is how the Python SDK executes workflow code. +Workflows run inside a sandboxed environment where the worker re-imports the +workflow module before every execution. The sandbox also enforces +determinism constraints, so tracing logic (which depends on `ddtrace`, a +module that schedules asyncio work during import) cannot live directly inside +it. + +### Difference from the upstream OpenTelemetry interceptor + +The upstream OpenTelemetry tracing interceptor works around the sandbox +limitation by emitting a zero-duration notification span whenever a workflow +execution occurs. Activities and other events attach to that span. While +functional, this approach does not produce actual workflow traces: the +`RunWorkflow` span has no duration and the trace does not survive a worker +restart. + +This implementation instead generates real workflow traces. The `RunWorkflow` +span starts when the first worker picks up the workflow and finishes only when +the workflow completes, matching Go's behaviour. Deterministic span IDs and +context propagation ensure the trace remains coherent even if a worker +restarts mid-execution. + +## Usage + +```python +import ddtrace +from temporalio.client import Client +from temporalio.contrib.datadog import DatadogTracingInterceptor + +# Inject dd.trace_id and dd.span_id into every log record so workflow and +# activity logs are correlated with their trace in Datadog Log Management. +ddtrace.patch(logging=True) + +interceptor = DatadogTracingInterceptor( + service_name="my-service", + extra_tags={"deployment.environment": "prod"}, +) +client = await Client.connect("localhost:7233", interceptors=[interceptor]) +``` + +**Important**: Passing the `interceptor` instance to the client is enough. +The worker will automatically pick up the interceptor from the client. + +## Deterministic span IDs + +The workflows' `RunWorkflow` spans may be long-running. If the worker +restarts and the workflow is replayed, the new execution recreates +the span with the **same span ID** so the trace remains coherent in APM. + +Span IDs for `RunWorkflow` (and any operation with an idempotency key) are +derived via FNV-1 64-bit hash of a key: + +``` +WorkflowInboundInterceptor:::: +``` + +This matches the Go SDK's algorithm byte-for-byte, so a workflow started by a +Go client and executed by a Python worker produces the same span ID. The +counter starts at 1 (reserved for RunWorkflow) and increments for each +subsequent handler span (HandleSignal, HandleUpdate) to give each a stable, +unique ID across worker restarts. + +## Replay safety + +Temporal replays workflow history on every new worker to rebuild execution +state. Without guards, replay would re-emit duplicate completed spans for +operations that already finished on the dead worker. + +Two mechanisms prevent duplicates: + +**Inbound handlers** (HandleSignal, HandleUpdate): suppressed during replay +via `temporalio.workflow.unsafe.is_replaying()`. A non-`None` idempotency key +signals that the span completed within a single workflow task and must not be +re-emitted. RunWorkflow is explicitly exempt — its span is in-flight and was +never sent by the dead worker, so the new worker recreates it. + +**Outbound operations** (StartActivity, StartLocalActivity, StartChildWorkflow, +SignalChildWorkflow, SignalExternalWorkflow): suppressed during replay by an +early `is_replaying()` check at the top of each outbound interceptor method. +The Temporal SDK matches these commands against history and returns cached +results, but the interceptor code runs first — without the guard, a fresh span +would be emitted for every replayed command. + +Queries and update validators are never suppressed: queries are not in +history (they run on demand), and validators do not execute during replay. + +## Workflow sandbox + +Temporal runs workflow code inside a restricted import sandbox. Importing +`ddtrace` from within that import path triggers an asyncio-loop conflict +during ddtrace init and a `builtins.open` restriction from pytest's assertion +rewriter. + +The interceptor works around this with an extern-function bridge: + +1. `DatadogTracingInterceptor` (host side) registers functions under + `unsafe_extern_functions` before the sandbox starts. +2. `DatadogTracingWorkflowInboundInterceptor` (sandbox side) retrieves those + functions via `temporalio.workflow.extern_functions()` at init time and + holds them as instance attributes. +3. All ddtrace calls (`start_span`, `finish_span`, baggage, annotation) go + through these externs, so the sandbox never imports ddtrace directly. + +Two `contextvars.ContextVar` values live on the host module and are accessed +by the sandbox exclusively through externs: + +- `_active_workflow_span` — the live `RunWorkflow` ddtrace span for the + current execution. Used by outbound operations to parent their spans to + RunWorkflow when no propagated header is available. +- `_trace_disconnected` — set by `disconnect_trace_span_from_workflow_context` + to suppress trace propagation into the next ContinueAsNew run. + +## ContinueAsNew + +`_WorkflowOutboundInterceptor.continue_as_new` injects the current RunWorkflow +span context into the ContinueAsNew headers so the next run's RunWorkflow span +is a child of the current one, forming a continuous trace across runs. + +Call `disconnect_trace_span_from_workflow_context()` before +`workflow.continue_as_new()` to start a fresh root trace for the next run +instead. diff --git a/temporalio/contrib/datadog/__init__.py b/temporalio/contrib/datadog/__init__.py new file mode 100644 index 000000000..57233ca07 --- /dev/null +++ b/temporalio/contrib/datadog/__init__.py @@ -0,0 +1,37 @@ +"""Datadog tracing integration for the Temporal Python SDK. + +This package provides a Datadog (``ddtrace``) tracing interceptor for the +Temporal Python SDK. + +Usage:: + + from ddtrace import patch + from temporalio.client import Client + from temporalio.contrib.datadog import DatadogTracingInterceptor + + patch(logging=True) # opt in to dd.trace_id log injection + + interceptor = DatadogTracingInterceptor( + service_name="my-service", + extra_tags={"deployment.environment": "prod"}, + ) + client = await Client.connect("localhost:7233", interceptors=[interceptor]) +""" + +from temporalio.contrib.datadog._interceptor import DatadogTracingInterceptor +from temporalio.contrib.datadog._workflow_interceptor import ( + disconnect_trace_span_from_workflow_context, + span_from_workflow_context, +) +from temporalio.contrib.datadog._wrapped_tracer import ( + FinishContext, + FinishResult, +) + +__all__ = [ + "DatadogTracingInterceptor", + "FinishContext", + "FinishResult", + "disconnect_trace_span_from_workflow_context", + "span_from_workflow_context", +] diff --git a/temporalio/contrib/datadog/_activity_interceptor.py b/temporalio/contrib/datadog/_activity_interceptor.py new file mode 100644 index 000000000..1327fc9aa --- /dev/null +++ b/temporalio/contrib/datadog/_activity_interceptor.py @@ -0,0 +1,67 @@ +"""Datadog tracing interceptor for Temporal activity inbound calls.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import temporalio.activity +import temporalio.worker +from temporalio.contrib.datadog._constants import ( + OperationNames, + SpanAttributes, +) +from temporalio.contrib.datadog._id_generator import gen_span_id +from temporalio.contrib.datadog._span_runner import _SpanRunner + +if TYPE_CHECKING: + from temporalio.contrib.datadog._interceptor import DatadogTracingInterceptor + + +class _ActivityInboundInterceptor( # type: ignore[reportUnsafeMultipleInheritance] + _SpanRunner, temporalio.worker.ActivityInboundInterceptor +): + def __init__( + self, + next: temporalio.worker.ActivityInboundInterceptor, + root: DatadogTracingInterceptor, + ) -> None: + temporalio.worker.ActivityInboundInterceptor.__init__(self, next) + _SpanRunner.__init__(self, root) + + async def execute_activity( + self, input: temporalio.worker.ExecuteActivityInput + ) -> Any: + return await self.run( + self._get_span(input), + OperationNames.RUN_ACTIVITY, + super().execute_activity(input), + ) + + def _get_span(self, input: temporalio.worker.ExecuteActivityInput) -> Any: + info = temporalio.activity.info() + return self.root.tracer.start_span( + operation_name=OperationNames.RUN_ACTIVITY, + parent_ctx=self.root.propagator.extract_headers(input.headers), + resource_name=info.activity_type, + activate=True, + span_id=gen_span_id( + f"{info.workflow_run_id}:{info.activity_id}:{info.attempt}" + ), + attributes=self._get_activity_attributes(info), + parent_from_header=True, + ) + + @staticmethod + def _get_activity_attributes(info: temporalio.activity.Info) -> dict[str, Any]: + attributes: dict[str, Any] = { + SpanAttributes.ACTIVITY_ID: info.activity_id, + SpanAttributes.ACTIVITY_TYPE: info.activity_type, + SpanAttributes.ATTEMPT: info.attempt, + } + if info.workflow_id: + attributes[SpanAttributes.WORKFLOW_ID] = info.workflow_id + if info.workflow_run_id: + attributes[SpanAttributes.RUN_ID] = info.workflow_run_id + if info.workflow_namespace: + attributes[SpanAttributes.NAMESPACE] = info.workflow_namespace + return attributes diff --git a/temporalio/contrib/datadog/_client_interceptor.py b/temporalio/contrib/datadog/_client_interceptor.py new file mode 100644 index 000000000..4a4bf10ef --- /dev/null +++ b/temporalio/contrib/datadog/_client_interceptor.py @@ -0,0 +1,184 @@ +"""Datadog tracing interceptor for Temporal client outbound calls.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Any + +import temporalio.client +from temporalio.contrib.datadog._constants import ( + COMMON_ATTRIBUTE_MAP, + OperationNames, + SpanAttributes, +) +from temporalio.contrib.datadog._span_runner import _SpanRunner + +if TYPE_CHECKING: + from temporalio.contrib.datadog._interceptor import DatadogTracingInterceptor + + +class _ClientOutboundInterceptor( # type: ignore[reportUnsafeMultipleInheritance] + _SpanRunner, temporalio.client.OutboundInterceptor +): + def __init__( + self, + next: temporalio.client.OutboundInterceptor, + root: DatadogTracingInterceptor, + ) -> None: + temporalio.client.OutboundInterceptor.__init__(self, next) + _SpanRunner.__init__(self, root) + + async def start_workflow( + self, input: temporalio.client.StartWorkflowInput + ) -> temporalio.client.WorkflowHandle[Any, Any]: + operation_name = ( + OperationNames.SIGNAL_WITH_START_WORKFLOW + if input.start_signal + else OperationNames.START_WORKFLOW + ) + return await self._start_operation( + operation_name, + input, + input.workflow, + self._get_workflow_attributes(input), + super().start_workflow, + ) + + async def signal_workflow( + self, input: temporalio.client.SignalWorkflowInput + ) -> None: + if self.root.workflow_tracing_config.disable_signal_tracing: + return await super().signal_workflow(input) + return await self._start_operation( + OperationNames.SIGNAL_WORKFLOW, + input, + input.signal, + self._get_workflow_attributes(input), + super().signal_workflow, + ) + + async def query_workflow(self, input: temporalio.client.QueryWorkflowInput) -> Any: + if self.root.workflow_tracing_config.disable_query_tracing: + return await super().query_workflow(input) + return await self._start_operation( + OperationNames.QUERY_WORKFLOW, + input, + input.query, + self._get_workflow_attributes(input), + super().query_workflow, + ) + + async def create_schedule( + self, input: temporalio.client.CreateScheduleInput + ) -> temporalio.client.ScheduleHandle: + span = self._get_span(OperationNames.CREATE_SCHEDULE, input.id) + return await self.run( + span, OperationNames.CREATE_SCHEDULE, super().create_schedule(input) + ) + + async def start_workflow_update( + self, input: temporalio.client.StartWorkflowUpdateInput + ) -> temporalio.client.WorkflowUpdateHandle[Any]: + if self.root.workflow_tracing_config.disable_update_tracing: + return await super().start_workflow_update(input) + return await self._start_operation( + OperationNames.UPDATE_WORKFLOW, + input, + input.update, + self._get_workflow_attributes(input), + super().start_workflow_update, + ) + + async def start_update_with_start_workflow( + self, input: temporalio.client.StartWorkflowUpdateWithStartInput + ) -> temporalio.client.WorkflowUpdateHandle[Any]: + if self.root.workflow_tracing_config.disable_update_tracing: + # Update tracing is disabled, but this call also starts a new workflow. + # Propagate the currently active trace into the workflow start headers so + # RunWorkflow is not an unparented root when workflow tracing is enabled. + active_ctx = self.root.tracer.tracer.context_provider.active() + input.start_workflow_input.headers = self.root.propagator.inject_headers( + input.start_workflow_input.headers, active_ctx + ) + return await super().start_update_with_start_workflow(input) + + operation_name = OperationNames.UPDATE_WITH_START_WORKFLOW + span = self._get_span( + operation_name, + input.start_workflow_input.workflow, + self._get_workflow_attributes(input.start_workflow_input), + ) + + input.start_workflow_input.headers = self.root.propagator.inject_headers( + input.start_workflow_input.headers, span.context + ) + input.update_workflow_input.headers = self.root.propagator.inject_headers( + input.update_workflow_input.headers, span.context + ) + + return await self.run( + span, + operation_name, + super().start_update_with_start_workflow(input), + ) + + async def start_activity( + self, input: temporalio.client.StartActivityInput + ) -> temporalio.client.ActivityHandle[Any]: + return await self._start_operation( + OperationNames.START_ACTIVITY, + input, + input.activity_type, + self._get_activity_attributes(input), + super().start_activity, + ) + + async def _start_operation( + self, + operation_name: str, + input: Any, + resource_name: str, + attributes: dict[str, Any], + awaitable: Callable[[Any], Awaitable[Any]], + ) -> Any: + span = self._get_span(operation_name, resource_name, attributes) + input.headers = self.root.propagator.inject_headers(input.headers, span.context) + return await self.run(span, operation_name, awaitable(input)) + + def _get_span( + self, + operation_name: str, + resource_name: str, + attributes: dict[str, Any] | None = None, + ) -> Any: + # Use the currently active ddtrace span as parent if one exists + parent_ctx = self.root.tracer.tracer.context_provider.active() + return self.root.tracer.start_span( + operation_name=operation_name, + parent_ctx=parent_ctx, + resource_name=resource_name, + activate=True, + attributes=attributes, + parent_from_header=False, + ) + + @classmethod + def _get_workflow_attributes(cls, input: Any) -> dict[str, Any]: + attributes: dict[str, Any] = {SpanAttributes.WORKFLOW_ID: input.id} + for field, span_key in COMMON_ATTRIBUTE_MAP: + if val := getattr(input, field, None): + attributes[span_key] = val + if getattr(input, "workflow", None): + attributes[SpanAttributes.WORKFLOW_TYPE] = input.workflow + if getattr(input, "update", None): + attributes[SpanAttributes.UPDATE_NAME] = input.update + if getattr(input, "update_id", None): + attributes[SpanAttributes.UPDATE_ID] = input.update_id + return attributes + + @classmethod + def _get_activity_attributes(cls, input: Any) -> dict[str, Any]: + return { + SpanAttributes.ACTIVITY_ID: input.id, + SpanAttributes.ACTIVITY_TYPE: input.activity_type, + } diff --git a/temporalio/contrib/datadog/_constants.py b/temporalio/contrib/datadog/_constants.py new file mode 100644 index 000000000..62ce33f7c --- /dev/null +++ b/temporalio/contrib/datadog/_constants.py @@ -0,0 +1,89 @@ +"""Package-wide constants for the Datadog tracing interceptor.""" + +from collections.abc import Mapping +from typing import TypeAlias + +import temporalio.api.common.v1 + +Carrier: TypeAlias = dict[str, str] +StringHeader: TypeAlias = Mapping[str, str] +TemporalHeader: TypeAlias = Mapping[str, temporalio.api.common.v1.Payload] + +BAGGAGE_ITEM_SERVICE = "servicename" +CONTINUE_AS_NEW_TAG = "temporal.continued_as_new" +DEFAULT_HEADER_KEY = "dd_trace_span" +TEMPORAL_TAG_PREFIX = "temporal." +_MANUAL_KEEP_TAG = "manual.keep" + + +class SpanAttributes: + ACTIVITY_ID = "ActivityID" + ACTIVITY_TYPE = "ActivityType" + ATTEMPT = "Attempt" + CHILD_WORKFLOW_ID = "ChildWorkflowID" + CHILD_WORKFLOW_TYPE = "ChildWorkflowType" + EXTERNAL_WORKFLOW_ID = "ExternalWorkflowID" + LOCAL = "Local" + NAMESPACE = "Namespace" + NEXUS_OPERATION = "NexusOperation" + NEXUS_SERVICE = "NexusService" + QUERY_TYPE = "QueryType" + RUN_ID = "RunID" + SIGNAL_NAME = "SignalName" + UPDATE_ID = "UpdateID" + UPDATE_NAME = "UpdateName" + WORKFLOW_ID = "WorkflowID" + WORKFLOW_TYPE = "WorkflowType" + + +COMMON_ATTRIBUTE_MAP: tuple[tuple[str, str], ...] = ( + ("signal", SpanAttributes.SIGNAL_NAME), + ("query", SpanAttributes.QUERY_TYPE), + ("activity", SpanAttributes.ACTIVITY_TYPE), + ("child_workflow_id", SpanAttributes.CHILD_WORKFLOW_ID), + ("workflow_id", SpanAttributes.EXTERNAL_WORKFLOW_ID), + ("service", SpanAttributes.NEXUS_SERVICE), + ("operation_name", SpanAttributes.NEXUS_OPERATION), +) + + +class OperationNames: + CREATE_SCHEDULE = "CreateSchedule" + HANDLE_QUERY = "HandleQuery" + HANDLE_SIGNAL = "HandleSignal" + HANDLE_UPDATE = "HandleUpdate" + QUERY_WORKFLOW = "QueryWorkflow" + RUN_ACTIVITY = "RunActivity" + RUN_WORKFLOW = "RunWorkflow" + SIGNAL_CHILD_WORKFLOW = "SignalChildWorkflow" + SIGNAL_EXTERNAL_WORKFLOW = "SignalExternalWorkflow" + SIGNAL_WITH_START_WORKFLOW = "SignalWithStartWorkflow" + SIGNAL_WORKFLOW = "SignalWorkflow" + START_ACTIVITY = "StartActivity" + START_CHILD_WORKFLOW = "StartChildWorkflow" + START_NEXUS_OPERATION = "StartNexusOperation" + RUN_NEXUS_OPERATION_START_HANDLER = "RunStartNexusOperationHandler" + RUN_NEXUS_OPERATION_CANCEL_HANDLER = "RunCancelNexusOperationHandler" + UPDATE_WITH_START_WORKFLOW = "UpdateWithStartWorkflow" + UPDATE_WORKFLOW = "UpdateWorkflow" + START_WORKFLOW = "StartWorkflow" + VALIDATE_UPDATE = "ValidateUpdate" + + +# Temporal entry-point operations whose spans are assigned USER_KEEP when they +# have no in-process parent. +# StartActivity is included because clients can start standalone activities +# without a workflow parent. +_MANUAL_KEEP_OPS: frozenset[str] = frozenset( + { + OperationNames.RUN_WORKFLOW, + OperationNames.START_WORKFLOW, + OperationNames.SIGNAL_WITH_START_WORKFLOW, + OperationNames.SIGNAL_WORKFLOW, + OperationNames.QUERY_WORKFLOW, + OperationNames.UPDATE_WORKFLOW, + OperationNames.UPDATE_WITH_START_WORKFLOW, + OperationNames.CREATE_SCHEDULE, + OperationNames.START_ACTIVITY, + } +) diff --git a/temporalio/contrib/datadog/_id_generator.py b/temporalio/contrib/datadog/_id_generator.py new file mode 100644 index 000000000..629477d63 --- /dev/null +++ b/temporalio/contrib/datadog/_id_generator.py @@ -0,0 +1,40 @@ +"""Deterministic span ID generation for Temporal Datadog tracing.""" + +_FNV_OFFSET_64 = 0xCBF29CE484222325 +_FNV_PRIME_64 = 0x100000001B3 +_FNV_MASK_64 = 0xFFFFFFFFFFFFFFFF + + +def gen_trace_id(key: str) -> int: + """Compute a deterministic 64-bit trace ID for a root RunWorkflow span. + + Used when no Datadog trace header is present (uninstrumented client), so + that the trace ID is stable across worker restarts for the same run. + + Uses the same FNV-1 algorithm as gen_span_id but prefixes the input with + ``trace:`` to keep trace-ID inputs distinct from span-ID inputs. + """ + return gen_span_id(f"trace:{key}") + + +def gen_span_id(key: str) -> int: + """Compute a 64-bit FNV-1 hash of a UTF-8 encoded string. + + Used to derive deterministic Datadog span IDs from the composite + idempotency keys built by the tracing interceptors. + + Matches the byte-for-byte output of Go's ``hash/fnv.New64()`` (which is + FNV-1, not FNV-1a), so Go and Python workers hashing the same + idempotency key produce the same span ID. + + Args: + key: The string to hash. + + Returns: + A 64-bit unsigned integer suitable for use as a Datadog span ID. + """ + h = _FNV_OFFSET_64 + for byte in key.encode("utf-8"): + h = (h * _FNV_PRIME_64) & _FNV_MASK_64 + h ^= byte + return h diff --git a/temporalio/contrib/datadog/_interceptor.py b/temporalio/contrib/datadog/_interceptor.py new file mode 100644 index 000000000..bf48917f0 --- /dev/null +++ b/temporalio/contrib/datadog/_interceptor.py @@ -0,0 +1,159 @@ +"""Datadog tracing interceptor for Temporal.""" + +from collections.abc import Callable, Mapping +from typing import Any + +import temporalio.client +import temporalio.converter +import temporalio.worker +import temporalio.workflow +from temporalio.contrib.datadog._activity_interceptor import _ActivityInboundInterceptor +from temporalio.contrib.datadog._client_interceptor import _ClientOutboundInterceptor +from temporalio.contrib.datadog._constants import ( + CONTINUE_AS_NEW_TAG, + DEFAULT_HEADER_KEY, + OperationNames, +) +from temporalio.contrib.datadog._id_generator import gen_span_id, gen_trace_id +from temporalio.contrib.datadog._nexus_interceptor import ( + _NexusOperationInboundInterceptor, +) +from temporalio.contrib.datadog._propagator import _Propagator +from temporalio.contrib.datadog._span_annotator import _SpanAnnotator +from temporalio.contrib.datadog._workflow_interceptor import ( + DatadogTracingWorkflowInboundInterceptor, + WorkflowTracingConfig, + _active_workflow_span, + _trace_disconnected, +) +from temporalio.contrib.datadog._wrapped_tracer import ( + FinishContext, + FinishResult, + WrappedTracer, +) + + +class DatadogTracingInterceptor( + temporalio.client.Interceptor, temporalio.worker.Interceptor +): + def __init__( # type: ignore[reportMissingSuperCall] + self, + tracer: Any | None = None, + *, + service_name: str | None = None, + header_key: str = DEFAULT_HEADER_KEY, + extra_tags: Mapping[str, str] | None = None, + on_span_finish: Callable[[FinishContext], FinishResult | None] | None = None, + workflow_tracing_config: WorkflowTracingConfig = WorkflowTracingConfig.default_config(), + allow_invalid_parent_spans: bool = False, + ) -> None: + self.workflow_tracing_config = workflow_tracing_config + + self.propagator = _Propagator( + header_key=header_key, + service_name=service_name, + payload_converter=temporalio.converter.PayloadConverter.default, + allow_invalid_parent_spans=allow_invalid_parent_spans, + ) + + self.tracer = WrappedTracer( + service_name=service_name, + tracer=tracer, + on_span_finish=on_span_finish, + annotator=_SpanAnnotator(service_name=service_name, extra_tags=extra_tags), + propagator=self.propagator, + ) + + def intercept_client( + self, next: temporalio.client.OutboundInterceptor + ) -> temporalio.client.OutboundInterceptor: + return _ClientOutboundInterceptor(next, self) + + def intercept_activity( + self, next: temporalio.worker.ActivityInboundInterceptor + ) -> temporalio.worker.ActivityInboundInterceptor: + return _ActivityInboundInterceptor(next, self) + + def intercept_nexus_operation( + self, next: temporalio.worker.NexusOperationInboundInterceptor + ) -> temporalio.worker.NexusOperationInboundInterceptor: + return _NexusOperationInboundInterceptor(next, self) + + def workflow_interceptor_class( + self, input: temporalio.worker.WorkflowInterceptorClassInput + ) -> type[DatadogTracingWorkflowInboundInterceptor]: + input.unsafe_extern_functions["__temporal_datadog_start_sandboxed_span"] = ( + self._start_sandboxed_span + ) + input.unsafe_extern_functions["__temporal_datadog_finish_sandboxed_span"] = ( + self._finish_sandboxed_span + ) + input.unsafe_extern_functions[ + "__temporal_datadog_configure_workflow_tracing" + ] = self._configure_workflow_tracing + # NOTE: fork-specific workaround; reconsider if this integration is ever proposed upstream. + # These two externs let workflow code cross the sandbox boundary to read/write + # ContextVars that live in the host module. Without them, the sandbox's own + # copy of workflow_interceptor has fresh ContextVars that are never set. + input.unsafe_extern_functions["__temporal_datadog_get_workflow_span"] = ( + _active_workflow_span.get + ) + input.unsafe_extern_functions["__temporal_datadog_set_trace_disconnected"] = ( + lambda: _trace_disconnected.set(True) + ) + return DatadogTracingWorkflowInboundInterceptor + + def _configure_workflow_tracing(self): + return self.propagator, self.workflow_tracing_config + + def _start_sandboxed_span( + self, + operation_name: str, + resource_name: str, + attributes: dict[str, Any] | None, + parent_ctx: Any | None, + idempotency_key: str | None, + start_time: int | None = None, + ) -> Any: + # No DD header (uninstrumented client): pass a deterministic trace_id to keep + # the RunWorkflow trace consistent if the worker restarts mid-run. + det_trace_id = ( + gen_trace_id(idempotency_key) + if operation_name == OperationNames.RUN_WORKFLOW + and parent_ctx is None + and idempotency_key is not None + else None + ) + span = self.tracer.start_span( + operation_name=operation_name, + parent_ctx=parent_ctx, + resource_name=resource_name, + activate=False, + start_time=start_time, + span_id=gen_span_id(idempotency_key) + if idempotency_key is not None + else None, + attributes=attributes, + parent_from_header=True, + trace_id=det_trace_id, + ) + # NOTE: fork-specific workaround; reconsider if this integration is ever proposed upstream. + # Expose the RunWorkflow span to the host-side ContextVar so that + # span_from_workflow_context() can return it via the extern. + if operation_name == OperationNames.RUN_WORKFLOW: + _active_workflow_span.set(span) + return span + + def _finish_sandboxed_span( + self, + operation_name: str, + span: Any | None, + operation_exc: BaseException | None, + ) -> None: + if span is None: + return + + if isinstance(operation_exc, temporalio.workflow.ContinueAsNewError): + span.set_tag(CONTINUE_AS_NEW_TAG, True) + + self.tracer.finish_span(span, operation_name, operation_exc) diff --git a/temporalio/contrib/datadog/_nexus_interceptor.py b/temporalio/contrib/datadog/_nexus_interceptor.py new file mode 100644 index 000000000..20dcc59af --- /dev/null +++ b/temporalio/contrib/datadog/_nexus_interceptor.py @@ -0,0 +1,64 @@ +"""Datadog tracing interceptor for Temporal Nexus operation inbound calls.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import nexusrpc.handler + +import temporalio.worker +from temporalio.contrib.datadog._constants import OperationNames, SpanAttributes +from temporalio.contrib.datadog._span_runner import _SpanRunner + +if TYPE_CHECKING: + from temporalio.contrib.datadog._interceptor import DatadogTracingInterceptor + + +class _NexusOperationInboundInterceptor( # type: ignore[reportUnsafeMultipleInheritance] + _SpanRunner, temporalio.worker.NexusOperationInboundInterceptor +): + def __init__( + self, + next: temporalio.worker.NexusOperationInboundInterceptor, + root: DatadogTracingInterceptor, + ) -> None: + temporalio.worker.NexusOperationInboundInterceptor.__init__(self, next) + _SpanRunner.__init__(self, root) + + async def execute_nexus_operation_start( + self, input: temporalio.worker.ExecuteNexusOperationStartInput + ) -> ( + nexusrpc.handler.StartOperationResultSync[Any] + | nexusrpc.handler.StartOperationResultAsync + ): + return await self.run( + self._get_span(input, OperationNames.RUN_NEXUS_OPERATION_START_HANDLER), + OperationNames.RUN_NEXUS_OPERATION_START_HANDLER, + super().execute_nexus_operation_start(input), + ) + + async def execute_nexus_operation_cancel( + self, input: temporalio.worker.ExecuteNexusOperationCancelInput + ) -> None: + return await self.run( + self._get_span(input, OperationNames.RUN_NEXUS_OPERATION_CANCEL_HANDLER), + OperationNames.RUN_NEXUS_OPERATION_CANCEL_HANDLER, + super().execute_nexus_operation_cancel(input), + ) + + def _get_span(self, input: Any, operation_name: str) -> Any: + return self.root.tracer.start_span( + operation_name=operation_name, + parent_ctx=self.root.propagator.extract(input.ctx.headers), + resource_name=f"{input.ctx.service}/{input.ctx.operation}", + activate=True, + attributes=self._get_nexus_attributes(input.ctx), + parent_from_header=True, + ) + + @staticmethod + def _get_nexus_attributes(nexus_ctx: Any) -> dict[str, Any]: + return { + SpanAttributes.NEXUS_SERVICE: nexus_ctx.service, + SpanAttributes.NEXUS_OPERATION: nexus_ctx.operation, + } diff --git a/temporalio/contrib/datadog/_propagator.py b/temporalio/contrib/datadog/_propagator.py new file mode 100644 index 000000000..d3074a300 --- /dev/null +++ b/temporalio/contrib/datadog/_propagator.py @@ -0,0 +1,113 @@ +"""Datadog propagation wrapper for Temporal headers.""" + +from typing import Any, cast + +import temporalio.api.common.v1 +import temporalio.converter +from temporalio.contrib.datadog._constants import ( + BAGGAGE_ITEM_SERVICE, + Carrier, + StringHeader, + TemporalHeader, +) + + +class _Propagator: # type: ignore[reportUnusedClass] + """Wraps HTTPPropagator with Temporal header encode/decode logic. + + The ``ddtrace`` import is deferred to ``__init__`` so sandbox re-importing + this module does not import ``ddtrace``. + """ + + def __init__( + self, + *, + header_key: str, + service_name: str | None, + payload_converter: temporalio.converter.PayloadConverter, + allow_invalid_parent_spans: bool = False, + ) -> None: + from ddtrace.propagation.http import HTTPPropagator + + self._propagator = HTTPPropagator + self.header_key = header_key + self.service_name = service_name + self._payload_converter = payload_converter + self.allow_invalid_parent_spans = allow_invalid_parent_spans + + @staticmethod + def get_baggage(ctx: Any) -> str | None: + if ctx is None: + return None + + getter = getattr(ctx, "get_baggage_item", None) + if callable(getter): + return cast(str | None, getter(BAGGAGE_ITEM_SERVICE)) + + return None + + def set_baggage(self, ctx: Any) -> None: + if self.service_name is None: + return + setter = getattr(ctx, "set_baggage_item", None) + if callable(setter): + setter(BAGGAGE_ITEM_SERVICE, self.service_name) + + def inject(self, context: Any) -> Carrier: + carrier: Carrier = {} + if context is None: + return carrier + self._propagator.inject(context, carrier) + return carrier + + def extract(self, header: StringHeader | None) -> Any: + if header is None: + return None + + try: + ctx = self._propagator.extract(header) + except Exception: + if self.allow_invalid_parent_spans: + return None + raise + + if ctx is None or getattr(ctx, "trace_id", None) is None: + return None + + return ctx + + def _carrier_to_payload(self, carrier: Carrier) -> temporalio.api.common.v1.Payload: + return self._payload_converter.to_payloads([carrier])[0] + + def _payload_to_carrier( + self, payload: temporalio.api.common.v1.Payload + ) -> Carrier | None: + decoded = self._payload_converter.from_payloads([payload])[0] + if not isinstance(decoded, dict): + return None + return {str(k): str(v) for k, v in decoded.items()} + + def inject_headers( + self, + headers: TemporalHeader, + context: Any, + ) -> TemporalHeader: + if context is None: + return headers + + self.set_baggage(context) + carrier = self.inject(context) + + return {**headers, self.header_key: self._carrier_to_payload(carrier)} + + def extract_headers(self, headers: TemporalHeader) -> Any: + payload = headers.get(self.header_key) + if payload is None: + return None + try: + carrier = self._payload_to_carrier(payload) + except Exception: + if self.allow_invalid_parent_spans: + return None + raise + return self.extract(carrier) diff --git a/temporalio/contrib/datadog/_span_annotator.py b/temporalio/contrib/datadog/_span_annotator.py new file mode 100644 index 000000000..42e03dc5c --- /dev/null +++ b/temporalio/contrib/datadog/_span_annotator.py @@ -0,0 +1,93 @@ +"""Span annotation logic for the Datadog tracing interceptor.""" + +from collections.abc import Mapping +from typing import Any + +from temporalio.contrib.datadog._constants import ( + _MANUAL_KEEP_OPS, + _MANUAL_KEEP_TAG, + TEMPORAL_TAG_PREFIX, + OperationNames, +) + + +class _SpanAnnotator: # type: ignore[reportUnusedClass] + _PEER_SERVICE_TAG = "peer.service" + _SPAN_KIND_TAG = "span.kind" + _PRODUCER = "producer" + _CONSUMER = "consumer" + + _SPAN_KIND: dict[str, str] = { + OperationNames.START_ACTIVITY: _PRODUCER, + OperationNames.RUN_ACTIVITY: _CONSUMER, + OperationNames.START_CHILD_WORKFLOW: _PRODUCER, + OperationNames.START_WORKFLOW: _PRODUCER, + OperationNames.SIGNAL_WITH_START_WORKFLOW: _PRODUCER, + OperationNames.RUN_WORKFLOW: _CONSUMER, + OperationNames.SIGNAL_WORKFLOW: _PRODUCER, + OperationNames.SIGNAL_CHILD_WORKFLOW: _PRODUCER, + OperationNames.SIGNAL_EXTERNAL_WORKFLOW: _PRODUCER, + OperationNames.HANDLE_SIGNAL: _CONSUMER, + OperationNames.QUERY_WORKFLOW: _PRODUCER, + OperationNames.HANDLE_QUERY: _CONSUMER, + OperationNames.UPDATE_WORKFLOW: _PRODUCER, + OperationNames.UPDATE_WITH_START_WORKFLOW: _PRODUCER, + OperationNames.VALIDATE_UPDATE: _CONSUMER, + OperationNames.HANDLE_UPDATE: _CONSUMER, + OperationNames.CREATE_SCHEDULE: _PRODUCER, + OperationNames.START_NEXUS_OPERATION: _PRODUCER, + OperationNames.RUN_NEXUS_OPERATION_START_HANDLER: _CONSUMER, + OperationNames.RUN_NEXUS_OPERATION_CANCEL_HANDLER: _CONSUMER, + } + + def __init__( + self, + *, + service_name: str | None = None, + extra_tags: Mapping[str, str] | None = None, + ) -> None: + self.service_name = service_name + self.extra_tags: Mapping[str, str] = extra_tags or {} + + @classmethod + def _normalize_key(cls, key: str) -> str: + if key.startswith(TEMPORAL_TAG_PREFIX): + return key + if key.lower().startswith("temporal"): + return TEMPORAL_TAG_PREFIX + key[len("temporal") :].lstrip(".") + return TEMPORAL_TAG_PREFIX + key + + def annotate( + self, + span: Any, + operation: str, + attributes: Mapping[str, Any] | None, + parent_service_name: str | None, + force_keep: bool = False, + ) -> None: + # User-defined global custom tags + for key, value in self.extra_tags.items(): + span.set_tag(key, value) + + # Attributes from the operation + if attributes: + for key, value in attributes.items(): + span.set_tag(self._normalize_key(key), value) + + # Force-keep entry-point operations that have no local parent. + # Two parent shapes qualify: no parent (nil/None — scheduled or standalone + # execution) and parents extracted from Temporal task headers (cross-process). + # Parents from context_provider.active() are in-process producer spans and + # should inherit the caller's sampling decision instead. + if operation in _MANUAL_KEEP_OPS and force_keep: + span.set_tag(_MANUAL_KEEP_TAG, True) + + kind = self._SPAN_KIND.get(operation) + if kind: + span.set_tag(self._SPAN_KIND_TAG, kind) + if ( + kind == self._CONSUMER + and parent_service_name + and parent_service_name != self.service_name + ): + span.set_tag(self._PEER_SERVICE_TAG, parent_service_name) diff --git a/temporalio/contrib/datadog/_span_runner.py b/temporalio/contrib/datadog/_span_runner.py new file mode 100644 index 000000000..11fbc9440 --- /dev/null +++ b/temporalio/contrib/datadog/_span_runner.py @@ -0,0 +1,27 @@ +"""Base span lifecycle management class for Datadog interceptors.""" + +from __future__ import annotations + +from collections.abc import Awaitable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from temporalio.contrib.datadog._interceptor import DatadogTracingInterceptor + + +class _SpanRunner: # type: ignore[reportUnusedClass] + def __init__(self, root: DatadogTracingInterceptor) -> None: + self.root = root + + async def run( + self, span: Any, operation_name: str, operation: Awaitable[Any] + ) -> Any: + operation_exc: BaseException | None = None + try: + result = await operation + except BaseException as exc: + operation_exc = exc + raise + finally: + self.root.tracer.finish_span(span, operation_name, operation_exc) + return result diff --git a/temporalio/contrib/datadog/_workflow_interceptor.py b/temporalio/contrib/datadog/_workflow_interceptor.py new file mode 100644 index 000000000..a35c4fbad --- /dev/null +++ b/temporalio/contrib/datadog/_workflow_interceptor.py @@ -0,0 +1,407 @@ +"""Datadog tracing interceptors for Temporal workflow inbound and outbound calls.""" + +import contextvars +import logging +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, NoReturn, cast + +import temporalio.worker +import temporalio.workflow +from temporalio.contrib.datadog._constants import ( + COMMON_ATTRIBUTE_MAP, + OperationNames, + SpanAttributes, +) +from temporalio.contrib.datadog._id_generator import gen_span_id +from temporalio.contrib.datadog._propagator import _Propagator + +# ContextVar keeps this flag task-local. Sandboxed workflows load +# non-passthrough modules into a per-instance module namespace; unsandboxed +# workflow tasks capture their own context. +_trace_disconnected: contextvars.ContextVar[bool] = contextvars.ContextVar( + "_trace_disconnected", default=False +) + +# Live RunWorkflow span for the current execution, set before user code runs. +# Never None — RunWorkflow is exempt from the replay guard in span_ctx. +# Isolated per execution by the same mechanism as _trace_disconnected. +_active_workflow_span: contextvars.ContextVar[Any] = contextvars.ContextVar( + "_active_workflow_span", default=None +) + +# Holds (trace_id, span_id) for the active RunWorkflow span so that the log +# filter below can inject dd.trace_id / dd.span_id into every workflow.logger +# call without activating the span (which is unsafe inside the sandbox). +_current_span_info: contextvars.ContextVar[tuple[str, str] | None] = ( + contextvars.ContextVar("_current_span_info", default=None) +) + + +class _DDTraceLogFilter(logging.Filter): + def filter(self, record: logging.LogRecord) -> bool: + info = _current_span_info.get() + if info is not None: + record.__dict__["dd.trace_id"] = info[0] + record.__dict__["dd.span_id"] = info[1] + return True + + +# Each sandbox instance can import this module again; avoid adding the same +# filter repeatedly to Temporal's shared workflow logger. +if not getattr( + temporalio.workflow.logger.base_logger, "_dd_trace_filter_installed", False +): + temporalio.workflow.logger.base_logger.addFilter(_DDTraceLogFilter()) + temporalio.workflow.logger.base_logger._dd_trace_filter_installed = True # type: ignore[attr-defined] + + +@dataclass +class WorkflowTracingConfig: + # Set the relevant flag to suppress signal, query, or update spans while + # keeping workflow and activity tracing enabled. + disable_signal_tracing: bool + disable_query_tracing: bool + disable_update_tracing: bool + + @staticmethod + def default_config() -> "WorkflowTracingConfig": + return WorkflowTracingConfig( + disable_signal_tracing=False, + disable_query_tracing=False, + disable_update_tracing=False, + ) + + +class DatadogTracingWorkflowInboundInterceptor( + temporalio.worker.WorkflowInboundInterceptor +): + def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None: + super().__init__(next) + self._span_counter = 1 # Reserve the 1 for the RunWorkflow span + + externs = temporalio.workflow.extern_functions() + self._start_span_extern = cast( + Callable[ + [str, str, dict[str, Any], Any | None, str | None, int | None], Any + ], + externs["__temporal_datadog_start_sandboxed_span"], + ) + self._finish_span_extern = cast( + Callable[[str, Any | None, BaseException | None], None], + externs["__temporal_datadog_finish_sandboxed_span"], + ) + config_func = cast( + Callable[[], tuple[_Propagator, WorkflowTracingConfig]], + externs["__temporal_datadog_configure_workflow_tracing"], + ) + self.propagator, self.config = config_func() + + def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None: + super().init(_WorkflowOutboundInterceptor(outbound, self)) + + @contextmanager + def span_ctx( + self, + operation_name: str, + resource_name: str, + input: Any, + idempotency_key: str | None = None, + start_time: int | None = None, + ) -> Generator[tuple[Any, Any], None, None]: + attributes = self._get_span_attributes(input) + parent_ctx = self._parent_ctx_for(operation_name, input) + + # Idempotency-keyed HandleSignal and HandleUpdate spans are suppressed + # while replaying. RunWorkflow is exempt so a restarted worker recreates + # its long-lived span. HandleQuery and ValidateUpdate pass no idempotency + # key, so this guard does not suppress them. + if ( + idempotency_key is not None + and operation_name != OperationNames.RUN_WORKFLOW + and temporalio.workflow.unsafe.is_replaying() + ): + span = None + else: + span = self._start_span_extern( + operation_name, + resource_name, + attributes, + parent_ctx, + idempotency_key, + start_time, + ) + exc: BaseException | None = None + try: + yield input, span + except BaseException as e: + exc = e + raise + finally: + self._finish_span_extern(operation_name, span, exc) + + def _parent_ctx_for(self, operation_name: str, input: Any) -> Any: + # Parent from workflow header + if operation_name == OperationNames.RUN_WORKFLOW: + return self.propagator.extract_headers(temporalio.workflow.info().headers) + + # Parent from input headers + ctx = self.propagator.extract_headers(input.headers) + if ctx is not None: + return ctx + + # Make RunWorkflow the parent + return self.recover_workflow_span() + + def recover_workflow_span(self) -> Any: + # Reconstruct the RunWorkflow context from the workflow start headers by + # reusing the trace_id from StartWorkflow and overriding the span_id with + # RunWorkflow's deterministic ID. Reliable across sandbox task boundaries + # because workflow.info().headers is always available. + ctx = self.propagator.extract_headers(temporalio.workflow.info().headers) + if ctx is not None: + ctx.span_id = gen_span_id(self._make_idempotency_key(1)) + return ctx + + # No start headers (uninstrumented client). Fall back to the live span. + span = span_from_workflow_context() + return span.context if span is not None else None + + def _get_span_attributes(self, input: Any) -> dict[str, Any]: + info = temporalio.workflow.info() + attrs: dict[str, Any] = { + SpanAttributes.WORKFLOW_ID: info.workflow_id, + SpanAttributes.RUN_ID: info.run_id, + SpanAttributes.WORKFLOW_TYPE: info.workflow_type, + } + for field, span_key in COMMON_ATTRIBUTE_MAP: + if val := getattr(input, field, None): + attrs[span_key] = val + if getattr(input, "update", None): + attrs[SpanAttributes.UPDATE_NAME] = input.update + if getattr(input, "id", None): + attrs[SpanAttributes.UPDATE_ID] = input.id + if getattr(input, "workflow", None): + attrs[SpanAttributes.CHILD_WORKFLOW_TYPE] = input.workflow + if getattr(input, "id", None): + attrs[SpanAttributes.CHILD_WORKFLOW_ID] = input.id + if isinstance(input, temporalio.worker.StartLocalActivityInput): + attrs[SpanAttributes.LOCAL] = True + return attrs + + def _make_idempotency_key(self, counter: int) -> str: + info = temporalio.workflow.info() + # Matches the Go SDK's idempotency key + return f"WorkflowInboundInterceptor:{info.namespace}:{info.workflow_id}:{info.run_id}:{counter}" + + def _next_idempotency_key(self) -> str: + self._span_counter += 1 + return self._make_idempotency_key(self._span_counter) + + async def execute_workflow( + self, input: temporalio.worker.ExecuteWorkflowInput + ) -> Any: + info = temporalio.workflow.info() + with self.span_ctx( + OperationNames.RUN_WORKFLOW, + info.workflow_type, + input, + idempotency_key=self._make_idempotency_key(1), + start_time=int(info.workflow_start_time.timestamp() * 1e9), + ) as (i, span): + if span is not None: + i.headers = self.propagator.inject_headers(i.headers, span.context) + tid = span.context.trace_id + # Match ddtrace's format_trace_id: decimal for 64-bit IDs, 32-char hex for 128-bit. + formatted_tid = f"{tid:032x}" if tid > (1 << 64) - 1 else str(tid) + _current_span_info.set((formatted_tid, str(span.span_id))) + _active_workflow_span.set(span) + return await super().execute_workflow(i) + + async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None: + if self.config.disable_signal_tracing: + return await super().handle_signal(input) + with self.span_ctx( + OperationNames.HANDLE_SIGNAL, + input.signal, + input, + self._next_idempotency_key(), + ) as (i, _): + await super().handle_signal(i) + + async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any: + if self.config.disable_query_tracing: + return await super().handle_query(input) + with self.span_ctx(OperationNames.HANDLE_QUERY, input.query, input) as (i, _): + return await super().handle_query(i) + + def handle_update_validator( + self, input: temporalio.worker.HandleUpdateInput + ) -> None: + if self.config.disable_update_tracing: + return super().handle_update_validator(input) + with self.span_ctx(OperationNames.VALIDATE_UPDATE, input.update, input) as ( + i, + _, + ): + super().handle_update_validator(i) + + async def handle_update_handler( + self, input: temporalio.worker.HandleUpdateInput + ) -> Any: + if self.config.disable_update_tracing: + return await super().handle_update_handler(input) + with self.span_ctx( + OperationNames.HANDLE_UPDATE, + input.update, + input, + self._next_idempotency_key(), + ) as (i, _): + return await super().handle_update_handler(i) + + +class _WorkflowOutboundInterceptor(temporalio.worker.WorkflowOutboundInterceptor): + def __init__( + self, + next: temporalio.worker.WorkflowOutboundInterceptor, + root: DatadogTracingWorkflowInboundInterceptor, + ) -> None: + super().__init__(next) + self.root = root + + def continue_as_new(self, input: temporalio.worker.ContinueAsNewInput) -> NoReturn: + if not _trace_disconnected.get(): + input.headers = self.root.propagator.inject_headers( + input.headers, self.root.recover_workflow_span() + ) + super().continue_as_new(input) + + async def signal_child_workflow( + self, input: temporalio.worker.SignalChildWorkflowInput + ) -> None: + if self.root.config.disable_signal_tracing: + return await super().signal_child_workflow(input) + if temporalio.workflow.unsafe.is_replaying(): + return await super().signal_child_workflow(input) + with self.root.span_ctx( + OperationNames.SIGNAL_CHILD_WORKFLOW, input.signal, input + ) as (i, span): + if span is not None: + i.headers = self.root.propagator.inject_headers(i.headers, span.context) + await super().signal_child_workflow(i) + + async def signal_external_workflow( + self, input: temporalio.worker.SignalExternalWorkflowInput + ) -> None: + if self.root.config.disable_signal_tracing: + return await super().signal_external_workflow(input) + if temporalio.workflow.unsafe.is_replaying(): + return await super().signal_external_workflow(input) + with self.root.span_ctx( + OperationNames.SIGNAL_EXTERNAL_WORKFLOW, input.signal, input + ) as (i, span): + if span is not None: + i.headers = self.root.propagator.inject_headers(i.headers, span.context) + await super().signal_external_workflow(i) + + def start_activity( + self, input: temporalio.worker.StartActivityInput + ) -> temporalio.workflow.ActivityHandle: + if temporalio.workflow.unsafe.is_replaying(): + return super().start_activity(input) + with self.root.span_ctx( + OperationNames.START_ACTIVITY, input.activity, input + ) as (i, span): + if span is not None: + i.headers = self.root.propagator.inject_headers(i.headers, span.context) + return super().start_activity(i) + + async def start_child_workflow( + self, input: temporalio.worker.StartChildWorkflowInput + ) -> temporalio.workflow.ChildWorkflowHandle: + if temporalio.workflow.unsafe.is_replaying(): + return await super().start_child_workflow(input) + with self.root.span_ctx( + OperationNames.START_CHILD_WORKFLOW, input.workflow, input + ) as (i, span): + if span is not None: + i.headers = self.root.propagator.inject_headers(i.headers, span.context) + return await super().start_child_workflow(i) + + def start_local_activity( + self, input: temporalio.worker.StartLocalActivityInput + ) -> temporalio.workflow.ActivityHandle: + if temporalio.workflow.unsafe.is_replaying(): + return super().start_local_activity(input) + with self.root.span_ctx( + OperationNames.START_ACTIVITY, input.activity, input + ) as (i, span): + if span is not None: + i.headers = self.root.propagator.inject_headers(i.headers, span.context) + return super().start_local_activity(i) + + async def start_nexus_operation( + self, input: temporalio.worker.StartNexusOperationInput[Any, Any] + ) -> temporalio.workflow.NexusOperationHandle[Any]: + # Skip Nexus tracing during workflow replay so replay does not emit a + # duplicate StartNexusOperation span. + if temporalio.workflow.unsafe.is_replaying(): + return await super().start_nexus_operation(input) + + with self.root.span_ctx( + OperationNames.START_NEXUS_OPERATION, + f"{input.service}/{input.operation_name}", + input, + ) as (i, span): + if span is not None: + # Nexus uses plain string headers, not Temporal payload headers. + carrier = self.root.propagator.inject(span.context) + i.headers = {**(i.headers or {}), **carrier} + return await super().start_nexus_operation(i) + + +def span_from_workflow_context() -> Any: + """Return the active RunWorkflow ddtrace span for this execution. + + Always returns a live span, including during replay on a new worker, so + custom tags set here survive a worker restart:: + + span = span_from_workflow_context() + if span is not None: + span.set_tag("my.tag", value) + + Python equivalent of the Go SDK's ``SpanFromWorkflowContext``. Unlike the + Go version, which takes a ``workflow.Context`` and can return any + operation's span, this always returns the RunWorkflow span. + """ + # NOTE: fork-specific workaround; reconsider if this integration is ever proposed upstream. + # Prefer the extern so that the call reaches the host-side ContextVar even + # when the workflow sandbox has reimported this module into its own namespace. + fn = temporalio.workflow.extern_functions().get( + "__temporal_datadog_get_workflow_span" + ) + if fn is not None: + return fn() + return _active_workflow_span.get() + + +def disconnect_trace_span_from_workflow_context() -> None: + """Prevent the current trace from propagating into the next ContinueAsNew execution. + + Call before ``workflow.continue_as_new()``; the next run starts a fresh + root span rather than continuing this trace:: + + disconnect_trace_span_from_workflow_context() + workflow.continue_as_new(count + 1) + + """ + _trace_disconnected.set(True) + # NOTE: fork-specific workaround; reconsider if this integration is ever proposed upstream. + # Call through the extern so that the flag is set on the host-side ContextVar + # where _WorkflowOutboundInterceptor.continue_as_new reads it. + fn = temporalio.workflow.extern_functions().get( + "__temporal_datadog_set_trace_disconnected" + ) + if fn is not None: + fn() diff --git a/temporalio/contrib/datadog/_wrapped_tracer.py b/temporalio/contrib/datadog/_wrapped_tracer.py new file mode 100644 index 000000000..839200dc5 --- /dev/null +++ b/temporalio/contrib/datadog/_wrapped_tracer.py @@ -0,0 +1,148 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +import temporalio.activity +import temporalio.workflow +from temporalio.contrib.datadog._constants import TEMPORAL_TAG_PREFIX + +# ``ddtrace`` is intentionally not imported at module level. ``WrappedTracer`` +# resolves and caches the tracer and Context class in the host process so sandbox +# extern calls do not import ``ddtrace``. + + +@dataclass(frozen=True) +class FinishContext: + """Context passed to a user-supplied ``on_span_finish`` callback. + + Attributes: + operation: The Temporal operation name (e.g. ``"RunWorkflow"``). + exception: The exception that caused the span to fail, or ``None``. + """ + + operation: str + exception: BaseException | None + + +@dataclass(frozen=True) +class FinishResult: + """Returned by ``on_span_finish`` to control how a span is finished. + + All fields default to leaving the interceptor's default behavior in place; + return ``None`` from the callback (or omit the callback) for the default. + + Attributes: + extra_tags: Tags applied to the span before it is finished. + """ + + extra_tags: Mapping[str, Any] | None = None + + +class WrappedTracer: + tracer: Any + + def __init__( + self, + *, + service_name: str | None, + tracer: Any, + on_span_finish: Callable[[FinishContext], FinishResult | None] | None = None, + annotator: Any, + propagator: Any, + ): + if tracer is None: + import ddtrace + + tracer = ddtrace.tracer # type: ignore[reportPrivateImportUsage] + + # Cache Context here so start_span never imports it as an extern call. + # Deferred imports inside externs go through the sandbox's restricted + # importer and fail with RestrictedWorkflowAccessError. + from ddtrace._trace.context import ( + Context, # type: ignore[reportPrivateImportUsage] + ) + + self.ctx_cls: Any = Context + else: + self.ctx_cls = None + self.tracer = tracer + self.service_name = service_name + self.on_span_finish = on_span_finish + self.annotator = annotator + self.propagator = propagator + + def start_span( + self, + *, + operation_name: str, + parent_ctx: Any, + resource_name: str, + activate: bool, + start_time: Any = None, + span_id: int | None = None, + attributes: Mapping[str, Any] | None = None, + parent_from_header: bool = False, + trace_id: int | None = None, + ) -> Any: + # No DD header (uninstrumented client) the call passes a deterministic trace_id + # to keep the RunWorkflow trace consistent if the worker restarts mid-run. + # Mutating span.trace_id after creation breaks the tracer's internal trace registry. + effective_parent = parent_ctx + if trace_id is not None and parent_ctx is None and self.ctx_cls is not None: + effective_parent = self.ctx_cls( + trace_id=trace_id, span_id=None, is_remote=True + ) + span = self.tracer.start_span( + name=f"{TEMPORAL_TAG_PREFIX}{operation_name}", + child_of=effective_parent, + service=self.service_name, + resource=resource_name, + activate=activate, + ) + if start_time is not None: + span.start_ns = start_time + if span_id is not None: + span.span_id = span_id + if getattr(span, "context", None) is not None: + span.context.span_id = span_id + force_keep = parent_ctx is None or parent_from_header + self.annotator.annotate( + span, + operation_name, + attributes, + self.propagator.get_baggage(parent_ctx), + force_keep, + ) + self.propagator.set_baggage(span.context) + return span + + def finish_span( + self, + span: Any, + operation_name: str, + exc: BaseException | None, + ) -> None: + try: + result: FinishResult | None = None + if self.on_span_finish is not None: + result = self.on_span_finish( + FinishContext(operation=operation_name, exception=exc) + ) + + if exc and not self._should_skip_error(exc): + span.set_exc_info(type(exc), exc, exc.__traceback__) + + if result is not None and result.extra_tags: + for key, value in result.extra_tags.items(): + span.set_tag(key, value) + finally: + span.finish() + + def _should_skip_error(self, exc: BaseException | None) -> bool: + if exc is None: + return True + if isinstance(exc, temporalio.workflow.ContinueAsNewError): + return True + if isinstance(exc, temporalio.activity._CompleteAsyncError): + return True + return False diff --git a/tests/contrib/datadog/__init__.py b/tests/contrib/datadog/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/contrib/datadog/_workflows.py b/tests/contrib/datadog/_workflows.py new file mode 100644 index 000000000..177a81237 --- /dev/null +++ b/tests/contrib/datadog/_workflows.py @@ -0,0 +1,309 @@ +"""Workflow and activity definitions for Datadog tracing tests. + +Kept separate from ``test_tracing.py`` so the workflow sandbox can re-import +workflow definitions without importing that test module's host-only ``ddtrace`` +setup. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, workflow +from temporalio.common import RetryPolicy +from temporalio.contrib.datadog import ( + disconnect_trace_span_from_workflow_context, + span_from_workflow_context, +) +from temporalio.exceptions import ApplicationError + + +@dataclass +class TestRequest: + use_activity: bool = False + use_local_activity: bool = False + use_child_workflow: bool = False + wait_for_kick: bool = False + + +@activity.defn +async def echo_activity(value: str) -> str: + return value + + +@activity.defn +async def failing_activity() -> str: + raise ApplicationError("boom") + + +@workflow.defn +class TestWorkflow: + def __init__(self) -> None: + self._kicked = False + + @workflow.run + async def run(self, param: TestRequest) -> str: + if param.use_activity: + await workflow.execute_activity( + echo_activity, + "hello", + start_to_close_timeout=timedelta(seconds=10), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + if param.use_local_activity: + await workflow.execute_local_activity( + echo_activity, + "hello", + start_to_close_timeout=timedelta(seconds=10), + ) + if param.use_child_workflow: + await workflow.execute_child_workflow( + ChildWorkflow.run, + "hello", + id=f"{workflow.info().workflow_id}-child", + ) + if param.wait_for_kick: + await workflow.wait_condition(lambda: self._kicked) + return "done" + + @workflow.signal + def kick(self) -> None: + self._kicked = True + + @workflow.query + def get_status(self) -> str: + return "running" + + +@workflow.defn +class UpdateTestWorkflow: + def __init__(self) -> None: + self._result: str | None = None + + @workflow.run + async def run(self) -> str: + await workflow.wait_condition(lambda: self._result is not None) + return self._result or "" + + @workflow.update + async def do_update(self, value: str) -> str: + self._result = value + return f"updated:{value}" + + @do_update.validator + def validate_do_update(self, value: str) -> None: + if value == "invalid": + raise ApplicationError("invalid value") + + +@workflow.defn +class ChildWorkflow: + @workflow.run + async def run(self, value: str) -> str: + return f"child:{value}" + + +@workflow.defn +class WaitingChildWorkflow: + def __init__(self) -> None: + self._done = False + + @workflow.run + async def run(self) -> str: + await workflow.wait_condition(lambda: self._done) + return "done" + + @workflow.signal + def finish(self) -> None: + self._done = True + + +@workflow.defn +class ParentWithSignalChildWorkflow: + """Starts a child, signals it via signal_child_workflow, awaits it, then waits for kick.""" + + def __init__(self) -> None: + self._kicked = False + + @workflow.run + async def run(self) -> str: + child_handle = await workflow.start_child_workflow( + WaitingChildWorkflow.run, + id=f"{workflow.info().workflow_id}-child", + ) + await child_handle.signal(WaitingChildWorkflow.finish) + await child_handle + await workflow.wait_condition(lambda: self._kicked) + return "done" + + @workflow.signal + def kick(self) -> None: + self._kicked = True + + @workflow.query + def get_status(self) -> str: + return "running" + + +@workflow.defn +class ParentWorkflow: + @workflow.run + async def run(self) -> str: + return await workflow.execute_child_workflow( + ChildWorkflow.run, + "hello", + id=f"{workflow.info().workflow_id}-child", + ) + + +@workflow.defn +class LocalActivityWorkflow: + @workflow.run + async def run(self) -> str: + return await workflow.execute_local_activity( + echo_activity, + "hello", + start_to_close_timeout=timedelta(seconds=10), + ) + + +@workflow.defn +class FailingWorkflow: + @workflow.run + async def run(self) -> str: + return await workflow.execute_activity( + failing_activity, + start_to_close_timeout=timedelta(seconds=10), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + + +@workflow.defn +class ContinueAsNewWorkflow: + @workflow.run + async def run(self, iteration: int = 0) -> str: + if iteration == 0: + workflow.continue_as_new(1) + return "done" + + +@workflow.defn +class DisconnectedContinueAsNewWorkflow: + @workflow.run + async def run(self, iteration: int = 0) -> str: + if iteration == 0: + disconnect_trace_span_from_workflow_context() + workflow.continue_as_new(1) + return "done" + + +@workflow.defn +class CustomTagWorkflow: + def __init__(self) -> None: + self._kicked = False + + @workflow.run + async def run(self, wait_for_kick: bool = False) -> str: + span = span_from_workflow_context() + if span is not None: + span.set_tag("custom.workflow.tag", "hello-from-workflow") + if wait_for_kick: + await workflow.wait_condition(lambda: self._kicked) + return "done" + + @workflow.signal + def kick(self) -> None: + self._kicked = True + + @workflow.query + def get_status(self) -> str: + return "running" + + +@workflow.defn +class DirectlyFailingWorkflow: + @workflow.run + async def run(self) -> str: + raise ApplicationError("workflow failed directly") + + +@activity.defn +async def logging_activity() -> str: + activity.logger.info("test log message from activity") + return "done" + + +@activity.defn +async def custom_span_activity() -> tuple[int | None, int | None]: + """Create a custom ddtrace span and return its (parent_id, trace_id). + + Uses tracer.trace() which auto-parents from the active context span. + """ + import ddtrace + + child = ddtrace.tracer.trace("custom.span") # type: ignore[reportPrivateImportUsage] + try: + return (child.parent_id, child.trace_id) + finally: + child.finish() + + +@workflow.defn +class LoggingWorkflow: + @workflow.run + async def run(self) -> str: + workflow.logger.info("test log message from workflow") + return "done" + + +@workflow.defn +class CustomSpanActivityWorkflow: + @workflow.run + async def run(self) -> tuple[int | None, int | None]: + return await workflow.execute_activity( + custom_span_activity, + start_to_close_timeout=timedelta(seconds=10), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + + +@workflow.defn +class LoggingActivityWorkflow: + @workflow.run + async def run(self) -> str: + return await workflow.execute_activity( + logging_activity, + start_to_close_timeout=timedelta(seconds=10), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + + +@workflow.defn +class ConcurrentLoggingWorkflow: + """Logs before and after a signal barrier so two instances provably overlap. + + The test starts both workflows, waits until both have logged "start: