WIP: otel baggage support initial PR

ddtrace/_trace/context.py

@@ -241,6 +241,19 @@ def _get_baggage_item(self, key: str) -> Optional[Any]:
         """Gets a baggage item in this span context."""
         return self._baggage.get(key, None)
 
+    def _get_all_baggage_items(self) -> Dict[str, Any]:
+        """Returns all baggage items in this span context."""
+        return self._baggage
+
+    def _remove_baggage_item(self, key: str) -> None:
+        """Remove a baggage item from this span context."""
+        if key in self._baggage:
+            del self._baggage[key]
+
+    def _remove_all_baggage_items(self) -> None:
+        """Removes all baggage items from this span context."""
+        self._baggage.clear()
+
     def __eq__(self, other: Any) -> bool:
         if isinstance(other, Context):
             with self._lock:

ddtrace/_trace/span.py

@@ -1,3 +1,4 @@
+import json
 import math
 import pprint
 import sys
@@ -508,6 +509,15 @@ def _get_baggage_item(self, key: str) -> Optional[Any]:
         """Gets a baggage item from the span context of this span."""
         return self.context._get_baggage_item(key)
 
+    def _get_all_baggage_items(self) -> str:
+        return json.dumps(self.context._get_all_baggage_items())
+
+    def _remove_baggage_item(self, key: str) -> None:
+        self.context._remove_baggage_item(key)
+
+    def _remove_all_baggage_items(self) -> None:
+        self.context._remove_all_baggage_items()
+
     def get_metrics(self) -> _MetricDictType:
         """Return all metrics."""
         return self._metrics.copy()

ddtrace/internal/constants.py

@@ -9,13 +9,15 @@
 PROPAGATION_STYLE_B3_SINGLE = "b3"
 _PROPAGATION_STYLE_W3C_TRACECONTEXT = "tracecontext"
 _PROPAGATION_STYLE_NONE = "none"
-_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext"
+_PROPAGATION_STYLE_DEFAULT = "datadog,tracecontext,baggage"
+_PROPAGATION_STYLE_BAGGAGE = "baggage"
 PROPAGATION_STYLE_ALL = (
     _PROPAGATION_STYLE_W3C_TRACECONTEXT,
     PROPAGATION_STYLE_DATADOG,
     PROPAGATION_STYLE_B3_MULTI,
     PROPAGATION_STYLE_B3_SINGLE,
     _PROPAGATION_STYLE_NONE,
+    _PROPAGATION_STYLE_BAGGAGE,
 )
 W3C_TRACESTATE_KEY = "tracestate"
 W3C_TRACEPARENT_KEY = "traceparent"

ddtrace/internal/opentelemetry/context.py

@@ -21,6 +21,7 @@ def attach(self, otel_context):
         Datadog representation in the Global DDtrace Trace Context Provider.
         """
         # Inline opentelemetry imports to avoid circular imports.
+        from opentelemetry.baggage import get_baggage
         from opentelemetry.trace import Span as OtelSpan
         from opentelemetry.trace import get_current_span
 
@@ -30,18 +31,28 @@ def attach(self, otel_context):
         if otel_span:
             if isinstance(otel_span, Span):
                 self._ddcontext_provider.activate(otel_span._ddspan)
+                ddcontext = otel_span._ddspan.context
             elif isinstance(otel_span, OtelSpan):
                 trace_id, span_id, _, tf, ts, _ = otel_span.get_span_context()
                 trace_state = ts.to_header() if ts else None
                 ddcontext = _TraceContext._get_context(trace_id, span_id, tf, trace_state)
                 self._ddcontext_provider.activate(ddcontext)
             else:
+                ddcontext = None
                 log.error(
                     "Programming ERROR: ddtrace does not support activiting spans with the type: %s. Please open a "
                     "github issue at: https://github.com/Datadog/dd-trace-py and set DD_TRACE_OTEL_ENABLED=True.",
                     type(otel_span),
                 )
 
+            # get current open telemetry baggage and store it on the datadog context object
+            otel_baggage = get_baggage(otel_context)
+            if ddcontext and otel_baggage:
+                for key, value in otel_baggage.items():
+                    # value is from opentelemetry and can be any object.
+                    # make sure this object is compatible w/ datadog internals
+                    ddcontext._baggage[key] = value  # potentially convert to json
+
         # A return value with the type `object` is required by the otel api to remove/deactivate spans.
         # Since manually deactivating spans is not supported by ddtrace this object will never be used.
         return object()
@@ -52,6 +63,7 @@ def get_current(self):
         in a format that can be parsed by the OpenTelemetry API.
         """
         # Inline opentelemetry imports to avoid circular imports.
+        from opentelemetry.baggage import set_baggage
         from opentelemetry.context.context import Context as OtelContext
         from opentelemetry.trace import NonRecordingSpan as OtelNonRecordingSpan
         from opentelemetry.trace import SpanContext as OtelSpanContext
@@ -61,8 +73,9 @@ def get_current(self):
 
         from .span import Span
 
+        # getting current active span
         ddactive = self._ddcontext_provider.active()
-        context = OtelContext()
+        context = OtelContext()  # store current datadog baggage into here
         if isinstance(ddactive, DDSpan):
             span = Span(ddactive)
             context = set_span_in_context(span, context)
@@ -72,6 +85,18 @@ def get_current(self):
             span_context = OtelSpanContext(ddactive.trace_id or 0, ddactive.span_id or 0, True, tf, ts)
             span = OtelNonRecordingSpan(span_context)
             context = set_span_in_context(span, context)
+
+        if isinstance(ddactive, DDSpan):
+            dd_baggage = ddactive.context._baggage
+        elif isinstance(ddactive, DDContext):
+            dd_baggage = ddactive._baggage
+        else:
+            dd_baggage = {}
+
+        # getting current active baggage
+        for key, value in dd_baggage.items():
+            set_baggage(key, value, context)
+
         return context
 
     def detach(self, token):

ddtrace/propagation/http.py

@@ -8,6 +8,7 @@
 from typing import Text  # noqa:F401
 from typing import Tuple  # noqa:F401
 from typing import cast  # noqa:F401
+import urllib.parse
 
 import ddtrace
 from ddtrace._trace.span import Span  # noqa:F401
@@ -38,6 +39,7 @@
 from ..internal._tagset import decode_tagset_string
 from ..internal._tagset import encode_tagset_values
 from ..internal.compat import ensure_text
+from ..internal.constants import _PROPAGATION_STYLE_BAGGAGE
 from ..internal.constants import _PROPAGATION_STYLE_NONE
 from ..internal.constants import _PROPAGATION_STYLE_W3C_TRACECONTEXT
 from ..internal.constants import HIGHER_ORDER_TRACE_ID_BITS as _HIGHER_ORDER_TRACE_ID_BITS
@@ -885,12 +887,57 @@ def _inject(span_context, headers):
         return headers
 
 
+class _BaggageHeader:
+    @staticmethod
+    def _encode_key(key):
+        safe_characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" "0123456789" "!#$%&'*+-.^_`|~"
+        encoded = urllib.parse.quote(key, safe=safe_characters)
+        return encoded
+
+    def _encode_value(value):
+        safe_characters = (
+            "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" "0123456789" "!#$%&'()*+-./:<>?@[]^_`{|}~"
+        )
+        encoded = urllib.parse.quote(value, safe=safe_characters)
+        return encoded
+
+    @staticmethod
+    def _inject(span_context, headers):
+        baggage_items = span_context._baggage.items()
+        if not baggage_items:
+            return
+
+        header_value = ",".join(
+            f"{_BaggageHeader._encode_key(str(key).strip())}={_BaggageHeader._encode_value(str(value).strip())}"
+            for key, value in baggage_items
+        )
+
+        headers["baggage"] = header_value
+
+    @staticmethod
+    def _extract(headers):
+        header_value = headers.get("baggage")
+        if not header_value:
+            return None
+
+        baggage = {}
+        baggages = header_value.split(",")
+        for key_value in baggages:
+            key, value = key_value.split("=", 1)
+            key = urllib.parse.unquote(key.strip())
+            value = urllib.parse.unquote(value.strip())
+            baggage[key] = value
+
+        return Context(baggage=baggage)
+
+
 _PROP_STYLES = {
     PROPAGATION_STYLE_DATADOG: _DatadogMultiHeader,
     PROPAGATION_STYLE_B3_MULTI: _B3MultiHeader,
     PROPAGATION_STYLE_B3_SINGLE: _B3SingleHeader,
     _PROPAGATION_STYLE_W3C_TRACECONTEXT: _TraceContext,
     _PROPAGATION_STYLE_NONE: _NOP_Propagator,
+    _PROPAGATION_STYLE_BAGGAGE: _BaggageHeader,
 }
 
 
@@ -915,8 +962,19 @@ def _extract_configured_contexts_avail(normalized_headers):
     def _resolve_contexts(contexts, styles_w_ctx, normalized_headers):
         primary_context = contexts[0]
         links = []
+
         for context in contexts[1:]:
             style_w_ctx = styles_w_ctx[contexts.index(context)]
+
+            # special case where baggage is first in propagation styles
+            # why would you do that? I don't know, but it's possible
+            if style_w_ctx[0] == _PROPAGATION_STYLE_BAGGAGE:
+                baggage_context = contexts[0]
+                contexts.append(baggage_context)
+                del contexts[0]
+                styles_w_ctx.append(_PROPAGATION_STYLE_BAGGAGE)
+                del style_w_ctx[0]
+
             # encoding expects at least trace_id and span_id
             if context.span_id and context.trace_id and context.trace_id != primary_context.trace_id:
                 links.append(
@@ -952,6 +1010,11 @@ def _resolve_contexts(contexts, styles_w_ctx, normalized_headers):
                         primary_context._meta[LAST_DD_PARENT_ID_KEY] = "{:016x}".format(dd_context.span_id)
                     # the span_id in tracecontext takes precedence over the first extracted propagation style
                     primary_context.span_id = context.span_id
+
+            # baggage is always merged into the primary context
+            if style_w_ctx == _PROPAGATION_STYLE_BAGGAGE:
+                primary_context._baggage.update(context._baggage)
+
         primary_context._span_links = links
         return primary_context
 
@@ -998,6 +1061,10 @@ def parent_call():
         else:
             log.error("ddtrace.tracer.sample is not available, unable to sample span.")
 
+        # baggage should be injected regardless of existing span or trace id
+        if _PROPAGATION_STYLE_BAGGAGE in config._propagation_style_inject:
+            _BaggageHeader._inject(span_context, headers)
+
         # Not a valid context to propagate
         if span_context.trace_id is None or span_context.span_id is None:
             log.debug("tried to inject invalid context %r", span_context)

ddtrace/settings/config.py

@@ -110,10 +110,11 @@ def _parse_propagation_styles(name, default):
     - "b3" (formerly 'b3 single header')
     - "b3 single header (deprecated for 'b3')"
     - "tracecontext"
+    - "baggage"
     - "none"
 
 
-    The default value is ``"datadog,tracecontext"``.
+    The default value is ``"datadog,tracecontext,baggage"``.
 
 
     Examples::

tests/tracer/test_propagation.py

@@ -15,6 +15,7 @@
 from ddtrace.constants import AUTO_REJECT
 from ddtrace.constants import USER_KEEP
 from ddtrace.constants import USER_REJECT
+from ddtrace.internal.constants import _PROPAGATION_STYLE_BAGGAGE
 from ddtrace.internal.constants import _PROPAGATION_STYLE_NONE
 from ddtrace.internal.constants import _PROPAGATION_STYLE_W3C_TRACECONTEXT
 from ddtrace.internal.constants import LAST_DD_PARENT_ID_KEY
@@ -2915,6 +2916,60 @@ def test_span_links_set_on_root_span_not_child(fastapi_client, tracer, fastapi_t
             _HTTP_HEADER_TRACESTATE: "",
         },
     ),
+    (
+        "only_baggage",
+        [
+            _PROPAGATION_STYLE_BAGGAGE,
+        ],
+        {
+            "baggage": {"foo": "bar"},
+        },
+        {
+            "baggage": "foo=bar",
+        },
+    ),
+    (
+        "baggage_and_datadog",
+        [
+            PROPAGATION_STYLE_DATADOG,
+            _PROPAGATION_STYLE_BAGGAGE,
+        ],
+        {
+            "trace_id": VALID_DATADOG_CONTEXT["trace_id"],
+            "span_id": VALID_DATADOG_CONTEXT["span_id"],
+            "baggage": {"foo": "bar"},
+        },
+        {
+            HTTP_HEADER_TRACE_ID: "13088165645273925489",
+            HTTP_HEADER_PARENT_ID: "8185124618007618416",
+            "baggage": "foo=bar",
+        },
+    ),
+    (
+        "baggage_order",
+        [
+            _PROPAGATION_STYLE_BAGGAGE,
+            PROPAGATION_STYLE_DATADOG,
+            PROPAGATION_STYLE_B3_MULTI,
+            PROPAGATION_STYLE_B3_SINGLE,
+            _PROPAGATION_STYLE_W3C_TRACECONTEXT,
+        ],
+        {
+            "baggage": {"foo": "bar"},
+            "trace_id": VALID_DATADOG_CONTEXT["trace_id"],
+            "span_id": VALID_DATADOG_CONTEXT["span_id"],
+        },
+        {
+            "baggage": "foo=bar",
+            HTTP_HEADER_TRACE_ID: "13088165645273925489",
+            HTTP_HEADER_PARENT_ID: "8185124618007618416",
+            _HTTP_HEADER_B3_TRACE_ID: "b5a2814f70060771",
+            _HTTP_HEADER_B3_SPAN_ID: "7197677932a62370",
+            _HTTP_HEADER_B3_SINGLE: "b5a2814f70060771-7197677932a62370",
+            _HTTP_HEADER_TRACEPARENT: "00-0000000000000000b5a2814f70060771-7197677932a62370-00",
+            _HTTP_HEADER_TRACESTATE: "",
+        },
+    ),
 ]
 
 
@@ -3037,3 +3092,52 @@ def test_llmobs_parent_id_not_injected_by_default():
         context = Context(trace_id=1, span_id=2)
         HTTPPropagator.inject(context, {})
         mock_llmobs_inject.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "span_context,expected_headers",
+    [
+        (Context(baggage={"key1": "val1"}), {"baggage": "key1=val1"}),
+        (Context(baggage={"key1": "val1", "key2": "val2"}), {"baggage": "key1=val1,key2=val2"}),
+        (Context(baggage={"serverNode": "DF 28"}), {"baggage": "serverNode=DF%2028"}),
+        (Context(baggage={"userId": "Amélie"}), {"baggage": "userId=Am%C3%A9lie"}),
+        (Context(baggage={"user!d(me)": "false"}), {"baggage": "user!d%28me%29=false"}),
+    ],
+    ids=[
+        "single_key_value",
+        "multiple_key_value_pairs",
+        "space_in_value",
+        "special_characters_in_value",
+        "special_characters_in_key",
+    ],
+)
+def test_baggageheader_inject(span_context, expected_headers):
+    from ddtrace.propagation.http import _BaggageHeader
+
+    headers = {}
+    _BaggageHeader._inject(span_context, headers)
+    assert headers == expected_headers
+
+
+@pytest.mark.parametrize(
+    "headers,expected_baggage",
+    [
+        ({"baggage": "key1=val1"}, {"key1": "val1"}),
+        ({"baggage": "key1=val1,key2=val2,foo=bar,x=y"}, {"key1": "val1", "key2": "val2", "foo": "bar", "x": "y"}),
+        ({"baggage": "user!d%28me%29=false"}, {"user!d(me)": "false"}),
+        ({"baggage": "userId=Am%C3%A9lie"}, {"userId": "Amélie"}),
+        ({"baggage": "serverNode=DF%2028"}, {"serverNode": "DF 28"}),
+    ],
+    ids=[
+        "single_key_value",
+        "multiple_key_value_pairs",
+        "special_characters_in_key",
+        "special_characters_in_value",
+        "space_in_value",
+    ],
+)
+def test_baggageheader_extract(headers, expected_baggage):
+    from ddtrace.propagation.http import _BaggageHeader
+
+    context = _BaggageHeader._extract(headers)
+    assert context._baggage == expected_baggage