Add distributed tracing using opentelemetry

andystaples · andystaples · commit 06f30643a9f2 · 2026-03-03T11:23:35.000-07:00
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -1,2 +1,4 @@
 grpcio-tools
+opentelemetry-api
+opentelemetry-sdk
 pymarkdownlnt
diff --git a/durabletask/client.py b/durabletask/client.py
@@ -16,6 +16,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
+import durabletask.internal.tracing as tracing
 from durabletask import task
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
 
@@ -177,7 +178,8 @@ def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInpu
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
             version=helpers.get_string_value(version if version else self.default_version),
             orchestrationIdReusePolicy=reuse_id_policy,
-            tags=tags
+            tags=tags,
+            parentTraceContext=tracing.get_current_trace_context(),
         )
 
         self._logger.info(f"Starting new '{name}' instance with ID = '{req.instanceId}'.")
@@ -355,7 +357,7 @@ def signal_entity(self,
             input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
             requestId=str(uuid.uuid4()),
             scheduledTime=None,
-            parentTraceContext=None,
+            parentTraceContext=tracing.get_current_trace_context(),
             requestTime=helpers.new_timestamp(datetime.now(timezone.utc))
         )
         self._logger.info(f"Signaling entity '{entity_instance_id}' operation '{operation_name}'.")
diff --git a/durabletask/internal/helpers.py b/durabletask/internal/helpers.py
@@ -219,11 +219,13 @@ def new_create_timer_action(id: int, fire_at: datetime) -> pb.OrchestratorAction
 
 
 def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
-                             tags: Optional[dict[str, str]]) -> pb.OrchestratorAction:
+                             tags: Optional[dict[str, str]],
+                             parent_trace_context: Optional[pb.TraceContext] = None) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, scheduleTask=pb.ScheduleTaskAction(
         name=name,
         input=get_string_value(encoded_input),
-        tags=tags
+        tags=tags,
+        parentTraceContext=parent_trace_context,
     ))
 
 
@@ -298,12 +300,14 @@ def new_create_sub_orchestration_action(
         name: str,
         instance_id: Optional[str],
         encoded_input: Optional[str],
-        version: Optional[str]) -> pb.OrchestratorAction:
+        version: Optional[str],
+        parent_trace_context: Optional[pb.TraceContext] = None) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, createSubOrchestration=pb.CreateSubOrchestrationAction(
         name=name,
         instanceId=instance_id,
         input=get_string_value(encoded_input),
-        version=get_string_value(version)
+        version=get_string_value(version),
+        parentTraceContext=parent_trace_context,
     ))
 
 
diff --git a/durabletask/internal/tracing.py b/durabletask/internal/tracing.py
@@ -0,0 +1,165 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""OpenTelemetry distributed tracing utilities for the Durable Task SDK.
+
+This module provides helpers for propagating W3C Trace Context between
+orchestrations, activities, sub-orchestrations, and entities via the
+``TraceContext`` protobuf message carried over gRPC.
+
+OpenTelemetry is an **optional** dependency.  When the ``opentelemetry-api``
+package is not installed every helper gracefully degrades to a no-op so
+that the rest of the SDK continues to work without any tracing overhead.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import Optional
+
+from google.protobuf import wrappers_pb2
+
+import durabletask.internal.orchestrator_service_pb2 as pb
+
+logger = logging.getLogger("durabletask-tracing")
+
+# ---------------------------------------------------------------------------
+# Lazy / optional OpenTelemetry imports
+# ---------------------------------------------------------------------------
+try:
+    from opentelemetry import context as otel_context
+    from opentelemetry import trace
+    from opentelemetry.trace import (
+        SpanKind,
+        StatusCode,
+    )
+    from opentelemetry.trace.propagation.tracecontext import (
+        TraceContextTextMapPropagator,
+    )
+
+    _OTEL_AVAILABLE = True
+except ImportError:  # pragma: no cover
+    _OTEL_AVAILABLE = False
+
+# Re-export so callers can check without importing opentelemetry themselves.
+OTEL_AVAILABLE = _OTEL_AVAILABLE
+
+# The instrumentation scope name used when creating spans.
+_TRACER_NAME = "durabletask"
+
+
+# ---------------------------------------------------------------------------
+# Public helpers – extracting / injecting trace context
+# ---------------------------------------------------------------------------
+
+def get_current_trace_context() -> Optional[pb.TraceContext]:
+    """Capture the current OpenTelemetry span context as a protobuf ``TraceContext``.
+
+    Returns ``None`` when OpenTelemetry is not installed or there is no
+    active span.
+    """
+    if not _OTEL_AVAILABLE:
+        return None
+
+    propagator = TraceContextTextMapPropagator()
+    carrier: dict[str, str] = {}
+    propagator.inject(carrier)
+
+    traceparent = carrier.get("traceparent")
+    if not traceparent:
+        return None
+
+    tracestate = carrier.get("tracestate")
+
+    # Extract the span ID from the traceparent header.
+    # Format: 00-<trace-id>-<span-id>-<flags>
+    parts = traceparent.split("-")
+    span_id = parts[2] if len(parts) >= 4 else ""
+
+    return pb.TraceContext(
+        traceParent=traceparent,
+        spanID=span_id,
+        traceState=wrappers_pb2.StringValue(value=tracestate) if tracestate else None,
+    )
+
+
+def extract_trace_context(proto_ctx: Optional[pb.TraceContext]) -> Optional[object]:
+    """Convert a protobuf ``TraceContext`` into an OpenTelemetry ``Context``.
+
+    Returns ``None`` when OpenTelemetry is not installed or the supplied
+    context is empty / ``None``.
+    """
+    if not _OTEL_AVAILABLE or proto_ctx is None:
+        return None
+
+    traceparent = proto_ctx.traceParent
+    if not traceparent:
+        return None
+
+    carrier: dict[str, str] = {"traceparent": traceparent}
+    if proto_ctx.HasField("traceState") and proto_ctx.traceState.value:
+        carrier["tracestate"] = proto_ctx.traceState.value
+
+    propagator = TraceContextTextMapPropagator()
+    ctx = propagator.extract(carrier)
+    return ctx
+
+
+@contextmanager
+def start_span(
+    name: str,
+    trace_context: Optional[pb.TraceContext] = None,
+    kind: Optional[object] = None,
+    attributes: Optional[dict[str, str]] = None,
+):
+    """Context manager that starts an OpenTelemetry span linked to a parent trace context.
+
+    If OpenTelemetry is not installed, the block executes without tracing.
+
+    Parameters
+    ----------
+    name:
+        Human-readable span name (e.g. ``"activity:say_hello"``).
+    trace_context:
+        The protobuf ``TraceContext`` received from the sidecar.  When
+        provided the new span will be created as a **child** of this
+        context.
+    kind:
+        The ``SpanKind`` for the new span.  Defaults to ``SpanKind.INTERNAL``.
+    attributes:
+        Optional dictionary of span attributes.
+    """
+    if not _OTEL_AVAILABLE:
+        yield None
+        return
+
+    parent_ctx = extract_trace_context(trace_context)
+
+    if kind is None:
+        kind = SpanKind.INTERNAL
+
+    tracer = trace.get_tracer(_TRACER_NAME)
+
+    if parent_ctx is not None:
+        token = otel_context.attach(parent_ctx)
+        try:
+            with tracer.start_as_current_span(
+                name, kind=kind, attributes=attributes
+            ) as span:
+                yield span
+        finally:
+            otel_context.detach(token)
+    else:
+        with tracer.start_as_current_span(
+            name, kind=kind, attributes=attributes
+        ) as span:
+            yield span
+
+
+def set_span_error(span, ex: Exception) -> None:
+    """Record an exception on the given span (if tracing is available)."""
+    if not _OTEL_AVAILABLE or span is None:
+        return
+    span.set_status(StatusCode.ERROR, str(ex))
+    span.record_exception(ex)
diff --git a/durabletask/worker.py b/durabletask/worker.py
@@ -32,6 +32,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
+import durabletask.internal.tracing as tracing
 from durabletask import task
 from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
 
@@ -641,6 +642,7 @@ def _execute_orchestrator(
                 actions=result.actions,
                 customStatus=ph.get_string_value(result.encoded_custom_status),
                 completionToken=completionToken,
+                orchestrationTraceContext=req.orchestrationTraceContext,
             )
         except pe.AbandonOrchestrationError:
             self._logger.info(
@@ -697,9 +699,20 @@ def _execute_activity(
         instance_id = req.orchestrationInstance.instanceId
         try:
             executor = _ActivityExecutor(self._registry, self._logger)
-            result = executor.execute(
-                instance_id, req.name, req.taskId, req.input.value
-            )
+            with tracing.start_span(
+                f"activity:{req.name}",
+                trace_context=req.parentTraceContext,
+                attributes={"durabletask.task.instance_id": instance_id,
+                            "durabletask.task.name": req.name,
+                            "durabletask.task.task_id": str(req.taskId)},
+            ) as span:
+                try:
+                    result = executor.execute(
+                        instance_id, req.name, req.taskId, req.input.value
+                    )
+                except Exception as ex:
+                    tracing.set_span_error(span, ex)
+                    raise
             res = pb.ActivityResponse(
                 instanceId=instance_id,
                 taskId=req.taskId,
@@ -759,30 +772,41 @@ def _execute_entity_batch(
 
             operation_result = None
 
-            try:
-                entity_result = executor.execute(
-                    instance_id, entity_instance_id, operation.operation, entity_state, operation.input.value
-                )
-
-                entity_result = ph.get_string_value_or_empty(entity_result)
-                operation_result = pb.OperationResult(success=pb.OperationResultSuccess(
-                    result=entity_result,
-                    startTimeUtc=new_timestamp(start_time),
-                    endTimeUtc=new_timestamp(datetime.now(timezone.utc))
-                ))
-                results.append(operation_result)
+            # Get the trace context for this operation, if available
+            op_trace_ctx = operation.traceContext if operation.HasField("traceContext") else None
 
-                entity_state.commit()
-            except Exception as ex:
-                self._logger.exception(ex)
-                operation_result = pb.OperationResult(failure=pb.OperationResultFailure(
-                    failureDetails=ph.new_failure_details(ex),
-                    startTimeUtc=new_timestamp(start_time),
-                    endTimeUtc=new_timestamp(datetime.now(timezone.utc))
-                ))
-                results.append(operation_result)
+            with tracing.start_span(
+                f"entity:{entity_instance_id.entity}:{operation.operation}",
+                trace_context=op_trace_ctx,
+                attributes={"durabletask.entity.instance_id": instance_id,
+                            "durabletask.entity.name": entity_instance_id.entity,
+                            "durabletask.entity.operation": operation.operation},
+            ) as span:
+                try:
+                    entity_result = executor.execute(
+                        instance_id, entity_instance_id, operation.operation, entity_state, operation.input.value
+                    )
 
-                entity_state.rollback()
+                    entity_result = ph.get_string_value_or_empty(entity_result)
+                    operation_result = pb.OperationResult(success=pb.OperationResultSuccess(
+                        result=entity_result,
+                        startTimeUtc=new_timestamp(start_time),
+                        endTimeUtc=new_timestamp(datetime.now(timezone.utc))
+                    ))
+                    results.append(operation_result)
+
+                    entity_state.commit()
+                except Exception as ex:
+                    tracing.set_span_error(span, ex)
+                    self._logger.exception(ex)
+                    operation_result = pb.OperationResult(failure=pb.OperationResultFailure(
+                        failureDetails=ph.new_failure_details(ex),
+                        startTimeUtc=new_timestamp(start_time),
+                        endTimeUtc=new_timestamp(datetime.now(timezone.utc))
+                    ))
+                    results.append(operation_result)
+
+                    entity_state.rollback()
 
         batch_result = pb.EntityBatchResult(
             results=results,
@@ -847,6 +871,7 @@ def __init__(self, instance_id: str, registry: _Registry):
         self._new_input: Optional[Any] = None
         self._save_events = False
         self._encoded_custom_status: Optional[str] = None
+        self._parent_trace_context: Optional[pb.TraceContext] = None
 
     def run(self, generator: Generator[task.Task, Any, Any]):
         self._generator = generator
@@ -1136,15 +1161,18 @@ def call_activity_function_helper(
                 if isinstance(activity_function, str)
                 else task.get_name(activity_function)
             )
-            action = ph.new_schedule_task_action(id, name, encoded_input, tags)
+            action = ph.new_schedule_task_action(
+                id, name, encoded_input, tags,
+                parent_trace_context=self._parent_trace_context)
         else:
             if instance_id is None:
                 # Create a deteministic instance ID based on the parent instance ID
                 instance_id = f"{self.instance_id}:{id:04x}"
             if not isinstance(activity_function, str):
                 raise ValueError("Orchestrator function name must be a string")
             action = ph.new_create_sub_orchestration_action(
-                id, activity_function, instance_id, encoded_input, version
+                id, activity_function, instance_id, encoded_input, version,
+                parent_trace_context=self._parent_trace_context
             )
         self._pending_actions[id] = action
 
@@ -1397,6 +1425,10 @@ def process_event(
                 if event.executionStarted.version:
                     ctx._version = event.executionStarted.version.value
 
+                # Store the parent trace context for propagation to child tasks
+                if event.executionStarted.HasField("parentTraceContext"):
+                    ctx._parent_trace_context = event.executionStarted.parentTraceContext
+
                 if self._registry.versioning:
                     version_failure = self.evaluate_orchestration_versioning(
                         self._registry.versioning,
diff --git a/pyproject.toml b/pyproject.toml
@@ -31,6 +31,11 @@ dependencies = [
     "packaging"
 ]
 
+[project.optional-dependencies]
+opentelemetry = [
+    "opentelemetry-api>=1.0.0"
+]
+
 [project.urls]
 repository = "https://github.com/microsoft/durabletask-python"
 changelog = "https://github.com/microsoft/durabletask-python/blob/main/CHANGELOG.md"
diff --git a/tests/durabletask/test_tracing.py b/tests/durabletask/test_tracing.py
