diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java index 6723dfdbc8c..b72bdd46d66 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java +++ b/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 invertedBaggageMapping) { - return new Injector(invertedBaggageMapping); + public static HttpCodec.Injector newInjector( + Map invertedBaggageMapping, boolean isInjectOtBaggage) { + return new Injector(invertedBaggageMapping, isInjectOtBaggage); } private static class Injector implements HttpCodec.Injector { private final Map invertedBaggageMapping; + private final boolean isInjectOtBaggage; - public Injector(Map invertedBaggageMapping) { + public Injector(Map invertedBaggageMapping, boolean isInjectOtBaggage) { assert invertedBaggageMapping != null; this.invertedBaggageMapping = invertedBaggageMapping; + this.isInjectOtBaggage = isInjectOtBaggage; } @Override @@ -66,6 +69,17 @@ public 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 void injectBaggage(DDSpanContext context, C carrier, CarrierSetter setter) { long e2eStart = context.getEndToEndStartTime(); if (e2eStart > 0) { setter.set(carrier, E2E_START_KEY, Long.toString(NANOSECONDS.toMillis(e2eStart))); @@ -76,12 +90,6 @@ public 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); - } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 9835d2e1c03..58d92f8b7ed 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -98,10 +98,12 @@ private static Map createInjectors( Set propagationStyles, Map reverseBaggageMapping) { EnumMap 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 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; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java index c201047fd01..4a01ae78429 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/W3CHttpCodec.java @@ -48,17 +48,20 @@ private W3CHttpCodec() { // This class should not be created. This also makes code coverage checks happy. } - public static HttpCodec.Injector newInjector(Map invertedBaggageMapping) { - return new Injector(invertedBaggageMapping); + public static HttpCodec.Injector newInjector( + Map invertedBaggageMapping, boolean isInjectBaggage) { + return new Injector(invertedBaggageMapping, isInjectBaggage); } private static class Injector implements HttpCodec.Injector { private final Map invertedBaggageMapping; + private final boolean isInjectBaggage; - public Injector(Map invertedBaggageMapping) { + public Injector(Map invertedBaggageMapping, boolean isInjectBaggage) { assert invertedBaggageMapping != null; this.invertedBaggageMapping = invertedBaggageMapping; + this.isInjectBaggage = isInjectBaggage; } @Override @@ -66,7 +69,9 @@ public void inject( final DDSpanContext context, final C carrier, final CarrierSetter setter) { injectTraceParent(context, carrier, setter); injectTraceState(context, carrier, setter); - injectBaggage(context, carrier, setter); + if (isInjectBaggage) { + injectBaggage(context, carrier, setter); + } } private void injectTraceParent(DDSpanContext context, C carrier, CarrierSetter setter) { diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/DatadogHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/DatadogHttpInjectorTest.groovy index a51a37fd929..40ad666d557 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/DatadogHttpInjectorTest.groovy +++ b/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: diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpInjectorTest.groovy index 4274be2af55..e7b641e2c42 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpInjectorTest.groovy +++ b/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 + 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 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 baggage) { return new DDSpanContext( traceId, diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CHttpInjectorTest.groovy index 06c41767168..d92082c405b 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/W3CHttpInjectorTest.groovy +++ b/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: