diff --git a/ddtrace/propagation/http.py b/ddtrace/propagation/http.py index 343f653f548..a1664664ace 100644 --- a/ddtrace/propagation/http.py +++ b/ddtrace/propagation/http.py @@ -1,3 +1,4 @@ +import itertools import re import sys from typing import Any # noqa:F401 @@ -912,21 +913,23 @@ def _inject(span_context: Context, headers: Dict[str, str]) -> None: 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 - + if len(baggage_items) > DD_TRACE_BAGGAGE_MAX_ITEMS: + log.warning("Baggage item limit exceeded, dropping excess items") + baggage_items = itertools.islice(baggage_items, DD_TRACE_BAGGAGE_MAX_ITEMS) # type: ignore + + encoded_items: List[str] = [] + total_size = 0 + for key, value in baggage_items: + item = f"{_BaggageHeader._encode_key(key)}={_BaggageHeader._encode_value(value)}" + item_size = len(item.encode("utf-8")) + (1 if encoded_items else 0) # +1 for comma if not first item + if total_size + item_size > DD_TRACE_BAGGAGE_MAX_BYTES: + log.warning("Baggage header size exceeded, dropping excess items") + break # stop adding items when size limit is reached + encoded_items.append(item) + total_size += item_size + + header_value = ",".join(encoded_items) headers[_HTTP_HEADER_BAGGAGE] = header_value except Exception: diff --git a/tests/tracer/test_propagation.py b/tests/tracer/test_propagation.py index e4bf9be6ef3..64ce8232524 100644 --- a/tests/tracer/test_propagation.py +++ b/tests/tracer/test_propagation.py @@ -38,6 +38,7 @@ from ddtrace.propagation.http import HTTP_HEADER_SAMPLING_PRIORITY from ddtrace.propagation.http import HTTP_HEADER_TRACE_ID from ddtrace.propagation.http import HTTPPropagator +from ddtrace.propagation.http import _BaggageHeader from ddtrace.propagation.http import _TraceContext from tests.contrib.fastapi.conftest import client as fastapi_client # noqa:F401 from tests.contrib.fastapi.conftest import fastapi_application # noqa:F401 @@ -3134,16 +3135,15 @@ def test_llmobs_parent_id_not_injected_by_default(): ], ) 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(): + import urllib.parse + from ddtrace.internal.constants import DD_TRACE_BAGGAGE_MAX_ITEMS - from ddtrace.propagation.http import _BaggageHeader headers = {} baggage_items = {} @@ -3151,18 +3151,45 @@ def test_baggageheader_maxitems_inject(): baggage_items[f"key{i}"] = f"val{i}" span_context = Context(baggage=baggage_items) _BaggageHeader._inject(span_context, headers) - assert "baggage" not in headers + assert "baggage" in headers + header_value = headers["baggage"] + items = header_value.split(",") + assert len(items) == DD_TRACE_BAGGAGE_MAX_ITEMS + + expected_keys = [f"key{i}" for i in range(DD_TRACE_BAGGAGE_MAX_ITEMS)] + for item in items: + key, value = item.split("=", 1) + key = urllib.parse.unquote(key) + assert key in expected_keys 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)} + # baggage item that exceeds the maximum byte size + baggage_items = {"foo": "a" * (DD_TRACE_BAGGAGE_MAX_BYTES + 1)} + span_context = Context(baggage=baggage_items) + _BaggageHeader._inject(span_context, headers) + # since the baggage item exceeds the max bytes, no header should be injected + header_value = headers["baggage"] + assert header_value == "" + + # multiple baggage items to test dropping items when the total size exceeds the limit + headers = {} + baggage_items = { + "key1": "a" * ((DD_TRACE_BAGGAGE_MAX_BYTES // 3)), + "key2": "b" * ((DD_TRACE_BAGGAGE_MAX_BYTES // 3)), + "key3": "c" * ((DD_TRACE_BAGGAGE_MAX_BYTES // 3)), + "key4": "d", + } span_context = Context(baggage=baggage_items) _BaggageHeader._inject(span_context, headers) - assert "baggage" not in headers + header_value = headers["baggage"] + header_size = len(header_value.encode("utf-8")) + assert header_size <= DD_TRACE_BAGGAGE_MAX_BYTES + assert "key4" not in header_value + assert "key2" in header_value @pytest.mark.parametrize( @@ -3188,8 +3215,6 @@ def test_baggageheader_maxbytes_inject(): ], ) def test_baggageheader_extract(headers, expected_baggage): - from ddtrace.propagation.http import _BaggageHeader - context = _BaggageHeader._extract(headers) assert context._baggage == expected_baggage @@ -3210,8 +3235,6 @@ def test_baggageheader_extract(headers, expected_baggage): ], ) def test_baggage_malformedheader_extract(headers, expected_baggage): - from ddtrace.propagation.http import _BaggageHeader - context = _BaggageHeader._extract(headers) assert context._baggage == expected_baggage