Support adding W3C baggage as span tags (#9169)

* WIP: baggage span tags

* simple test

* fix build errors

* adding more tests

* Add baggage span tags functionality and tests

* Move baggage tags feature to HttpServerDecorator.getExtractedSpanContext and update it to check the W3C baggage held in the new Context API

(the baggage map in CoreTracer is OpenTracing only, and won't contain W3C baggage)

* Fix potential corner-case where TagContext.getTags() may return an immutable EMPTY TagMap

* Move baggage tags unit test to HttpServerDecoratorTest

* Baggage tags unit tests have been moved to HttpServerDecoratorTest

* Move W3C baggage tag injection to just before the span is serialized

* Update test trace writer so injected baggage metadata can be asserted as tags

* spotless format

* Remove unused method

* Fix test code

* Support recording of baggage, that would be sent to agent as trace metadata, during tests

---------

Co-authored-by: Stuart McCulloch <[email protected]>

Author: rachelyangdog

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy

@@ -39,6 +39,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
 import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.core.DDSpan
+import datadog.trace.core.Metadata
 import datadog.trace.core.datastreams.StatsGroup
 import datadog.trace.test.util.Flaky
 import groovy.json.JsonOutput
@@ -778,6 +779,62 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
     'GET'  | null | 'x-datadog-test-request-header' | 'bar' | ['request_header_tag': 'bar']
   }
 
+  def "test baggage span tags are properly added"() {
+    setup:
+    def recordedBaggageTags = [:]
+    TEST_WRITER.metadataConsumer = { Metadata md ->
+      md.baggage.forEach { k, v ->
+        // record non-internal baggage sent to agent as trace metadata
+        if (!k.startsWith("_dd.")) {
+          recordedBaggageTags.put(k, v)
+        }
+      }
+    }
+    // Use default configuration for TRACE_BAGGAGE_TAG_KEYS (user.id, session.id, account.id)
+    def baggageHeader = "user.id=test-user,session.id=test-session,account.id=test-account,language=en"
+    def request = request(SUCCESS, 'GET', null)
+    .header("baggage", baggageHeader)
+    .build()
+    def response = client.newCall(request).execute()
+    if (isDataStreamsEnabled()) {
+      TEST_DATA_STREAMS_WRITER.waitForGroups(1)
+    }
+
+    expect:
+    response.code() == SUCCESS.status
+    response.body().string() == SUCCESS.body
+
+    and:
+    assertTraces(1) {
+      trace(spanCount(SUCCESS)) {
+        sortSpansByStart()
+        // Verify baggage tags are added for default configured keys only
+        serverSpan(it, null, null, 'GET', SUCCESS)
+        if (hasHandlerSpan()) {
+          handlerSpan(it)
+        }
+        controllerSpan(it)
+        if (hasResponseSpan(SUCCESS)) {
+          responseSpan(it, SUCCESS)
+        }
+      }
+    }
+    recordedBaggageTags == [
+      "baggage.user.id": "test-user",
+      "baggage.session.id": "test-session",
+      "baggage.account.id": "test-account"
+      // "baggage.language" should NOT be present since it's not in default config
+    ]
+
+    and:
+    if (isDataStreamsEnabled()) {
+      StatsGroup first = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == 0 }
+      verifyAll(first) {
+        tags == DSM_EDGE_TAGS
+      }
+    }
+  }
+
   @Flaky(value = "https://github.com/DataDog/dd-trace-java/issues/4690", suites = ["MuleHttpServerForkedTest"])
   def "test QUERY_ENCODED_BOTH with response header x-ig-response-header tag mapping"() {
     setup:

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -84,6 +84,8 @@ public final class ConfigDefaults {
       new LinkedHashSet<>(asList(PropagationStyle.DATADOG));
   static final int DEFAULT_TRACE_BAGGAGE_MAX_ITEMS = 64;
   static final int DEFAULT_TRACE_BAGGAGE_MAX_BYTES = 8192;
+  static final List<String> DEFAULT_TRACE_BAGGAGE_TAG_KEYS =
+      Arrays.asList("user.id", "session.id", "account.id");
   static final boolean DEFAULT_JMX_FETCH_ENABLED = true;
   static final boolean DEFAULT_TRACE_AGENT_V05_ENABLED = false;
 

dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java

@@ -98,6 +98,7 @@ public final class TracerConfig {
   public static final String TRACE_PROPAGATION_EXTRACT_FIRST = "trace.propagation.extract.first";
   public static final String TRACE_BAGGAGE_MAX_ITEMS = "trace.baggage.max.items";
   public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes";
+  public static final String TRACE_BAGGAGE_TAG_KEYS = "trace.baggage.tag.keys";
 
   public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled";
 

dd-trace-core/src/main/java/datadog/trace/common/writer/ListWriter.java

@@ -26,6 +26,8 @@ public class ListWriter extends CopyOnWriteArrayList<List<DDSpan>> implements Wr
 
   private Filter filter = ACCEPT_ALL;
 
+  private MetadataConsumer metadataConsumer = MetadataConsumer.NO_OP;
+
   public List<DDSpan> firstTrace() {
     return get(0);
   }
@@ -38,7 +40,7 @@ public void write(List<DDSpan> trace) {
     for (DDSpan span : trace) {
       // This is needed to properly do all delayed processing to make this writer even
       // remotely realistic so the test actually test something
-      span.processTagsAndBaggage(MetadataConsumer.NO_OP);
+      span.processTagsAndBaggage(metadataConsumer);
     }
 
     add(trace);
@@ -117,6 +119,11 @@ public void setFilter(Filter filter) {
     this.filter = filter;
   }
 
+  /** Set a {@link MetadataConsumer} to capture what trace metadata would be sent to the agent. */
+  public void setMetadataConsumer(MetadataConsumer metadataConsumer) {
+    this.metadataConsumer = metadataConsumer;
+  }
+
   private boolean isReported(DDSpan span) {
     for (List<DDSpan> trace : this) {
       for (DDSpan aSpan : trace) {

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -59,6 +59,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Baggage;
 import datadog.trace.bootstrap.instrumentation.api.BlackHoleSpan;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
@@ -1591,6 +1592,7 @@ private DDSpanContext buildSpanContext() {
       final long spanId;
       final long parentSpanId;
       final Map<String, String> baggage;
+      final Baggage w3cBaggage;
       final TraceCollector parentTraceCollector;
       final int samplingPriority;
       final CharSequence origin;
@@ -1645,6 +1647,7 @@ private DDSpanContext buildSpanContext() {
         traceId = ddsc.getTraceId();
         parentSpanId = ddsc.getSpanId();
         baggage = ddsc.getBaggageItems();
+        w3cBaggage = null;
         parentTraceCollector = ddsc.getTraceCollector();
         samplingPriority = PrioritySampling.UNSET;
         origin = null;
@@ -1706,6 +1709,7 @@ private DDSpanContext buildSpanContext() {
           coreTagsNeedsIntercept = true; // maybe intercept isn't needed?
           origin = tc.getOrigin();
           baggage = tc.getBaggage();
+          w3cBaggage = tc.getW3CBaggage();
           requestContextDataAppSec = tc.getRequestContextDataAppSec();
           requestContextDataIast = tc.getRequestContextDataIast();
           ciVisibilityContextData = tc.getCiVisibilityContextData();
@@ -1715,6 +1719,7 @@ private DDSpanContext buildSpanContext() {
           coreTagsNeedsIntercept = false;
           origin = null;
           baggage = null;
+          w3cBaggage = null;
           requestContextDataAppSec = null;
           requestContextDataIast = null;
           ciVisibilityContextData = null;
@@ -1805,6 +1810,7 @@ private DDSpanContext buildSpanContext() {
               samplingPriority,
               origin,
               baggage,
+              w3cBaggage,
               errorFlag,
               spanType,
               tagsSize,

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -5,6 +5,7 @@
 import static datadog.trace.api.cache.RadixTreeCache.HTTP_STATUSES;
 import static datadog.trace.bootstrap.instrumentation.api.ErrorPriorities.UNSET;
 
+import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.DDTraceId;
@@ -23,6 +24,7 @@
 import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
+import datadog.trace.bootstrap.instrumentation.api.Baggage;
 import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
@@ -77,6 +79,8 @@ public class DDSpanContext
   /** Baggage is associated with the whole trace and shared with other spans */
   private volatile Map<String, String> baggageItems;
 
+  private final Baggage w3cBaggage;
+
   // Not Shared with other span contexts
   private final DDTraceId traceId;
   private final long spanId;
@@ -187,6 +191,7 @@ public DDSpanContext(
         samplingPriority,
         origin,
         baggageItems,
+        null,
         errorFlag,
         spanType,
         tagsSize,
@@ -233,6 +238,7 @@ public DDSpanContext(
         samplingPriority,
         origin,
         baggageItems,
+        null,
         errorFlag,
         spanType,
         tagsSize,
@@ -279,6 +285,7 @@ public DDSpanContext(
         samplingPriority,
         origin,
         baggageItems,
+        null,
         errorFlag,
         spanType,
         tagsSize,
@@ -304,6 +311,7 @@ public DDSpanContext(
       final int samplingPriority,
       final CharSequence origin,
       final Map<String, String> baggageItems,
+      final Baggage w3cBaggage,
       final boolean errorFlag,
       final CharSequence spanType,
       final int tagsSize,
@@ -331,6 +339,7 @@ public DDSpanContext(
     } else {
       this.baggageItems = new ConcurrentHashMap<>(baggageItems);
     }
+    this.w3cBaggage = w3cBaggage;
 
     this.requestContextDataAppSec = requestContextDataAppSec;
     this.requestContextDataIast = requestContextDataIast;
@@ -926,6 +935,9 @@ public void processTagsAndBaggage(
       Map<String, String> baggageItemsWithPropagationTags;
       if (injectBaggageAsTags) {
         baggageItemsWithPropagationTags = new HashMap<>(baggageItems);
+        if (w3cBaggage != null) {
+          injectW3CBaggageTags(baggageItemsWithPropagationTags);
+        }
         propagationTags.fillTagMap(baggageItemsWithPropagationTags);
       } else {
         baggageItemsWithPropagationTags = propagationTags.createTagMap();
@@ -948,6 +960,27 @@ public void processTagsAndBaggage(
     }
   }
 
+  void injectW3CBaggageTags(Map<String, String> baggageItemsWithPropagationTags) {
+    List<String> baggageTagKeys = Config.get().getTraceBaggageTagKeys();
+    if (baggageTagKeys.isEmpty()) {
+      return;
+    }
+    Map<String, String> w3cBaggageItems = w3cBaggage.asMap();
+    for (String key : baggageTagKeys) {
+      if ("*".equals(key)) {
+        // If the key is "*", we add all baggage items
+        for (Map.Entry<String, String> entry : w3cBaggageItems.entrySet()) {
+          baggageItemsWithPropagationTags.put("baggage." + entry.getKey(), entry.getValue());
+        }
+        break;
+      }
+      String value = w3cBaggageItems.get(key);
+      if (value != null) {
+        baggageItemsWithPropagationTags.put("baggage." + key, value);
+      }
+    }
+  }
+
   @Override
   public void setIntegrationName(CharSequence integrationName) {
     this.integrationName = integrationName;

dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java

@@ -7,7 +7,10 @@
 import datadog.context.propagation.CarrierVisitor;
 import datadog.context.propagation.Propagator;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.Baggage;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.core.util.PercentEscaper;
 import datadog.trace.core.util.PercentEscaper.Escaped;
 import java.io.UnsupportedEncodingException;
@@ -109,7 +112,21 @@ public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
     }
     BaggageExtractor baggageExtractor = new BaggageExtractor();
     visitor.forEachKeyValue(carrier, baggageExtractor);
-    return baggageExtractor.extracted == null ? context : context.with(baggageExtractor.extracted);
+    Baggage baggage = baggageExtractor.extracted;
+    if (baggage == null) {
+      return context;
+    }
+
+    // TODO: consider a better way to link baggage with the extracted (legacy) TagContext
+    AgentSpan extractedSpan = AgentSpan.fromContext(context);
+    if (extractedSpan != null) {
+      AgentSpanContext extractedSpanContext = extractedSpan.context();
+      if (extractedSpanContext instanceof TagContext) {
+        ((TagContext) extractedSpanContext).setW3CBaggage(baggage);
+      }
+    }
+
+    return context.with(baggage);
   }
 
   private class BaggageExtractor implements BiConsumer<String, String> {

dd-trace-core/src/test/groovy/datadog/trace/core/test/DDCoreSpecification.groovy

@@ -70,6 +70,7 @@ abstract class DDCoreSpecification extends DDSpecification {
       prioritySampling,
       null,
       [:],
+      null,
       false,
       spanType,
       0,

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -151,6 +151,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_ANALYTICS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_BAGGAGE_MAX_BYTES;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_BAGGAGE_MAX_ITEMS;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_BAGGAGE_TAG_KEYS;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_CLOUD_PAYLOAD_TAGGING_SERVICES;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_EXPERIMENTAL_FEATURES_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH;
@@ -583,6 +584,7 @@
 import static datadog.trace.api.config.TracerConfig.TRACE_ANALYTICS_ENABLED;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_BYTES;
 import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_MAX_ITEMS;
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_TAG_KEYS;
 import static datadog.trace.api.config.TracerConfig.TRACE_CLIENT_IP_HEADER;
 import static datadog.trace.api.config.TracerConfig.TRACE_CLIENT_IP_RESOLVER_ENABLED;
 import static datadog.trace.api.config.TracerConfig.TRACE_CLOUD_PAYLOAD_TAGGING_MAX_DEPTH;
@@ -823,6 +825,7 @@ public static String getHostName() {
   private final boolean tracePropagationExtractFirst;
   private final int traceBaggageMaxItems;
   private final int traceBaggageMaxBytes;
+  private final List<String> traceBaggageTagKeys;
   private final int clockSyncPeriod;
   private final boolean logsInjectionEnabled;
 
@@ -1694,6 +1697,11 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
         // If we have a new setting, we log a warning
         logOverriddenDeprecatedSettingWarning(PROPAGATION_STYLE_INJECT, injectOrigin, inject);
       }
+
+      // Parse the baggage tag keys configuration
+      traceBaggageTagKeys =
+          configProvider.getList(TRACE_BAGGAGE_TAG_KEYS, DEFAULT_TRACE_BAGGAGE_TAG_KEYS);
+
       // Now we can check if we should pick the default injection/extraction
 
       tracePropagationStylesToExtract =
@@ -2966,6 +2974,10 @@ public Map<String, String> getBaggageMapping() {
     return baggageMapping;
   }
 
+  public List<String> getTraceBaggageTagKeys() {
+    return traceBaggageTagKeys;
+  }
+
   public Map<String, String> getHttpServerPathResourceNameMapping() {
     return httpServerPathResourceNameMapping;
   }
@@ -5370,6 +5382,8 @@ public String toString() {
         + traceKeepLatencyThreshold
         + ", traceStrictWritesEnabled="
         + traceStrictWritesEnabled
+        + ", traceBaggageTagKeys="
+        + traceBaggageTagKeys
         + ", tracePropagationStylesToExtract="
         + tracePropagationStylesToExtract
         + ", tracePropagationStylesToInject="

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java

@@ -14,7 +14,9 @@
 
 public final class AgentPropagation {
   public static final Concern TRACING_CONCERN = named("tracing");
-  public static final Concern BAGGAGE_CONCERN = named("baggage");
+  // TODO: Baggage propagator should run after tracing so it can link baggage with the span context
+  // TODO: remove this priority once we have a story for replacing TagContext with the Context API
+  public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105);
   public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray");
 
   // TODO DSM propagator should run after the other propagators as it stores the pathway context