Skip to content

Commit 85d573f

Browse files
mccullsmcculls
committed
Move W3C baggage tag injection to just before the span is serialized
1 parent 2d45a1f commit 85d573f

File tree

7 files changed

+70
-87
lines changed

7 files changed

+70
-87
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java‎

Lines changed: 1 addition & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2727
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
2828
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
29-
import datadog.trace.bootstrap.instrumentation.api.Baggage;
3029
import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
3130
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
3231
import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
@@ -38,7 +37,6 @@
3837
import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver;
3938
import java.net.InetAddress;
4039
import java.util.BitSet;
41-
import java.util.List;
4240
import java.util.Locale;
4341
import java.util.Map;
4442
import java.util.concurrent.ExecutionException;
@@ -151,17 +149,7 @@ public AgentSpan startSpan(REQUEST_CARRIER carrier, Context context) {
151149

152150
public AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
153151
AgentSpan extractedSpan = AgentSpan.fromContext(context);
154-
AgentSpanContext.Extracted extractedContext = null;
155-
if (extractedSpan != null) {
156-
extractedContext = (AgentSpanContext.Extracted) extractedSpan.context();
157-
if (extractedContext instanceof TagContext) {
158-
Baggage baggage = Baggage.fromContext(context);
159-
if (baggage != null) {
160-
setBaggageTags((TagContext) extractedContext, baggage.asMap());
161-
}
162-
}
163-
}
164-
return extractedContext;
152+
return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
165153
}
166154

167155
public AgentSpan onRequest(
@@ -656,24 +644,4 @@ public boolean accept(String key, String value) {
656644
return true;
657645
}
658646
}
659-
660-
static void setBaggageTags(TagContext tagContext, Map<String, String> baggage) {
661-
List<String> baggageTagKeys = Config.get().getTraceBaggageTagKeys();
662-
if (baggageTagKeys.isEmpty()) {
663-
return;
664-
}
665-
for (String key : baggageTagKeys) {
666-
if ("*".equals(key)) {
667-
// If the key is "*", we add all baggage items
668-
for (Map.Entry<String, String> entry : baggage.entrySet()) {
669-
tagContext.putTagIfAbsent("baggage." + entry.getKey(), entry.getValue());
670-
}
671-
break;
672-
}
673-
String value = baggage.get(key);
674-
if (value != null) {
675-
tagContext.putTagIfAbsent("baggage." + key, value);
676-
}
677-
}
678-
}
679647
}

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy‎

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI
1717
import datadog.trace.bootstrap.instrumentation.api.ContextVisitors
1818
import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities
1919
import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities
20-
import datadog.trace.bootstrap.instrumentation.api.TagContext
2120
import datadog.trace.bootstrap.instrumentation.api.Tags
2221
import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
2322
import datadog.trace.bootstrap.instrumentation.api.URIDefaultDataAdapter
@@ -30,7 +29,6 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_DE
3029
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_QUERY_STRING
3130
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_RAW_RESOURCE
3231
import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_SERVER_TAG_QUERY_STRING
33-
import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_TAG_KEYS
3432
import static datadog.trace.api.gateway.Events.EVENTS
3533

3634
class HttpServerDecoratorTest extends ServerDecoratorTest {
@@ -568,54 +566,4 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
568566
return reqHeaderDoneCount
569567
}
570568
}
571-
572-
def "test setBaggageTags default"() {
573-
setup:
574-
def tags = [:]
575-
576-
when:
577-
def tagContext = new TagContext()
578-
HttpServerDecorator.setBaggageTags(tagContext, ["user.id": "doggo", "account.id": "test", "session.id": "1234", "region": "us -east - 1"])
579-
tagContext.tags.fillMap(tags)
580-
581-
then:
582-
tags == ["baggage.user.id": "doggo", "baggage.account.id": "test", "baggage.session.id": "1234"]
583-
}
584-
585-
def "test setBaggageTags does not overwrite"() {
586-
setup:
587-
def tags = [:]
588-
589-
when:
590-
def tagContext = new TagContext()
591-
tagContext.putTag("baggage.account.id", "staging")
592-
HttpServerDecorator.setBaggageTags(tagContext, ["user.id": "doggo", "account.id": "test", "session.id": "1234", "region": "us -east - 1"])
593-
tagContext.tags.fillMap(tags)
594-
595-
then:
596-
tags == ["baggage.user.id": "doggo", "baggage.account.id": "staging", "baggage.session.id": "1234"]
597-
}
598-
599-
def "test setBaggageTags with different configurations"() {
600-
setup:
601-
injectSysConfig(TRACE_BAGGAGE_TAG_KEYS, baggageTagKeysConfig)
602-
def tags = [:]
603-
604-
when:
605-
def tagContext = new TagContext()
606-
HttpServerDecorator.setBaggageTags(tagContext, ["user.id": "doggo", "foo": "bar", "language": "en"])
607-
tagContext.tags.fillMap(tags)
608-
609-
then:
610-
tags == expectedTags
611-
612-
where:
613-
baggageTagKeysConfig | expectedTags
614-
"*" | ["baggage.user.id": "doggo", "baggage.foo": "bar", "baggage.language": "en"]
615-
" " | [:]
616-
"system.id" | [:]
617-
"user.id" | ["baggage.user.id": "doggo"]
618-
"foo" | ["baggage.foo": "bar"]
619-
"user.id,foo" | ["baggage.user.id": "doggo", "baggage.foo": "bar"]
620-
}
621569
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
6060
import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
6161
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
62+
import datadog.trace.bootstrap.instrumentation.api.Baggage;
6263
import datadog.trace.bootstrap.instrumentation.api.BlackHoleSpan;
6364
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
6465
import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
@@ -1591,6 +1592,7 @@ private DDSpanContext buildSpanContext() {
15911592
final long spanId;
15921593
final long parentSpanId;
15931594
final Map<String, String> baggage;
1595+
final Baggage w3cBaggage;
15941596
final TraceCollector parentTraceCollector;
15951597
final int samplingPriority;
15961598
final CharSequence origin;
@@ -1645,6 +1647,7 @@ private DDSpanContext buildSpanContext() {
16451647
traceId = ddsc.getTraceId();
16461648
parentSpanId = ddsc.getSpanId();
16471649
baggage = ddsc.getBaggageItems();
1650+
w3cBaggage = null;
16481651
parentTraceCollector = ddsc.getTraceCollector();
16491652
samplingPriority = PrioritySampling.UNSET;
16501653
origin = null;
@@ -1706,6 +1709,7 @@ private DDSpanContext buildSpanContext() {
17061709
coreTagsNeedsIntercept = true; // maybe intercept isn't needed?
17071710
origin = tc.getOrigin();
17081711
baggage = tc.getBaggage();
1712+
w3cBaggage = tc.getW3CBaggage();
17091713
requestContextDataAppSec = tc.getRequestContextDataAppSec();
17101714
requestContextDataIast = tc.getRequestContextDataIast();
17111715
ciVisibilityContextData = tc.getCiVisibilityContextData();
@@ -1715,6 +1719,7 @@ private DDSpanContext buildSpanContext() {
17151719
coreTagsNeedsIntercept = false;
17161720
origin = null;
17171721
baggage = null;
1722+
w3cBaggage = null;
17181723
requestContextDataAppSec = null;
17191724
requestContextDataIast = null;
17201725
ciVisibilityContextData = null;
@@ -1805,6 +1810,7 @@ private DDSpanContext buildSpanContext() {
18051810
samplingPriority,
18061811
origin,
18071812
baggage,
1813+
w3cBaggage,
18081814
errorFlag,
18091815
spanType,
18101816
tagsSize,

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import static datadog.trace.api.cache.RadixTreeCache.HTTP_STATUSES;
66
import static datadog.trace.bootstrap.instrumentation.api.ErrorPriorities.UNSET;
77

8+
import datadog.trace.api.Config;
89
import datadog.trace.api.DDSpanId;
910
import datadog.trace.api.DDTags;
1011
import datadog.trace.api.DDTraceId;
@@ -23,6 +24,7 @@
2324
import datadog.trace.api.sampling.SamplingMechanism;
2425
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
2526
import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
27+
import datadog.trace.bootstrap.instrumentation.api.Baggage;
2628
import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
2729
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
2830
import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities;
@@ -77,6 +79,8 @@ public class DDSpanContext
7779
/** Baggage is associated with the whole trace and shared with other spans */
7880
private volatile Map<String, String> baggageItems;
7981

82+
private final Baggage w3cBaggage;
83+
8084
// Not Shared with other span contexts
8185
private final DDTraceId traceId;
8286
private final long spanId;
@@ -187,6 +191,7 @@ public DDSpanContext(
187191
samplingPriority,
188192
origin,
189193
baggageItems,
194+
null,
190195
errorFlag,
191196
spanType,
192197
tagsSize,
@@ -233,6 +238,7 @@ public DDSpanContext(
233238
samplingPriority,
234239
origin,
235240
baggageItems,
241+
null,
236242
errorFlag,
237243
spanType,
238244
tagsSize,
@@ -279,6 +285,7 @@ public DDSpanContext(
279285
samplingPriority,
280286
origin,
281287
baggageItems,
288+
null,
282289
errorFlag,
283290
spanType,
284291
tagsSize,
@@ -304,6 +311,7 @@ public DDSpanContext(
304311
final int samplingPriority,
305312
final CharSequence origin,
306313
final Map<String, String> baggageItems,
314+
final Baggage w3cBaggage,
307315
final boolean errorFlag,
308316
final CharSequence spanType,
309317
final int tagsSize,
@@ -331,6 +339,7 @@ public DDSpanContext(
331339
} else {
332340
this.baggageItems = new ConcurrentHashMap<>(baggageItems);
333341
}
342+
this.w3cBaggage = w3cBaggage;
334343

335344
this.requestContextDataAppSec = requestContextDataAppSec;
336345
this.requestContextDataIast = requestContextDataIast;
@@ -926,6 +935,9 @@ public void processTagsAndBaggage(
926935
Map<String, String> baggageItemsWithPropagationTags;
927936
if (injectBaggageAsTags) {
928937
baggageItemsWithPropagationTags = new HashMap<>(baggageItems);
938+
if (w3cBaggage != null) {
939+
injectW3CBaggageTags(baggageItemsWithPropagationTags);
940+
}
929941
propagationTags.fillTagMap(baggageItemsWithPropagationTags);
930942
} else {
931943
baggageItemsWithPropagationTags = propagationTags.createTagMap();
@@ -948,6 +960,27 @@ public void processTagsAndBaggage(
948960
}
949961
}
950962

963+
void injectW3CBaggageTags(Map<String, String> baggageItemsWithPropagationTags) {
964+
List<String> baggageTagKeys = Config.get().getTraceBaggageTagKeys();
965+
if (baggageTagKeys.isEmpty()) {
966+
return;
967+
}
968+
Map<String, String> w3cBaggageItems = w3cBaggage.asMap();
969+
for (String key : baggageTagKeys) {
970+
if ("*".equals(key)) {
971+
// If the key is "*", we add all baggage items
972+
for (Map.Entry<String, String> entry : w3cBaggageItems.entrySet()) {
973+
baggageItemsWithPropagationTags.put("baggage." + entry.getKey(), entry.getValue());
974+
}
975+
break;
976+
}
977+
String value = w3cBaggageItems.get(key);
978+
if (value != null) {
979+
baggageItemsWithPropagationTags.put("baggage." + key, value);
980+
}
981+
}
982+
}
983+
951984
@Override
952985
public void setIntegrationName(CharSequence integrationName) {
953986
this.integrationName = integrationName;

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
import datadog.context.propagation.CarrierVisitor;
88
import datadog.context.propagation.Propagator;
99
import datadog.trace.api.Config;
10+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
1012
import datadog.trace.bootstrap.instrumentation.api.Baggage;
13+
import datadog.trace.bootstrap.instrumentation.api.TagContext;
1114
import datadog.trace.core.util.PercentEscaper;
1215
import datadog.trace.core.util.PercentEscaper.Escaped;
1316
import java.io.UnsupportedEncodingException;
@@ -109,7 +112,21 @@ public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
109112
}
110113
BaggageExtractor baggageExtractor = new BaggageExtractor();
111114
visitor.forEachKeyValue(carrier, baggageExtractor);
112-
return baggageExtractor.extracted == null ? context : context.with(baggageExtractor.extracted);
115+
Baggage baggage = baggageExtractor.extracted;
116+
if (baggage == null) {
117+
return context;
118+
}
119+
120+
// TODO: consider a better way to link baggage with the extracted (legacy) TagContext
121+
AgentSpan extractedSpan = AgentSpan.fromContext(context);
122+
if (extractedSpan != null) {
123+
AgentSpanContext extractedSpanContext = extractedSpan.context();
124+
if (extractedSpanContext instanceof TagContext) {
125+
((TagContext) extractedSpanContext).setW3CBaggage(baggage);
126+
}
127+
}
128+
129+
return context.with(baggage);
113130
}
114131

115132
private class BaggageExtractor implements BiConsumer<String, String> {

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
public final class AgentPropagation {
1616
public static final Concern TRACING_CONCERN = named("tracing");
17-
public static final Concern BAGGAGE_CONCERN = named("baggage");
17+
// TODO: Baggage propagator should run after tracing so it can link baggage with the span context
18+
// TODO: remove this priority once we have a story for replacing TagContext with the Context API
19+
public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105);
1820
public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray");
1921

2022
// TODO DSM propagator should run after the other propagators as it stores the pathway context

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ public class TagContext implements AgentSpanContext.Extracted {
3232
private PathwayContext pathwayContext;
3333
private final HttpHeaders httpHeaders;
3434
private final Map<String, String> baggage;
35+
private Baggage w3cBaggage;
3536
private final int samplingPriority;
3637
private final TraceConfig traceConfig;
3738
private final TracePropagationStyle propagationStyle;
@@ -198,6 +199,14 @@ public Iterable<Map.Entry<String, String>> baggageItems() {
198199
return baggage.entrySet();
199200
}
200201

202+
public final Baggage getW3CBaggage() {
203+
return w3cBaggage;
204+
}
205+
206+
public void setW3CBaggage(Baggage w3cBaggage) {
207+
this.w3cBaggage = w3cBaggage;
208+
}
209+
201210
@Override
202211
public DDTraceId getTraceId() {
203212
return traceId;

0 commit comments

Comments
 (0)