Baggage creation for head of trace (#1589)

Author: sl0thentr0py

sentry_sdk/hub.py

@@ -717,6 +717,19 @@ def iter_trace_propagation_headers(self, span=None):
         for header in span.iter_headers():
             yield header
 
+    def trace_propagation_meta(self, span=None):
+        # type: (Optional[Span]) -> str
+        """
+        Return meta tags which should be injected into the HTML template
+        to allow propagation of trace data.
+        """
+        meta = ""
+
+        for name, content in self.iter_trace_propagation_headers(span):
+            meta += '<meta name="%s" content="%s">' % (name, content)
+
+        return meta
+
 
 GLOBAL_HUB = Hub()
 _local.set(GLOBAL_HUB)

sentry_sdk/tracing.py

@@ -35,6 +35,11 @@
 TRANSACTION_SOURCE_COMPONENT = "component"
 TRANSACTION_SOURCE_TASK = "task"
 
+# These are typically high cardinality and the server hates them
+LOW_QUALITY_TRANSACTION_SOURCES = [
+    TRANSACTION_SOURCE_URL,
+]
+
 SOURCE_FOR_STYLE = {
     "endpoint": TRANSACTION_SOURCE_COMPONENT,
     "function_name": TRANSACTION_SOURCE_COMPONENT,
@@ -281,6 +286,10 @@ def continue_from_headers(
 
         if sentrytrace_kwargs is not None:
             kwargs.update(sentrytrace_kwargs)
+
+            # If there's an incoming sentry-trace but no incoming baggage header,
+            # for instance in traces coming from older SDKs,
+            # baggage will be empty and immutable and won't be populated as head SDK.
             baggage.freeze()
 
         kwargs.update(extract_tracestate_data(headers.get("tracestate")))
@@ -309,8 +318,8 @@ def iter_headers(self):
         if tracestate:
             yield "tracestate", tracestate
 
-        if self.containing_transaction and self.containing_transaction._baggage:
-            baggage = self.containing_transaction._baggage.serialize()
+        if self.containing_transaction:
+            baggage = self.containing_transaction.get_baggage().serialize()
             if baggage:
                 yield "baggage", baggage
 
@@ -513,11 +522,10 @@ def get_trace_context(self):
         if sentry_tracestate:
             rv["tracestate"] = sentry_tracestate
 
-        # TODO-neel populate fresh if head SDK
-        if self.containing_transaction and self.containing_transaction._baggage:
+        if self.containing_transaction:
             rv[
                 "dynamic_sampling_context"
-            ] = self.containing_transaction._baggage.dynamic_sampling_context()
+            ] = self.containing_transaction.get_baggage().dynamic_sampling_context()
 
         return rv
 
@@ -527,6 +535,8 @@ class Transaction(Span):
         "name",
         "source",
         "parent_sampled",
+        # used to create baggage value for head SDKs in dynamic sampling
+        "sample_rate",
         # the sentry portion of the `tracestate` header used to transmit
         # correlation context for server-side dynamic sampling, of the form
         # `sentry=xxxxx`, where `xxxxx` is the base64-encoded json of the
@@ -562,6 +572,7 @@ def __init__(
         Span.__init__(self, **kwargs)
         self.name = name
         self.source = source
+        self.sample_rate = None  # type: Optional[float]
         self.parent_sampled = parent_sampled
         # if tracestate isn't inherited and set here, it will get set lazily,
         # either the first time an outgoing request needs it for a header or the
@@ -570,7 +581,7 @@ def __init__(
         self._third_party_tracestate = third_party_tracestate
         self._measurements = {}  # type: Dict[str, Any]
         self._profile = None  # type: Optional[Sampler]
-        self._baggage = baggage
+        self._baggage = baggage  # type: Optional[Baggage]
 
     def __repr__(self):
         # type: () -> str
@@ -708,6 +719,17 @@ def to_json(self):
 
         return rv
 
+    def get_baggage(self):
+        # type: () -> Baggage
+        """
+        The first time a new baggage with sentry items is made,
+        it will be frozen.
+        """
+        if not self._baggage or self._baggage.mutable:
+            self._baggage = Baggage.populate_from_transaction(self)
+
+        return self._baggage
+
     def _set_initial_sampling_decision(self, sampling_context):
         # type: (SamplingContext) -> None
         """
@@ -745,6 +767,7 @@ def _set_initial_sampling_decision(self, sampling_context):
         # if the user has forced a sampling decision by passing a `sampled`
         # value when starting the transaction, go with that
         if self.sampled is not None:
+            self.sample_rate = float(self.sampled)
             return
 
         # we would have bailed already if neither `traces_sampler` nor
@@ -773,6 +796,8 @@ def _set_initial_sampling_decision(self, sampling_context):
             self.sampled = False
             return
 
+        self.sample_rate = float(sample_rate)
+
         # if the function returned 0 (or false), or if `traces_sample_rate` is
         # 0, it's a sign the transaction should be dropped
         if not sample_rate:

sentry_sdk/tracing_utils.py

@@ -470,6 +470,54 @@ def from_incoming_header(cls, header):
 
         return Baggage(sentry_items, third_party_items, mutable)
 
+    @classmethod
+    def populate_from_transaction(cls, transaction):
+        # type: (Transaction) -> Baggage
+        """
+        Populate fresh baggage entry with sentry_items and make it immutable
+        if this is the head SDK which originates traces.
+        """
+        hub = transaction.hub or sentry_sdk.Hub.current
+        client = hub.client
+        sentry_items = {}  # type: Dict[str, str]
+
+        if not client:
+            return Baggage(sentry_items)
+
+        options = client.options or {}
+        user = (hub.scope and hub.scope._user) or {}
+
+        sentry_items["trace_id"] = transaction.trace_id
+
+        if options.get("environment"):
+            sentry_items["environment"] = options["environment"]
+
+        if options.get("release"):
+            sentry_items["release"] = options["release"]
+
+        if options.get("dsn"):
+            sentry_items["public_key"] = Dsn(options["dsn"]).public_key
+
+        if (
+            transaction.name
+            and transaction.source not in LOW_QUALITY_TRANSACTION_SOURCES
+        ):
+            sentry_items["transaction"] = transaction.name
+
+        if user.get("segment"):
+            sentry_items["user_segment"] = user["segment"]
+
+        if transaction.sample_rate is not None:
+            sentry_items["sample_rate"] = str(transaction.sample_rate)
+
+        # there's an existing baggage but it was mutable,
+        # which is why we are creating this new baggage.
+        # However, if by chance the user put some sentry items in there, give them precedence.
+        if transaction._baggage and transaction._baggage.sentry_items:
+            sentry_items.update(transaction._baggage.sentry_items)
+
+        return Baggage(sentry_items, mutable=False)
+
     def freeze(self):
         # type: () -> None
         self.mutable = False
@@ -500,6 +548,7 @@ def serialize(self, include_third_party=False):
 
 
 # Circular imports
+from sentry_sdk.tracing import LOW_QUALITY_TRANSACTION_SOURCES
 
 if MYPY:
-    from sentry_sdk.tracing import Span
+    from sentry_sdk.tracing import Span, Transaction

tests/integrations/sqlalchemy/test_sqlalchemy.py

@@ -191,14 +191,6 @@ def processor(event, hint):
     # Some spans are discarded.
     assert len(event["spans"]) == 1000
 
-    # Some spans have their descriptions truncated. Because the test always
-    # generates the same amount of descriptions and truncation is deterministic,
-    # the number here should never change across test runs.
-    #
-    # Which exact span descriptions are truncated depends on the span durations
-    # of each SQL query and is non-deterministic.
-    assert len(event["_meta"]["spans"]) == 537
-
     for i, span in enumerate(event["spans"]):
         description = span["description"]
 

tests/integrations/stdlib/test_httplib.py

@@ -1,6 +1,6 @@
 import platform
 import sys
-
+import random
 import pytest
 
 try:
@@ -122,9 +122,7 @@ def test_httplib_misuse(sentry_init, capture_events, request):
     }
 
 
-def test_outgoing_trace_headers(
-    sentry_init, monkeypatch, StringContaining  # noqa: N803
-):
+def test_outgoing_trace_headers(sentry_init, monkeypatch):
     # HTTPSConnection.send is passed a string containing (among other things)
     # the headers on the request. Mock it so we can check the headers, and also
     # so it doesn't try to actually talk to the internet.
@@ -176,3 +174,46 @@ def test_outgoing_trace_headers(
         assert sorted(request_headers["baggage"].split(",")) == sorted(
             expected_outgoing_baggage_items
         )
+
+
+def test_outgoing_trace_headers_head_sdk(sentry_init, monkeypatch):
+    # HTTPSConnection.send is passed a string containing (among other things)
+    # the headers on the request. Mock it so we can check the headers, and also
+    # so it doesn't try to actually talk to the internet.
+    mock_send = mock.Mock()
+    monkeypatch.setattr(HTTPSConnection, "send", mock_send)
+
+    # make sure transaction is always sampled
+    monkeypatch.setattr(random, "random", lambda: 0.1)
+
+    sentry_init(traces_sample_rate=0.5, release="foo")
+    transaction = Transaction.continue_from_headers({})
+
+    with start_transaction(transaction=transaction, name="Head SDK tx") as transaction:
+        HTTPSConnection("www.squirrelchasers.com").request("GET", "/top-chasers")
+
+        (request_str,) = mock_send.call_args[0]
+        request_headers = {}
+        for line in request_str.decode("utf-8").split("\r\n")[1:]:
+            if line:
+                key, val = line.split(": ")
+                request_headers[key] = val
+
+        request_span = transaction._span_recorder.spans[-1]
+        expected_sentry_trace = "{trace_id}-{parent_span_id}-{sampled}".format(
+            trace_id=transaction.trace_id,
+            parent_span_id=request_span.span_id,
+            sampled=1,
+        )
+        assert request_headers["sentry-trace"] == expected_sentry_trace
+
+        expected_outgoing_baggage_items = [
+            "sentry-trace_id=%s" % transaction.trace_id,
+            "sentry-sample_rate=0.5",
+            "sentry-release=foo",
+            "sentry-environment=production",
+        ]
+
+        assert sorted(request_headers["baggage"].split(",")) == sorted(
+            expected_outgoing_baggage_items
+        )

tests/tracing/test_integration_tests.py

@@ -1,7 +1,9 @@
 # coding: utf-8
 import weakref
 import gc
+import re
 import pytest
+import random
 
 from sentry_sdk import (
     capture_message,
@@ -142,6 +144,61 @@ def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_r
     assert message_payload["message"] == "hello"
 
 
+@pytest.mark.parametrize("sample_rate", [0.5, 1.0])
+def test_dynamic_sampling_head_sdk_creates_dsc(
+    sentry_init, capture_envelopes, sample_rate, monkeypatch
+):
+    sentry_init(traces_sample_rate=sample_rate, release="foo")
+    envelopes = capture_envelopes()
+
+    # make sure transaction is sampled for both cases
+    monkeypatch.setattr(random, "random", lambda: 0.1)
+
+    transaction = Transaction.continue_from_headers({}, name="Head SDK tx")
+
+    # will create empty mutable baggage
+    baggage = transaction._baggage
+    assert baggage
+    assert baggage.mutable
+    assert baggage.sentry_items == {}
+    assert baggage.third_party_items == ""
+
+    with start_transaction(transaction):
+        with start_span(op="foo", description="foodesc"):
+            pass
+
+    # finish will create a new baggage entry
+    baggage = transaction._baggage
+    trace_id = transaction.trace_id
+
+    assert baggage
+    assert not baggage.mutable
+    assert baggage.third_party_items == ""
+    assert baggage.sentry_items == {
+        "environment": "production",
+        "release": "foo",
+        "sample_rate": str(sample_rate),
+        "transaction": "Head SDK tx",
+        "trace_id": trace_id,
+    }
+
+    expected_baggage = (
+        "sentry-environment=production,sentry-release=foo,sentry-sample_rate=%s,sentry-transaction=Head%%20SDK%%20tx,sentry-trace_id=%s"
+        % (sample_rate, trace_id)
+    )
+    assert sorted(baggage.serialize().split(",")) == sorted(expected_baggage.split(","))
+
+    (envelope,) = envelopes
+    assert envelope.headers["trace"] == baggage.dynamic_sampling_context()
+    assert envelope.headers["trace"] == {
+        "environment": "production",
+        "release": "foo",
+        "sample_rate": str(sample_rate),
+        "transaction": "Head SDK tx",
+        "trace_id": trace_id,
+    }
+
+
 @pytest.mark.parametrize(
     "args,expected_refcount",
     [({"traces_sample_rate": 1.0}, 100), ({"traces_sample_rate": 0.0}, 0)],
@@ -201,3 +258,27 @@ def capture_event(self, event):
             pass
 
     assert len(events) == 1
+
+
+def test_trace_propagation_meta_head_sdk(sentry_init):
+    sentry_init(traces_sample_rate=1.0, release="foo")
+
+    transaction = Transaction.continue_from_headers({}, name="Head SDK tx")
+    meta = None
+    span = None
+
+    with start_transaction(transaction):
+        with start_span(op="foo", description="foodesc") as current_span:
+            span = current_span
+            meta = Hub.current.trace_propagation_meta()
+
+    ind = meta.find(">") + 1
+    sentry_trace, baggage = meta[:ind], meta[ind:]
+
+    assert 'meta name="sentry-trace"' in sentry_trace
+    sentry_trace_content = re.findall('content="([^"]*)"', sentry_trace)[0]
+    assert sentry_trace_content == span.to_traceparent()
+
+    assert 'meta name="baggage"' in baggage
+    baggage_content = re.findall('content="([^"]*)"', baggage)[0]
+    assert baggage_content == transaction.get_baggage().serialize()