feat(tracing): Dynamic Sampling Context / Baggage continuation (#1485)

* `Baggage` class implementing sentry/third party/mutable logic with parsing from header and serialization
* Parse incoming `baggage` header while starting transaction and store it on the transaction
* Extract `dynamic_sampling_context` fields and add to the `trace` field in the envelope header while sending the transaction
* Propagate the `baggage` header (only sentry fields / no third party as per spec)

[DSC Spec](https://develop.sentry.dev/sdk/performance/dynamic-sampling-context/)

docs/conf.py

@@ -25,9 +25,9 @@
 
 # -- Project information -----------------------------------------------------
 
-project = u"sentry-python"
-copyright = u"2019, Sentry Team and Contributors"
-author = u"Sentry Team and Contributors"
+project = "sentry-python"
+copyright = "2019, Sentry Team and Contributors"
+author = "Sentry Team and Contributors"
 
 release = "1.6.0"
 version = ".".join(release.split(".")[:2])  # The short X.Y version.
@@ -72,7 +72,7 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = [u"_build", "Thumbs.db", ".DS_Store"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = None
@@ -140,8 +140,8 @@
     (
         master_doc,
         "sentry-python.tex",
-        u"sentry-python Documentation",
-        u"Sentry Team and Contributors",
+        "sentry-python Documentation",
+        "Sentry Team and Contributors",
         "manual",
     )
 ]
@@ -151,7 +151,7 @@
 
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
-man_pages = [(master_doc, "sentry-python", u"sentry-python Documentation", [author], 1)]
+man_pages = [(master_doc, "sentry-python", "sentry-python Documentation", [author], 1)]
 
 
 # -- Options for Texinfo output ----------------------------------------------
@@ -163,7 +163,7 @@
     (
         master_doc,
         "sentry-python",
-        u"sentry-python Documentation",
+        "sentry-python Documentation",
         author,
         "sentry-python",
         "One line description of project.",

sentry_sdk/client.py

@@ -373,6 +373,12 @@ def capture_event(
             event_opt.get("contexts", {}).get("trace", {}).pop("tracestate", "")
         )
 
+        dynamic_sampling_context = (
+            event_opt.get("contexts", {})
+            .get("trace", {})
+            .pop("dynamic_sampling_context", {})
+        )
+
         # Transactions or events with attachments should go to the /envelope/
         # endpoint.
         if is_transaction or attachments:
@@ -382,11 +388,15 @@ def capture_event(
                 "sent_at": format_timestamp(datetime.utcnow()),
             }
 
-            tracestate_data = raw_tracestate and reinflate_tracestate(
-                raw_tracestate.replace("sentry=", "")
-            )
-            if tracestate_data and has_tracestate_enabled():
-                headers["trace"] = tracestate_data
+            if has_tracestate_enabled():
+                tracestate_data = raw_tracestate and reinflate_tracestate(
+                    raw_tracestate.replace("sentry=", "")
+                )
+
+                if tracestate_data:
+                    headers["trace"] = tracestate_data
+            elif dynamic_sampling_context:
+                headers["trace"] = dynamic_sampling_context
 
             envelope = Envelope(headers=headers)
 

sentry_sdk/tracing.py

@@ -215,7 +215,7 @@ def continue_from_environ(
         # type: (...) -> Transaction
         """
         Create a Transaction with the given params, then add in data pulled from
-        the 'sentry-trace' and 'tracestate' headers from the environ (if any)
+        the 'sentry-trace', 'baggage' and 'tracestate' headers from the environ (if any)
         before returning the Transaction.
 
         This is different from `continue_from_headers` in that it assumes header
@@ -238,7 +238,7 @@ def continue_from_headers(
         # type: (...) -> Transaction
         """
         Create a transaction with the given params (including any data pulled from
-        the 'sentry-trace' and 'tracestate' headers).
+        the 'sentry-trace', 'baggage' and 'tracestate' headers).
         """
         # TODO move this to the Transaction class
         if cls is Span:
@@ -247,7 +247,17 @@ def continue_from_headers(
                 "instead of Span.continue_from_headers."
             )
 
-        kwargs.update(extract_sentrytrace_data(headers.get("sentry-trace")))
+        # TODO-neel move away from this kwargs stuff, it's confusing and opaque
+        # make more explicit
+        baggage = Baggage.from_incoming_header(headers.get("baggage"))
+        kwargs.update({"baggage": baggage})
+
+        sentrytrace_kwargs = extract_sentrytrace_data(headers.get("sentry-trace"))
+
+        if sentrytrace_kwargs is not None:
+            kwargs.update(sentrytrace_kwargs)
+            baggage.freeze
+
         kwargs.update(extract_tracestate_data(headers.get("tracestate")))
 
         transaction = Transaction(**kwargs)
@@ -258,7 +268,7 @@ def continue_from_headers(
     def iter_headers(self):
         # type: () -> Iterator[Tuple[str, str]]
         """
-        Creates a generator which returns the span's `sentry-trace` and
+        Creates a generator which returns the span's `sentry-trace`, `baggage` and
         `tracestate` headers.
 
         If the span's containing transaction doesn't yet have a
@@ -274,6 +284,9 @@ def iter_headers(self):
         if tracestate:
             yield "tracestate", tracestate
 
+        if self.containing_transaction and self.containing_transaction._baggage:
+            yield "baggage", self.containing_transaction._baggage.serialize()
+
     @classmethod
     def from_traceparent(
         cls,
@@ -460,7 +473,7 @@ def get_trace_context(self):
             "parent_span_id": self.parent_span_id,
             "op": self.op,
             "description": self.description,
-        }
+        }  # type: Dict[str, Any]
         if self.status:
             rv["status"] = self.status
 
@@ -473,6 +486,12 @@ 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:
+            rv[
+                "dynamic_sampling_context"
+            ] = self.containing_transaction._baggage.dynamic_sampling_context()
+
         return rv
 
 
@@ -488,6 +507,7 @@ class Transaction(Span):
         # tracestate data from other vendors, of the form `dogs=yes,cats=maybe`
         "_third_party_tracestate",
         "_measurements",
+        "_baggage",
     )
 
     def __init__(
@@ -496,6 +516,7 @@ def __init__(
         parent_sampled=None,  # type: Optional[bool]
         sentry_tracestate=None,  # type: Optional[str]
         third_party_tracestate=None,  # type: Optional[str]
+        baggage=None,  # type: Optional[Baggage]
         **kwargs  # type: Any
     ):
         # type: (...) -> None
@@ -517,6 +538,7 @@ def __init__(
         self._sentry_tracestate = sentry_tracestate
         self._third_party_tracestate = third_party_tracestate
         self._measurements = {}  # type: Dict[str, Any]
+        self._baggage = baggage
 
     def __repr__(self):
         # type: () -> str
@@ -734,6 +756,7 @@ def _set_initial_sampling_decision(self, sampling_context):
 # Circular imports
 
 from sentry_sdk.tracing_utils import (
+    Baggage,
     EnvironHeaders,
     compute_tracestate_entry,
     extract_sentrytrace_data,

sentry_sdk/tracing_utils.py

@@ -16,13 +16,15 @@
     to_string,
     from_base64,
 )
-from sentry_sdk._compat import PY2
+from sentry_sdk._compat import PY2, iteritems
 from sentry_sdk._types import MYPY
 
 if PY2:
     from collections import Mapping
+    from urllib import quote, unquote
 else:
     from collections.abc import Mapping
+    from urllib.parse import quote, unquote
 
 if MYPY:
     import typing
@@ -211,27 +213,29 @@ def maybe_create_breadcrumbs_from_span(hub, span):
 
 
 def extract_sentrytrace_data(header):
-    # type: (Optional[str]) -> typing.Mapping[str, Union[str, bool, None]]
+    # type: (Optional[str]) -> Optional[typing.Mapping[str, Union[str, bool, None]]]
     """
     Given a `sentry-trace` header string, return a dictionary of data.
     """
-    trace_id = parent_span_id = parent_sampled = None
+    if not header:
+        return None
 
-    if header:
-        if header.startswith("00-") and header.endswith("-00"):
-            header = header[3:-3]
+    if header.startswith("00-") and header.endswith("-00"):
+        header = header[3:-3]
 
-        match = SENTRY_TRACE_REGEX.match(header)
+    match = SENTRY_TRACE_REGEX.match(header)
+    if not match:
+        return None
 
-        if match:
-            trace_id, parent_span_id, sampled_str = match.groups()
+    trace_id, parent_span_id, sampled_str = match.groups()
+    parent_sampled = None
 
-            if trace_id:
-                trace_id = "{:032x}".format(int(trace_id, 16))
-            if parent_span_id:
-                parent_span_id = "{:016x}".format(int(parent_span_id, 16))
-            if sampled_str:
-                parent_sampled = sampled_str != "0"
+    if trace_id:
+        trace_id = "{:032x}".format(int(trace_id, 16))
+    if parent_span_id:
+        parent_span_id = "{:016x}".format(int(parent_span_id, 16))
+    if sampled_str:
+        parent_sampled = sampled_str != "0"
 
     return {
         "trace_id": trace_id,
@@ -413,6 +417,86 @@ def has_custom_measurements_enabled():
     return bool(options and options["_experiments"].get("custom_measurements"))
 
 
+class Baggage(object):
+    __slots__ = ("sentry_items", "third_party_items", "mutable")
+
+    SENTRY_PREFIX = "sentry-"
+    SENTRY_PREFIX_REGEX = re.compile("^sentry-")
+
+    # DynamicSamplingContext
+    DSC_KEYS = [
+        "trace_id",
+        "public_key",
+        "sample_rate",
+        "release",
+        "environment",
+        "transaction",
+        "user_id",
+        "user_segment",
+    ]
+
+    def __init__(
+        self,
+        sentry_items,  # type: Dict[str, str]
+        third_party_items="",  # type: str
+        mutable=True,  # type: bool
+    ):
+        self.sentry_items = sentry_items
+        self.third_party_items = third_party_items
+        self.mutable = mutable
+
+    @classmethod
+    def from_incoming_header(cls, header):
+        # type: (Optional[str]) -> Baggage
+        """
+        freeze if incoming header already has sentry baggage
+        """
+        sentry_items = {}
+        third_party_items = ""
+        mutable = True
+
+        if header:
+            for item in header.split(","):
+                item = item.strip()
+                key, val = item.split("=")
+                if Baggage.SENTRY_PREFIX_REGEX.match(key):
+                    baggage_key = unquote(key.split("-")[1])
+                    sentry_items[baggage_key] = unquote(val)
+                    mutable = False
+                else:
+                    third_party_items += ("," if third_party_items else "") + item
+
+        return Baggage(sentry_items, third_party_items, mutable)
+
+    def freeze(self):
+        # type: () -> None
+        self.mutable = False
+
+    def dynamic_sampling_context(self):
+        # type: () -> Dict[str, str]
+        header = {}
+
+        for key in Baggage.DSC_KEYS:
+            item = self.sentry_items.get(key)
+            if item:
+                header[key] = item
+
+        return header
+
+    def serialize(self, include_third_party=False):
+        # type: (bool) -> str
+        items = []
+
+        for key, val in iteritems(self.sentry_items):
+            item = Baggage.SENTRY_PREFIX + quote(key) + "=" + quote(val)
+            items.append(item)
+
+        if include_third_party:
+            items.append(self.third_party_items)
+
+        return ",".join(items)
+
+
 # Circular imports
 
 if MYPY: