Merge baggage headers (incoming and new created ones) (#3001)

If the incoming headers include a baggage header that also includes sentry trace data, use this incoming trace information and merge with other baggage information of the current transaction. (Before it was just concatenating incoming baggage information with the baggage information from the current transaction which lead sometimes to ever growing baggage headers.)

Author: antonpirker

sentry_sdk/integrations/celery/__init__.py

@@ -16,6 +16,7 @@
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, TRANSACTION_SOURCE_TASK
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.scope import Scope
+from sentry_sdk.tracing_utils import Baggage
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     ensure_integration_enabled,
@@ -168,6 +169,7 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks):
             headers = dict(
                 Scope.get_current_scope().iter_trace_propagation_headers(span=span)
             )
+
         if monitor_beat_tasks:
             headers.update(
                 {
@@ -182,10 +184,23 @@ def _update_celery_task_headers(original_headers, span, monitor_beat_tasks):
 
             combined_baggage = sentry_baggage or existing_baggage
             if sentry_baggage and existing_baggage:
-                combined_baggage = "{},{}".format(
-                    existing_baggage,
-                    sentry_baggage,
+                # Merge incoming and sentry baggage, where the sentry trace information
+                # in the incoming baggage takes precedence and the third-party items
+                # are concatenated.
+                incoming = Baggage.from_incoming_header(existing_baggage)
+                combined = Baggage.from_incoming_header(sentry_baggage)
+                combined.sentry_items.update(incoming.sentry_items)
+                combined.third_party_items = ",".join(
+                    [
+                        x
+                        for x in [
+                            combined.third_party_items,
+                            incoming.third_party_items,
+                        ]
+                        if x is not None and x != ""
+                    ]
                 )
+                combined_baggage = combined.serialize(include_third_party=True)
 
             updated_headers.update(headers)
             if combined_baggage:

tests/integrations/celery/test_update_celery_task_headers.py

@@ -1,9 +1,11 @@
+from copy import copy
 import pytest
 
 from unittest import mock
 
 from sentry_sdk.integrations.celery import _update_celery_task_headers
 import sentry_sdk
+from sentry_sdk.tracing_utils import Baggage
 
 
 BAGGAGE_VALUE = (
@@ -115,17 +117,25 @@ def test_span_with_transaction_custom_headers(sentry_init):
 
             assert updated_headers["sentry-trace"] == span.to_traceparent()
             assert updated_headers["headers"]["sentry-trace"] == span.to_traceparent()
-            # This is probably the cause for https://github.com/getsentry/sentry-python/issues/2916
-            # If incoming baggage includes sentry data, we should not concatenate a new baggage value to it
-            # but just keep the incoming sentry baggage values and concatenate new third-party items to the baggage
-            # I have some code somewhere where I have implemented this.
-            assert (
-                updated_headers["baggage"]
-                == headers["baggage"] + "," + transaction.get_baggage().serialize()
+
+            incoming_baggage = Baggage.from_incoming_header(headers["baggage"])
+            combined_baggage = copy(transaction.get_baggage())
+            combined_baggage.sentry_items.update(incoming_baggage.sentry_items)
+            combined_baggage.third_party_items = ",".join(
+                [
+                    x
+                    for x in [
+                        combined_baggage.third_party_items,
+                        incoming_baggage.third_party_items,
+                    ]
+                    if x is not None and x != ""
+                ]
             )
-            assert (
-                updated_headers["headers"]["baggage"]
-                == headers["baggage"] + "," + transaction.get_baggage().serialize()
+            assert updated_headers["baggage"] == combined_baggage.serialize(
+                include_third_party=True
+            )
+            assert updated_headers["headers"]["baggage"] == combined_baggage.serialize(
+                include_third_party=True
             )