Skip to content

Commit 31247ab

Browse files
ygreeygree
committed
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).
1 parent 562e533 commit 31247ab

6 files changed

Lines changed: 62 additions & 17 deletions

File tree

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

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,17 +40,20 @@ private DatadogHttpCodec() {
4040
// This class should not be created. This also makes code coverage checks happy.
4141
}
4242

43-
public static HttpCodec.Injector newInjector(Map<String, String> invertedBaggageMapping) {
44-
return new Injector(invertedBaggageMapping);
43+
public static HttpCodec.Injector newInjector(
44+
Map<String, String> invertedBaggageMapping, boolean isInjectOtBaggage) {
45+
return new Injector(invertedBaggageMapping, isInjectOtBaggage);
4546
}
4647

4748
private static class Injector implements HttpCodec.Injector {
4849

4950
private final Map<String, String> invertedBaggageMapping;
51+
private final boolean isInjectOtBaggage;
5052

51-
public Injector(Map<String, String> invertedBaggageMapping) {
53+
public Injector(Map<String, String> invertedBaggageMapping, boolean isInjectOtBaggage) {
5254
assert invertedBaggageMapping != null;
5355
this.invertedBaggageMapping = invertedBaggageMapping;
56+
this.isInjectOtBaggage = isInjectOtBaggage;
5457
}
5558

5659
@Override
@@ -66,6 +69,17 @@ public <C> void inject(
6669
if (origin != null) {
6770
setter.set(carrier, ORIGIN_KEY, origin.toString());
6871
}
72+
if (isInjectOtBaggage) {
73+
injectBaggage(context, carrier, setter);
74+
}
75+
// inject x-datadog-tags
76+
String datadogTags = context.getPropagationTags().headerValue(HeaderType.DATADOG);
77+
if (datadogTags != null) {
78+
setter.set(carrier, DATADOG_TAGS_KEY, datadogTags);
79+
}
80+
}
81+
82+
private <C> void injectBaggage(DDSpanContext context, C carrier, CarrierSetter<C> setter) {
6983
long e2eStart = context.getEndToEndStartTime();
7084
if (e2eStart > 0) {
7185
setter.set(carrier, E2E_START_KEY, Long.toString(NANOSECONDS.toMillis(e2eStart)));
@@ -76,12 +90,6 @@ public <C> void inject(
7690
header = header != null ? header : OT_BAGGAGE_PREFIX + entry.getKey();
7791
setter.set(carrier, header, HttpCodec.encodeBaggage(entry.getValue()));
7892
}
79-
80-
// inject x-datadog-tags
81-
String datadogTags = context.getPropagationTags().headerValue(HeaderType.DATADOG);
82-
if (datadogTags != null) {
83-
setter.set(carrier, DATADOG_TAGS_KEY, datadogTags);
84-
}
8593
}
8694
}
8795

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,10 +98,12 @@ private static Map<TracePropagationStyle, Injector> createInjectors(
9898
Set<TracePropagationStyle> propagationStyles,
9999
Map<String, String> reverseBaggageMapping) {
100100
EnumMap<TracePropagationStyle, Injector> result = new EnumMap<>(TracePropagationStyle.class);
101+
boolean isInjectOtBaggage = true;
101102
for (TracePropagationStyle style : propagationStyles) {
102103
switch (style) {
103104
case DATADOG:
104-
result.put(style, DatadogHttpCodec.newInjector(reverseBaggageMapping));
105+
result.put(style, DatadogHttpCodec.newInjector(reverseBaggageMapping, isInjectOtBaggage));
106+
isInjectOtBaggage = false;
105107
break;
106108
case B3SINGLE:
107109
result.put(
@@ -123,7 +125,8 @@ private static Map<TracePropagationStyle, Injector> createInjectors(
123125
result.put(style, NoneCodec.INJECTOR);
124126
break;
125127
case TRACECONTEXT:
126-
result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping));
128+
result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping, isInjectOtBaggage));
129+
isInjectOtBaggage = false;
127130
break;
128131
case BAGGAGE:
129132
break;

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,25 +48,30 @@ private W3CHttpCodec() {
4848
// This class should not be created. This also makes code coverage checks happy.
4949
}
5050

51-
public static HttpCodec.Injector newInjector(Map<String, String> invertedBaggageMapping) {
52-
return new Injector(invertedBaggageMapping);
51+
public static HttpCodec.Injector newInjector(
52+
Map<String, String> invertedBaggageMapping, boolean isInjectBaggage) {
53+
return new Injector(invertedBaggageMapping, isInjectBaggage);
5354
}
5455

5556
private static class Injector implements HttpCodec.Injector {
5657

5758
private final Map<String, String> invertedBaggageMapping;
59+
private final boolean isInjectBaggage;
5860

59-
public Injector(Map<String, String> invertedBaggageMapping) {
61+
public Injector(Map<String, String> invertedBaggageMapping, boolean isInjectBaggage) {
6062
assert invertedBaggageMapping != null;
6163
this.invertedBaggageMapping = invertedBaggageMapping;
64+
this.isInjectBaggage = isInjectBaggage;
6265
}
6366

6467
@Override
6568
public <C> void inject(
6669
final DDSpanContext context, final C carrier, final CarrierSetter<C> setter) {
6770
injectTraceParent(context, carrier, setter);
6871
injectTraceState(context, carrier, setter);
69-
injectBaggage(context, carrier, setter);
72+
if (isInjectBaggage) {
73+
injectBaggage(context, carrier, setter);
74+
}
7075
}
7176

7277
private <C> void injectTraceParent(DDSpanContext context, C carrier, CarrierSetter<C> setter) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import static datadog.trace.core.propagation.DatadogHttpCodec.*
1717

1818
class DatadogHttpInjectorTest extends DDCoreSpecification {
1919

20-
HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"])
20+
HttpCodec.Injector injector = newInjector(["some-baggage-key": "SOME_CUSTOM_HEADER"], true)
2121

2222
def "inject http headers"() {
2323
setup:

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,35 @@ class HttpInjectorTest extends DDCoreSpecification {
218218
style << [DATADOG, TRACECONTEXT, HAYSTACK]
219219
}
220220

221+
def "inject ot-baggage-* in http headers only once"() {
222+
setup:
223+
Config config = Mock(Config)
224+
def baggage = [
225+
'k1': 'v1',
226+
'k2': 'v2',
227+
]
228+
def mapping = baggage.keySet().collectEntries { [(it):it]} as Map<String, String>
229+
def injector = HttpCodec.createInjector(config, [DATADOG, TRACECONTEXT].toSet(), mapping)
230+
def traceId = DDTraceId.ONE
231+
def spanId = 2
232+
def writer = new ListWriter()
233+
def tracer = tracerBuilder().writer(writer).build()
234+
final DDSpanContext mockedContext = mockedContext(tracer, traceId, spanId, UNSET, null, baggage)
235+
final Map<String, String> carrier = Mock()
236+
mockedContext.beginEndToEnd()
237+
238+
when:
239+
injector.inject(mockedContext, carrier, MapSetter.INSTANCE)
240+
241+
then:
242+
1 * carrier.put('k1', 'v1')
243+
1 * carrier.put('k2', 'v2')
244+
1 * carrier.put('ot-baggage-t0', "${(long) (mockedContext.endToEndStartTime / 1000000L)}")
245+
246+
cleanup:
247+
tracer.close()
248+
}
249+
221250
static DDSpanContext mockedContext(CoreTracer tracer, DDTraceId traceId, long spanId, int samplingPriority, String origin, Map<String, String> baggage) {
222251
return new DDSpanContext(
223252
traceId,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import static datadog.trace.core.propagation.W3CHttpCodec.newInjector
2121

2222
class W3CHttpInjectorTest extends DDCoreSpecification {
2323

24-
HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"])
24+
HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"], true)
2525

2626
def "inject http headers"() {
2727
setup:

0 commit comments

Comments
 (0)