add to_tracestate and use it in iter_headers

lobsterkatie · lobsterkatie · commit b2b51d1727a4 · 2021-01-25T14:07:26.000-08:00
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -693,6 +693,10 @@ def iter_trace_propagation_headers(self):
 
         yield "sentry-trace", span.to_traceparent()
 
+        tracestate = span.to_tracestate()
+        if tracestate:
+            yield "tracestate", tracestate
+
 
 GLOBAL_HUB = Hub()
 _local.set(GLOBAL_HUB)
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -254,6 +254,9 @@ def continue_from_headers(
     def iter_headers(self):
         # type: () -> Generator[Tuple[str, str], None, None]
         yield "sentry-trace", self.to_traceparent()
+        tracestate = self.to_tracestate()
+        if tracestate:
+            yield "tracestate", tracestate
 
     @classmethod
     def from_traceparent(
@@ -313,6 +316,22 @@ def to_traceparent(self):
             sampled = "0"
         return "%s-%s-%s" % (self.trace_id, self.span_id, sampled)
 
+    def to_tracestate(self):
+        # type: () -> Optional[str]
+        transaction = self._containing_transaction
+
+        if not transaction:
+            return None
+
+        # it should never happen that thre's no tracestate value, so this "or"
+        # is just insurance
+        header_value = transaction._sentry_tracestate_value or "{}"
+
+        if transaction._third_party_tracestate:
+            header_value = ",".join([header_value, transaction._third_party_tracestate])
+
+        return header_value
+
     def set_tag(self, key, value):
         # type: (str, Any) -> None
         self._tags[key] = value
@@ -419,8 +438,13 @@ def get_trace_context(self):
         if self.status:
             rv["status"] = self.status
 
-        if self._containing_transaction:
-            rv["tracestate"] = self._containing_transaction._sentry_tracestate_value
+        if isinstance(self, Transaction):
+            transaction = self  # type: Optional[Transaction]
+        else:
+            transaction = self._containing_transaction
+
+        if transaction:
+            rv["tracestate"] = transaction._sentry_tracestate_value
 
         return rv
 
