Skip to content

Commit

Permalink
Merge branch 'main' into kylev/flaky-tests-should-be-shot-into-the-sun
Browse files Browse the repository at this point in the history
  • Loading branch information
Kyle-Verhoog authored Nov 27, 2024
2 parents 92ee3b7 + 21f43e5 commit 9e13565
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 27 deletions.
4 changes: 2 additions & 2 deletions ddtrace/_trace/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

def record_span_pointer_calculation(context: str, span_pointer_count: int) -> None:
telemetry_writer.add_count_metric(
namespace="tracer",
namespace="tracers",
name="span_pointer_calculation",
value=1,
tags=(("context", context), ("count", _span_pointer_count_to_tag(span_pointer_count))),
Expand Down Expand Up @@ -45,7 +45,7 @@ def record_span_pointer_calculation_issue(
tags += additional_tags

telemetry_writer.add_count_metric(
namespace="tracer",
namespace="tracers",
name="span_pointer_calculation.issue",
value=1,
tags=tags,
Expand Down
31 changes: 17 additions & 14 deletions ddtrace/propagation/http.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import itertools
import re
import sys
from typing import Any # noqa:F401
Expand Down Expand Up @@ -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:
Expand Down
45 changes: 34 additions & 11 deletions tests/tracer/test_propagation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -3163,35 +3164,61 @@ 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 = {}
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
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(
Expand All @@ -3217,8 +3244,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

Expand All @@ -3239,8 +3264,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

Expand Down

0 comments on commit 9e13565

Please sign in to comment.