[Core] Fix traceparent header in DistributedTracingPolicy (#40074)

When we get the trace context headers, we should first ensure that we are
inside the HTTP span's context to ensure that the traceparent header
contains the correct span ID.

Signed-off-by: Paul Van Eck <[email protected]>

* Fix mypy

Signed-off-by: Paul Van Eck <[email protected]>

---------

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -29,6 +29,8 @@
 
 ### Bugs Fixed
 
+- Fixed an issue where the `traceparent` header was not being set correctly in the `DistributedTracingPolicy`. The `traceparent` header will now set based on the context of the HTTP client span. #40074
+
 ### Other Changes
 
 - Added `opentelemetry-api` as an optional dependency for tracing. This can be installed with `pip install azure-core[tracing]`. #39563

sdk/core/azure-core/azure/core/pipeline/policies/_distributed_tracing.py

@@ -39,6 +39,7 @@
 from azure.core.rest import HttpResponse, HttpRequest
 from azure.core.settings import settings
 from azure.core.tracing import SpanKind
+from azure.core.tracing.common import change_context
 from azure.core.instrumentation import get_tracer
 from azure.core.tracing._models import TracingOptions
 
@@ -127,8 +128,9 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
                 for attr, value in span_attributes.items():
                     span.add_attribute(attr, value)  # type: ignore
 
-                headers = span.to_header()
-                request.http_request.headers.update(headers)
+                with change_context(span.span_instance):
+                    headers = span.to_header()
+                    request.http_request.headers.update(headers)
                 request.context[self.TRACING_CONTEXT] = span
             else:
                 # Otherwise, use the core tracing.
@@ -151,10 +153,11 @@ def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
                     attributes=span_attributes,
                 )
 
-                trace_context_headers = tracer.get_trace_context()
-                request.http_request.headers.update(trace_context_headers)
-                request.context[self.TRACING_CONTEXT] = otel_span
+                with tracer.use_span(otel_span, end_on_exit=False):
+                    trace_context_headers = tracer.get_trace_context()
+                    request.http_request.headers.update(trace_context_headers)
 
+                request.context[self.TRACING_CONTEXT] = otel_span
                 token = tracer._suppress_auto_http_instrumentation()  # pylint: disable=protected-access
                 request.context[self._SUPPRESSION_TOKEN] = token
 

sdk/core/azure-core/azure/core/tracing/opentelemetry.py

@@ -4,6 +4,7 @@
 # ------------------------------------
 from __future__ import annotations
 from contextlib import contextmanager
+from contextvars import Token
 from typing import Optional, Dict, Sequence, cast, Callable, Iterator, TYPE_CHECKING
 
 from opentelemetry import context as otel_context_module, trace
@@ -221,23 +222,23 @@ def get_trace_context(cls) -> Dict[str, str]:
         return trace_context
 
     @classmethod
-    def _suppress_auto_http_instrumentation(cls) -> object:
+    def _suppress_auto_http_instrumentation(cls) -> Token:
         """Enabled automatic HTTP instrumentation suppression.
 
         Since azure-core already instruments HTTP calls, we need to suppress any automatic HTTP
         instrumentation provided by other libraries to prevent duplicate spans. This has no effect if no
         automatic HTTP instrumentation libraries are being used.
 
         :return: A token that can be used to detach the suppression key from the context
-        :rtype: object
+        :rtype: ~contextvars.Token
         """
         return otel_context_module.attach(otel_context_module.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True))
 
     @classmethod
-    def _detach_from_context(cls, token: object) -> None:
+    def _detach_from_context(cls, token: Token) -> None:
         """Detach a token from the context.
 
         :param token: The token to detach
-        :type token: object
+        :type token: ~contextvars.Token
         """
         otel_context_module.detach(token)

sdk/core/azure-core/tests/test_tracing_policy.py

@@ -15,6 +15,7 @@
 from azure.core.tracing._models import SpanKind
 from azure.core.tracing._abstract_span import HttpSpanMixin
 import pytest
+from opentelemetry.trace import format_span_id, format_trace_id
 from opentelemetry.instrumentation.requests import RequestsInstrumentor
 
 from utils import HTTP_RESPONSES, HTTP_REQUESTS, create_http_response, request_and_responses_product
@@ -40,7 +41,7 @@ def test_distributed_tracing_policy_solo(self, tracing_implementation, http_requ
             response.status_code = 202
             response.headers[policy._RESPONSE_ID] = "some request id"
 
-            assert request.headers.get("traceparent") == "123456789"
+            assert request.headers.get("traceparent") == "00-12345-GET-01"
 
             policy.on_response(pipeline_request, PipelineResponse(request, response, PipelineContext(None)))
             time.sleep(0.001)
@@ -197,7 +198,7 @@ def test_distributed_tracing_policy_with_user_agent(self, tracing_implementation
                 response.headers[policy._RESPONSE_ID] = "some request id"
                 pipeline_response = PipelineResponse(request, response, PipelineContext(None))
 
-                assert request.headers.get("traceparent") == "123456789"
+                assert request.headers.get("traceparent") == "00-12345-GET-01"
 
                 policy.on_response(pipeline_request, pipeline_response)
 
@@ -381,6 +382,10 @@ def test_distributed_tracing_policy(self, tracing_helper, http_request, http_res
         assert finished_spans[0].name == "GET"
         assert finished_spans[0].parent is root_span.get_span_context()
 
+        span_context = finished_spans[0].get_span_context()
+        assert traceparent.split("-")[1] == format_trace_id(span_context.trace_id)
+        assert traceparent.split("-")[2] == format_span_id(span_context.span_id)
+
         assert finished_spans[0].attributes.get(policy._HTTP_REQUEST_METHOD) == "GET"
         assert finished_spans[0].attributes.get(policy._URL_FULL) == "http://localhost/temp?query=query"
         assert finished_spans[0].attributes.get(policy._SERVER_ADDRESS) == "localhost"
@@ -498,6 +503,10 @@ def test_distributed_tracing_policy_with_user_agent_policy(self, tracing_helper,
         assert finished_spans[0].name == "GET"
         assert finished_spans[0].parent is root_span.get_span_context()
 
+        span_context = finished_spans[0].get_span_context()
+        assert traceparent.split("-")[1] == format_trace_id(span_context.trace_id)
+        assert traceparent.split("-")[2] == format_span_id(span_context.span_id)
+
         assert finished_spans[0].attributes.get(policy._HTTP_REQUEST_METHOD) == "GET"
         assert finished_spans[0].attributes.get(policy._URL_FULL) == "http://localhost/temp?query=query"
         assert finished_spans[0].attributes.get(policy._SERVER_ADDRESS) == "localhost"

sdk/core/azure-core/tests/tracing_common.py

@@ -35,14 +35,12 @@ def __init__(self, span=None, name="span", kind=None):
         :param name: The name of the OpenTelemetry span to create if a new span is needed
         :type name: str
         """
-        self._span = span
         self.name = name
         self._kind = kind or SpanKind.UNSPECIFIED
         self.attributes = {}
         self.children = []
         if self.CONTEXT:
             self.CONTEXT[-1].children.append(self)
-        self.CONTEXT.append(self)
         self.status = None
 
     def __str__(self):
@@ -58,7 +56,7 @@ def span_instance(self):
         """
         :return: The OpenTelemetry span that is being wrapped.
         """
-        return self._span
+        return self
 
     def span(self, name="span"):
         # type: (Optional[str]) -> OpenCensusSpan
@@ -84,13 +82,15 @@ def kind(self, value):
 
     def __enter__(self):
         """Start a span."""
+        self.CONTEXT.append(self)
         return self
 
     def __exit__(self, exception_type, exception_value, traceback):
         """Finish a span."""
         if exception_value:
             self.status = exception_value.args[0]
-        self.CONTEXT.pop()
+        if self.CONTEXT and self.CONTEXT[-1] == self:
+            self.CONTEXT.pop()
 
     def start(self):
         # type: () -> None
@@ -100,15 +100,17 @@ def start(self):
     def finish(self):
         # type: () -> None
         """Set the end time for a span."""
-        self.CONTEXT.pop()
+        if self.CONTEXT and self.CONTEXT[-1] == self:
+            self.CONTEXT.pop()
 
     def to_header(self):
         # type: () -> Dict[str, str]
         """
         Returns a dictionary with the header labels and values.
         :return: A key value pair dictionary
         """
-        return {"traceparent": "123456789"}
+        current_span = self.get_current_span()
+        return {"traceparent": f"00-12345-{current_span.name}-01"}
 
     def add_attribute(self, key, value):
         # type: (str, Union[str, int]) -> None