From 2b410f95802b75528b3fb25023a942a5495b7f4a Mon Sep 17 00:00:00 2001 From: Rachel Yang Date: Wed, 23 Oct 2024 16:02:43 -0400 Subject: [PATCH] feat: support baggage (#10389) First PR introducing baggage support ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) --------- Co-authored-by: Emmett Butler <723615+emmettbutler@users.noreply.github.com> Co-authored-by: Zachary Groves <32471391+ZStriker19@users.noreply.github.com> Co-authored-by: Munir Abdinur --- ddtrace/_trace/context.py | 17 +- ddtrace/_trace/span.py | 11 - ddtrace/internal/constants.py | 8 +- ddtrace/llmobs/_llmobs.py | 10 +- ddtrace/propagation/http.py | 111 +++++++++- ddtrace/settings/config.py | 3 +- .../baggage-support-be7eed26293f1216.yaml | 3 + tests/tracer/test_propagation.py | 193 +++++++++++++++++- tests/tracer/test_span.py | 44 +++- 9 files changed, 357 insertions(+), 43 deletions(-) create mode 100644 releasenotes/notes/baggage-support-be7eed26293f1216.yaml diff --git a/ddtrace/_trace/context.py b/ddtrace/_trace/context.py index 4bc36527338..07bc3960b56 100644 --- a/ddtrace/_trace/context.py +++ b/ddtrace/_trace/context.py @@ -218,7 +218,7 @@ def _trace_id_64bits(self): else: return _MAX_UINT_64BITS & self.trace_id - def _set_baggage_item(self, key: str, value: Any) -> None: + def set_baggage_item(self, key: str, value: Any) -> None: """Sets a baggage item in this span context. Note that this operation mutates the baggage of this span context """ @@ -237,10 +237,23 @@ def _with_baggage_item(self, key: str, value: Any) -> "Context": ctx._baggage = new_baggage return ctx - def _get_baggage_item(self, key: str) -> Optional[Any]: + 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: diff --git a/ddtrace/_trace/span.py b/ddtrace/_trace/span.py index ae642b97a19..db90a769cd9 100644 --- a/ddtrace/_trace/span.py +++ b/ddtrace/_trace/span.py @@ -497,23 +497,12 @@ def get_metric(self, key: _TagNameType) -> Optional[NumericType]: """Return the given metric or None if it doesn't exist.""" return self._metrics.get(key) - def _set_baggage_item(self, key: str, value: Any) -> "Span": - """Sets a baggage item in the span context of this span. - Baggage is used to propagate state between spans (in-process, http/https). - """ - self._context = self.context._with_baggage_item(key, value) - return self - def _add_event( self, name: str, attributes: Optional[Dict[str, str]] = None, timestamp: Optional[int] = None ) -> None: """Add an event to the span.""" self._events.append(SpanEvent(name, attributes, timestamp)) - 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_metrics(self) -> _MetricDictType: """Return all metrics.""" return self._metrics.copy() diff --git a/ddtrace/internal/constants.py b/ddtrace/internal/constants.py index e20f73e852a..bb55e657204 100644 --- a/ddtrace/internal/constants.py +++ b/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" @@ -67,6 +69,10 @@ _HTTPLIB_NO_TRACE_REQUEST = "_dd_no_trace" DEFAULT_TIMEOUT = 2.0 +# baggage +DD_TRACE_BAGGAGE_MAX_ITEMS = 64 +DD_TRACE_BAGGAGE_MAX_BYTES = 8192 + class _PRIORITY_CATEGORY: USER = "user" diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 42e27d33fb4..14ae70c1214 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -112,7 +112,7 @@ def _do_annotations(self, span): if span.span_type != SpanTypes.LLM: # do this check to avoid the warning log in `annotate` return current_context = self._instance.tracer.current_trace_context() - current_context_id = current_context._get_baggage_item(ANNOTATIONS_CONTEXT_ID) + current_context_id = current_context.get_baggage_item(ANNOTATIONS_CONTEXT_ID) with self._annotation_context_lock: for _, context_id, annotation_kwargs in self._instance._annotations: if current_context_id == context_id: @@ -301,12 +301,12 @@ def get_annotations_context_id(): ctx_id = annotation_id if current_ctx is None: current_ctx = Context(is_remote=False) - current_ctx._set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id) + current_ctx.set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id) cls._instance.tracer.context_provider.activate(current_ctx) - elif not current_ctx._get_baggage_item(ANNOTATIONS_CONTEXT_ID): - current_ctx._set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id) + elif not current_ctx.get_baggage_item(ANNOTATIONS_CONTEXT_ID): + current_ctx.set_baggage_item(ANNOTATIONS_CONTEXT_ID, ctx_id) else: - ctx_id = current_ctx._get_baggage_item(ANNOTATIONS_CONTEXT_ID) + ctx_id = current_ctx.get_baggage_item(ANNOTATIONS_CONTEXT_ID) return ctx_id def register_annotation(): diff --git a/ddtrace/propagation/http.py b/ddtrace/propagation/http.py index abbb8e46c1a..343f653f548 100644 --- a/ddtrace/propagation/http.py +++ b/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,8 +39,11 @@ 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 DD_TRACE_BAGGAGE_MAX_BYTES +from ..internal.constants import DD_TRACE_BAGGAGE_MAX_ITEMS from ..internal.constants import HIGHER_ORDER_TRACE_ID_BITS as _HIGHER_ORDER_TRACE_ID_BITS from ..internal.constants import LAST_DD_PARENT_ID_KEY from ..internal.constants import MAX_UINT_64BITS as _MAX_UINT_64BITS @@ -74,6 +78,7 @@ _HTTP_HEADER_TAGS: Literal["x-datadog-tags"] = "x-datadog-tags" _HTTP_HEADER_TRACEPARENT: Literal["traceparent"] = "traceparent" _HTTP_HEADER_TRACESTATE: Literal["tracestate"] = "tracestate" +_HTTP_HEADER_BAGGAGE: Literal["baggage"] = "baggage" def _possible_header(header): @@ -127,7 +132,7 @@ def _attach_baggage_to_context(headers: Dict[str, str], context: Context): if context is not None: for key, value in headers.items(): if key[: len(_HTTP_BAGGAGE_PREFIX)] == _HTTP_BAGGAGE_PREFIX: - context._set_baggage_item(key[len(_HTTP_BAGGAGE_PREFIX) :], value) + context.set_baggage_item(key[len(_HTTP_BAGGAGE_PREFIX) :], value) def _hex_id_to_dd_id(hex_id): @@ -885,12 +890,77 @@ def _inject(span_context, headers): return headers +class _BaggageHeader: + """Helper class to inject/extract Baggage Headers""" + + SAFE_CHARACTERS_KEY = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" "0123456789" "!#$%&'*+-.^_`|~" + SAFE_CHARACTERS_VALUE = ( + "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" "0123456789" "!#$%&'()*+-./:<>?@[]^_`{|}~" + ) + + @staticmethod + def _encode_key(key: str) -> str: + return urllib.parse.quote(str(key).strip(), safe=_BaggageHeader.SAFE_CHARACTERS_KEY) + + @staticmethod + def _encode_value(value: str) -> str: + return urllib.parse.quote(str(value).strip(), safe=_BaggageHeader.SAFE_CHARACTERS_VALUE) + + @staticmethod + def _inject(span_context: Context, headers: Dict[str, str]) -> None: + baggage_items = span_context._baggage.items() + if not baggage_items: + return + + if len(baggage_items) > DD_TRACE_BAGGAGE_MAX_ITEMS: + log.warning("Baggage item limit exceeded") + return + + try: + header_value = ",".join( + f"{_BaggageHeader._encode_key(key)}={_BaggageHeader._encode_value(value)}" + for key, value in baggage_items + ) + + buf = bytes(header_value, "utf-8") + if len(buf) > DD_TRACE_BAGGAGE_MAX_BYTES: + log.warning("Baggage header size exceeded") + return + + headers[_HTTP_HEADER_BAGGAGE] = header_value + + except Exception: + log.warning("Failed to encode and inject baggage header") + + @staticmethod + def _extract(headers: Dict[str, str]) -> Context: + header_value = headers.get(_HTTP_HEADER_BAGGAGE) + + if not header_value: + return Context(baggage={}) + + baggage = {} + baggages = header_value.split(",") + for key_value in baggages: + if "=" not in key_value: + return Context(baggage={}) + key, value = key_value.split("=", 1) + key = urllib.parse.unquote(key.strip()) + value = urllib.parse.unquote(value.strip()) + if not key or not value: + return Context(baggage={}) + 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, } @@ -906,6 +976,9 @@ def _extract_configured_contexts_avail(normalized_headers): for prop_style in config._propagation_style_extract: propagator = _PROP_STYLES[prop_style] context = propagator._extract(normalized_headers) + # baggage is handled separately + if prop_style == _PROPAGATION_STYLE_BAGGAGE: + continue if context: contexts.append(context) styles_w_ctx.append(prop_style) @@ -915,6 +988,7 @@ 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)] # encoding expects at least trace_id and span_id @@ -924,9 +998,11 @@ def _resolve_contexts(contexts, styles_w_ctx, normalized_headers): context.trace_id, context.span_id, flags=1 if context.sampling_priority and context.sampling_priority > 0 else 0, - tracestate=context._meta.get(W3C_TRACESTATE_KEY, "") - if style_w_ctx == _PROPAGATION_STYLE_W3C_TRACECONTEXT - else None, + tracestate=( + context._meta.get(W3C_TRACESTATE_KEY, "") + if style_w_ctx == _PROPAGATION_STYLE_W3C_TRACECONTEXT + else None + ), attributes={ "reason": "terminated_context", "context_headers": style_w_ctx, @@ -952,6 +1028,7 @@ 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 + primary_context._span_links = links return primary_context @@ -999,6 +1076,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) @@ -1024,7 +1105,6 @@ def parent_call(): @staticmethod def extract(headers): - # type: (Dict[str,str]) -> Context """Extract a Context from HTTP headers into a new Context. For tracecontext propagation we extract tracestate headers for propagation even if another propagation style is specified before tracecontext, @@ -1050,16 +1130,17 @@ def my_controller(url, headers): return Context() try: normalized_headers = {name.lower(): v for name, v in headers.items()} - + context = Context() # tracer configured to extract first only if config._propagation_extract_first: # loop through the extract propagation styles specified in order, return whatever context we get first for prop_style in config._propagation_style_extract: propagator = _PROP_STYLES[prop_style] - context = propagator._extract(normalized_headers) # type: ignore - if config._propagation_http_baggage_enabled is True: + context = propagator._extract(normalized_headers) + if config.propagation_http_baggage_enabled is True: _attach_baggage_to_context(normalized_headers, context) - return context + break + # loop through all extract propagation styles else: contexts, styles_w_ctx = HTTPPropagator._extract_configured_contexts_avail(normalized_headers) @@ -1068,7 +1149,17 @@ def my_controller(url, headers): context = HTTPPropagator._resolve_contexts(contexts, styles_w_ctx, normalized_headers) if config._propagation_http_baggage_enabled is True: _attach_baggage_to_context(normalized_headers, context) - return context + + # baggage headers are handled separately from the other propagation styles + if _PROPAGATION_STYLE_BAGGAGE in config._propagation_style_extract: + baggage_context = _BaggageHeader._extract(normalized_headers) + if baggage_context._baggage != {}: + if context: + context._baggage = baggage_context._baggage + else: + context = baggage_context + + return context except Exception: log.debug("error while extracting context propagation headers", exc_info=True) diff --git a/ddtrace/settings/config.py b/ddtrace/settings/config.py index 178145601a1..33105142b0b 100644 --- a/ddtrace/settings/config.py +++ b/ddtrace/settings/config.py @@ -112,10 +112,11 @@ def _parse_propagation_styles(styles_str): - "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:: diff --git a/releasenotes/notes/baggage-support-be7eed26293f1216.yaml b/releasenotes/notes/baggage-support-be7eed26293f1216.yaml new file mode 100644 index 00000000000..95f44846b13 --- /dev/null +++ b/releasenotes/notes/baggage-support-be7eed26293f1216.yaml @@ -0,0 +1,3 @@ +features: + - | + tracing: Introduces support for Baggage as defined by the `OpenTelemetry specification `_. \ No newline at end of file diff --git a/tests/tracer/test_propagation.py b/tests/tracer/test_propagation.py index 62f93ff2d04..d47fef3652d 100644 --- a/tests/tracer/test_propagation.py +++ b/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 @@ -28,6 +29,7 @@ from ddtrace.propagation.http import _HTTP_HEADER_B3_SINGLE from ddtrace.propagation.http import _HTTP_HEADER_B3_SPAN_ID from ddtrace.propagation.http import _HTTP_HEADER_B3_TRACE_ID +from ddtrace.propagation.http import _HTTP_HEADER_BAGGAGE from ddtrace.propagation.http import _HTTP_HEADER_TAGS from ddtrace.propagation.http import _HTTP_HEADER_TRACEPARENT from ddtrace.propagation.http import _HTTP_HEADER_TRACESTATE @@ -68,7 +70,7 @@ def test_inject(tracer): # noqa: F811 def test_inject_with_baggage_http_propagation(tracer): # noqa: F811 with override_global_config(dict(_propagation_http_baggage_enabled=True)): ctx = Context(trace_id=1234, sampling_priority=2, dd_origin="synthetics") - ctx._set_baggage_item("key1", "val1") + ctx.set_baggage_item("key1", "val1") tracer.context_provider.activate(ctx) with tracer.trace("global_root_span") as span: headers = {} @@ -282,6 +284,7 @@ def test_extract(tracer): # noqa: F811 "x-datadog-origin": "synthetics", "x-datadog-tags": "_dd.p.test=value,any=tag", "ot-baggage-key1": "value1", + "baggage": "foo=bar,racoon=cute,serverNode=DF%2028", } context = HTTPPropagator.extract(headers) @@ -308,6 +311,10 @@ def test_extract(tracer): # noqa: F811 "_dd.p.dm": "-3", "_dd.p.test": "value", } + assert context.get_baggage_item("foo") == "bar" + assert context.get_baggage_item("racoon") == "cute" + assert context.get_baggage_item("serverNode") == "DF 28" + assert len(context.get_all_baggage_items()) == 3 def test_asm_standalone_minimum_trace_per_minute_has_no_downstream_propagation(tracer): # noqa: F811 @@ -649,9 +656,9 @@ def test_extract_with_baggage_http_propagation(tracer): # noqa: F811 tracer.context_provider.activate(context) with tracer.trace("local_root_span") as span: - assert span._get_baggage_item("key1") == "value1" + assert span.context.get_baggage_item("key1") == "value1" with tracer.trace("child_span") as child_span: - assert child_span._get_baggage_item("key1") == "value1" + assert child_span.context.get_baggage_item("key1") == "value1" @pytest.mark.subprocess( @@ -2839,8 +2846,15 @@ def test_span_links_set_on_root_span_not_child(fastapi_client, tracer, fastapi_t PROPAGATION_STYLE_B3_MULTI, PROPAGATION_STYLE_B3_SINGLE, _PROPAGATION_STYLE_W3C_TRACECONTEXT, + _PROPAGATION_STYLE_BAGGAGE, ], - VALID_DATADOG_CONTEXT, + { + "trace_id": 13088165645273925489, + "span_id": 8185124618007618416, + "sampling_priority": 1, + "dd_origin": "synthetics", + "baggage": {"foo": "bar"}, + }, { HTTP_HEADER_TRACE_ID: "13088165645273925489", HTTP_HEADER_PARENT_ID: "8185124618007618416", @@ -2852,6 +2866,7 @@ def test_span_links_set_on_root_span_not_child(fastapi_client, tracer, fastapi_t _HTTP_HEADER_B3_SINGLE: "b5a2814f70060771-7197677932a62370-1", _HTTP_HEADER_TRACEPARENT: "00-0000000000000000b5a2814f70060771-7197677932a62370-01", _HTTP_HEADER_TRACESTATE: "dd=s:1;o:synthetics", + _HTTP_HEADER_BAGGAGE: "foo=bar", }, ), ( @@ -2918,6 +2933,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"}, + }, + { + _HTTP_HEADER_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", + _HTTP_HEADER_BAGGAGE: "foo=bar", + }, + ), + ( + "baggage_order_first", + [ + _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"], + }, + { + _HTTP_HEADER_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: "", + }, + ), ] @@ -3040,3 +3109,119 @@ 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"}), + ( + Context(baggage={'",;\\()/:<=>?@[]{}': '",;\\'}), + {"baggage": "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C"}, + ), + ], + ids=[ + "single_key_value", + "multiple_key_value_pairs", + "space_in_value", + "special_characters_in_value", + "special_characters_in_key", + "special_characters_in_key_and_value", + ], +) +def test_baggageheader_inject(span_context, expected_headers): + from ddtrace.propagation.http import _BaggageHeader + + headers = {} + _BaggageHeader._inject(span_context, headers) + assert headers == expected_headers + + +def test_baggageheader_maxitems_inject(): + from ddtrace.internal.constants import DD_TRACE_BAGGAGE_MAX_ITEMS + from ddtrace.propagation.http import _BaggageHeader + + headers = {} + baggage_items = {} + for i in range(DD_TRACE_BAGGAGE_MAX_ITEMS + 1): + baggage_items[f"key{i}"] = f"val{i}" + span_context = Context(baggage=baggage_items) + _BaggageHeader._inject(span_context, headers) + assert "baggage" not in headers + + +def test_baggageheader_maxbytes_inject(): + from ddtrace.internal.constants import DD_TRACE_BAGGAGE_MAX_BYTES + from ddtrace.propagation.http import _BaggageHeader + + headers = {} + baggage_items = {"foo": ("a" * DD_TRACE_BAGGAGE_MAX_BYTES)} + span_context = Context(baggage=baggage_items) + _BaggageHeader._inject(span_context, headers) + assert "baggage" not in 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"}), + ( + {"baggage": "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C"}, + {'",;\\()/:<=>?@[]{}': '",;\\'}, + ), + ], + ids=[ + "single_key_value", + "multiple_key_value_pairs", + "special_characters_in_key", + "special_characters_in_value", + "space_in_value", + "special_characters_in_key_and_value", + ], +) +def test_baggageheader_extract(headers, expected_baggage): + from ddtrace.propagation.http import _BaggageHeader + + context = _BaggageHeader._extract(headers) + assert context._baggage == expected_baggage + + +@pytest.mark.parametrize( + "headers,expected_baggage", + [ + ({"baggage": "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed"}, {}), + ({"baggage": "foo=gets-dropped-because-subsequent-pair-is-malformed,="}, {}), + ({"baggage": "=no-key"}, {}), + ({"baggage": "no-value="}, {}), + ], + ids=[ + "no-equal-sign-prev", + "no-equal-sign-subsequent", + "no-key", + "no-value", + ], +) +def test_baggage_malformedheader_extract(headers, expected_baggage): + from ddtrace.propagation.http import _BaggageHeader + + context = _BaggageHeader._extract(headers) + assert context._baggage == expected_baggage + + +@pytest.mark.parametrize( + "headers", + [ + {"baggage": "key1=val1,key2=val2,foo=bar,x=y"}, + ], +) +def test_http_propagator_baggage_extract(headers): + context = HTTPPropagator.extract(headers) + assert context._baggage == {"key1": "val1", "key2": "val2", "foo": "bar", "x": "y"} diff --git a/tests/tracer/test_span.py b/tests/tracer/test_span.py index fd8f995f569..e746632037e 100644 --- a/tests/tracer/test_span.py +++ b/tests/tracer/test_span.py @@ -98,20 +98,46 @@ def test_set_tag_bool(self): def test_set_baggage_item(self): s = Span(name="test.span") - s._set_baggage_item("custom.key", "123") - assert s._get_baggage_item("custom.key") == "123" + s.context.set_baggage_item("custom.key", "123") + assert s.context.get_baggage_item("custom.key") == "123" - def test_baggage_propagation(self): + def test_baggage_get(self): span1 = Span(name="test.span1") - span1._set_baggage_item("item1", "123") + span1.context.set_baggage_item("item1", "123") span2 = Span(name="test.span2", context=span1.context) - span2._set_baggage_item("item2", "456") + span2.context.set_baggage_item("item2", "456") - assert span2._get_baggage_item("item1") == "123" - assert span2._get_baggage_item("item2") == "456" - assert span1._get_baggage_item("item1") == "123" - assert span1._get_baggage_item("item2") is None + assert span2.context.get_baggage_item("item1") == "123" + assert span2.context.get_baggage_item("item2") == "456" + assert span1.context.get_baggage_item("item1") == "123" + + def test_baggage_remove(self): + span1 = Span(name="test.span1") + span1.context.set_baggage_item("item1", "123") + span1.context.set_baggage_item("item2", "456") + + span1.context.remove_baggage_item("item1") + assert span1.context.get_baggage_item("item1") is None + assert span1.context.get_baggage_item("item2") == "456" + span1.context.remove_baggage_item("item2") + assert span1.context.get_baggage_item("item2") is None + + def test_baggage_remove_all(self): + span1 = Span(name="test.span1") + span1.context.set_baggage_item("item1", "123") + span1.context.set_baggage_item("item2", "456") + + span1.context.remove_all_baggage_items() + assert span1.context.get_baggage_item("item1") is None + assert span1.context.get_baggage_item("item2") is None + + def test_baggage_get_all(self): + span1 = Span(name="test.span1") + span1.context.set_baggage_item("item1", "123") + span1.context.set_baggage_item("item2", "456") + + assert span1.context.get_all_baggage_items() == {"item1": "123", "item2": "456"} def test_set_tag_metric(self): s = Span(name="test.span")