Only inject ot-baggage-* into the HTTP headers once.

This resolves the issue that occurs when both DATADOG and TRACECONTEXT propagation styles are active and the carrier injects the same key multiple times (e.g. GrpcHttp2OutboundHeaders).

Author: ygree

dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java

@@ -40,17 +40,20 @@ private DatadogHttpCodec() {
     // This class should not be created. This also makes code coverage checks happy.
   }
 
-  public static HttpCodec.Injector newInjector(Map<String, String> invertedBaggageMapping) {
-    return new Injector(invertedBaggageMapping);
+  public static HttpCodec.Injector newInjector(
+      Map<String, String> invertedBaggageMapping, boolean isInjectOtBaggage) {
+    return new Injector(invertedBaggageMapping, isInjectOtBaggage);
   }
 
   private static class Injector implements HttpCodec.Injector {
 
     private final Map<String, String> invertedBaggageMapping;
+    private final boolean isInjectOtBaggage;
 
-    public Injector(Map<String, String> invertedBaggageMapping) {
+    public Injector(Map<String, String> invertedBaggageMapping, boolean isInjectOtBaggage) {
       assert invertedBaggageMapping != null;
       this.invertedBaggageMapping = invertedBaggageMapping;
+      this.isInjectOtBaggage = isInjectOtBaggage;
     }
 
     @Override
@@ -66,6 +69,17 @@ public <C> void inject(
       if (origin != null) {
         setter.set(carrier, ORIGIN_KEY, origin.toString());
       }
+      if (isInjectOtBaggage) {
+        injectBaggage(context, carrier, setter);
+      }
+      // inject x-datadog-tags
+      String datadogTags = context.getPropagationTags().headerValue(HeaderType.DATADOG);
+      if (datadogTags != null) {
+        setter.set(carrier, DATADOG_TAGS_KEY, datadogTags);
+      }
+    }
+
+    private <C> void injectBaggage(DDSpanContext context, C carrier, CarrierSetter<C> setter) {
       long e2eStart = context.getEndToEndStartTime();
       if (e2eStart > 0) {
         setter.set(carrier, E2E_START_KEY, Long.toString(NANOSECONDS.toMillis(e2eStart)));
@@ -76,12 +90,6 @@ public <C> void inject(
         header = header != null ? header : OT_BAGGAGE_PREFIX + entry.getKey();
         setter.set(carrier, header, HttpCodec.encodeBaggage(entry.getValue()));
       }
-
-      // inject x-datadog-tags
-      String datadogTags = context.getPropagationTags().headerValue(HeaderType.DATADOG);
-      if (datadogTags != null) {
-        setter.set(carrier, DATADOG_TAGS_KEY, datadogTags);
-      }
     }
   }
 

dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java

@@ -98,10 +98,12 @@ private static Map<TracePropagationStyle, Injector> createInjectors(
       Set<TracePropagationStyle> propagationStyles,
       Map<String, String> reverseBaggageMapping) {
     EnumMap<TracePropagationStyle, Injector> result = new EnumMap<>(TracePropagationStyle.class);
+    boolean isInjectOtBaggage = true;
     for (TracePropagationStyle style : propagationStyles) {
       switch (style) {
         case DATADOG:
-          result.put(style, DatadogHttpCodec.newInjector(reverseBaggageMapping));
+          result.put(style, DatadogHttpCodec.newInjector(reverseBaggageMapping, isInjectOtBaggage));
+          isInjectOtBaggage = false;
           break;
         case B3SINGLE:
           result.put(
@@ -123,7 +125,8 @@ private static Map<TracePropagationStyle, Injector> createInjectors(
           result.put(style, NoneCodec.INJECTOR);
           break;
         case TRACECONTEXT:
-          result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping));
+          result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping, isInjectOtBaggage));
+          isInjectOtBaggage = false;
           break;
         case BAGGAGE:
           break;

dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java

@@ -48,25 +48,30 @@ private W3CHttpCodec() {
     // This class should not be created. This also makes code coverage checks happy.
   }
 
-  public static HttpCodec.Injector newInjector(Map<String, String> invertedBaggageMapping) {
-    return new Injector(invertedBaggageMapping);
+  public static HttpCodec.Injector newInjector(
+      Map<String, String> invertedBaggageMapping, boolean isInjectBaggage) {
+    return new Injector(invertedBaggageMapping, isInjectBaggage);
   }
 
   private static class Injector implements HttpCodec.Injector {
 
     private final Map<String, String> invertedBaggageMapping;
+    private final boolean isInjectBaggage;
 
-    public Injector(Map<String, String> invertedBaggageMapping) {
+    public Injector(Map<String, String> invertedBaggageMapping, boolean isInjectBaggage) {
       assert invertedBaggageMapping != null;
       this.invertedBaggageMapping = invertedBaggageMapping;
+      this.isInjectBaggage = isInjectBaggage;
     }
 
     @Override
     public <C> void inject(
         final DDSpanContext context, final C carrier, final CarrierSetter<C> setter) {
       injectTraceParent(context, carrier, setter);
       injectTraceState(context, carrier, setter);
-      injectBaggage(context, carrier, setter);
+      if (isInjectBaggage) {
+        injectBaggage(context, carrier, setter);
+      }
     }
 
     private <C> void injectTraceParent(DDSpanContext context, C carrier, CarrierSetter<C> setter) {

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/DatadogHttpInjectorTest.groovy

@@ -17,7 +17,7 @@ import static datadog.trace.core.propagation.DatadogHttpCodec.*
 
 class DatadogHttpInjectorTest extends DDCoreSpecification {
 
-  HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"])
+  HttpCodec.Injector injector = newInjector(["some-baggage-key": "SOME_CUSTOM_HEADER"], true)
 
   def "inject http headers"() {
     setup:

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpInjectorTest.groovy

@@ -218,6 +218,35 @@ class HttpInjectorTest extends DDCoreSpecification {
     style << [DATADOG, TRACECONTEXT, HAYSTACK]
   }
 
+  def "inject ot-baggage-* in http headers only once"() {
+    setup:
+    Config config = Mock(Config)
+    def baggage = [
+      'k1': 'v1',
+      'k2': 'v2',
+    ]
+    def mapping = baggage.keySet().collectEntries { [(it):it]} as Map<String, String>
+    def injector = HttpCodec.createInjector(config, [DATADOG, TRACECONTEXT].toSet(), mapping)
+    def traceId = DDTraceId.ONE
+    def spanId = 2
+    def writer = new ListWriter()
+    def tracer = tracerBuilder().writer(writer).build()
+    final DDSpanContext mockedContext = mockedContext(tracer, traceId, spanId, UNSET, null, baggage)
+    final Map<String, String> carrier = Mock()
+    mockedContext.beginEndToEnd()
+
+    when:
+    injector.inject(mockedContext, carrier, MapSetter.INSTANCE)
+
+    then:
+    1 * carrier.put('k1', 'v1')
+    1 * carrier.put('k2', 'v2')
+    1 * carrier.put('ot-baggage-t0', "${(long) (mockedContext.endToEndStartTime / 1000000L)}")
+
+    cleanup:
+    tracer.close()
+  }
+
   static DDSpanContext mockedContext(CoreTracer tracer, DDTraceId traceId, long spanId, int samplingPriority, String origin, Map<String, String> baggage) {
     return new DDSpanContext(
       traceId,

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CHttpInjectorTest.groovy

@@ -21,7 +21,7 @@ import static datadog.trace.core.propagation.W3CHttpCodec.newInjector
 
 class W3CHttpInjectorTest extends DDCoreSpecification {
 
-  HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"])
+  HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"], true)
 
   def "inject http headers"() {
     setup: