Skip to content

Commit ae71f51

Browse files
mhliddmhlidd
committed
initial implementation and unit tests for inject/extract
1 parent 8d0bb34 commit ae71f51

File tree

8 files changed

+425
-1
lines changed

8 files changed

+425
-1
lines changed

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.api;
22

3+
import static datadog.trace.api.TracePropagationStyle.BAGGAGE;
34
import static datadog.trace.api.TracePropagationStyle.DATADOG;
45
import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
56
import static java.util.Arrays.asList;
@@ -78,9 +79,11 @@ public final class ConfigDefaults {
7879
static final int DEFAULT_PARTIAL_FLUSH_MIN_SPANS = 1000;
7980
static final boolean DEFAULT_PROPAGATION_EXTRACT_LOG_HEADER_NAMES_ENABLED = false;
8081
static final Set<TracePropagationStyle> DEFAULT_TRACE_PROPAGATION_STYLE =
81-
new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT));
82+
new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT, BAGGAGE));
8283
static final Set<PropagationStyle> DEFAULT_PROPAGATION_STYLE =
8384
new LinkedHashSet<>(asList(PropagationStyle.DATADOG));
85+
static final int DEFAULT_TRACE_BAGGAGE_MAX_ITEMS = 64;
86+
static final int DEFAULT_TRACE_BAGGAGE_MAX_BYTES = 8192;
8487
static final boolean DEFAULT_JMX_FETCH_ENABLED = true;
8588
static final boolean DEFAULT_TRACE_AGENT_V05_ENABLED = false;
8689

‎dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ public enum TracePropagationStyle {
2121
// W3C trace context propagation style
2222
// https://www.w3.org/TR/trace-context-1/
2323
TRACECONTEXT,
24+
// OTEL baggage support
25+
// https://www.w3.org/TR/baggage/
26+
BAGGAGE,
2427
// None does not extract or inject
2528
NONE;
2629

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ public final class TracerConfig {
9191
public static final String TRACE_PROPAGATION_STYLE_EXTRACT = "trace.propagation.style.extract";
9292
public static final String TRACE_PROPAGATION_STYLE_INJECT = "trace.propagation.style.inject";
9393
public static final String TRACE_PROPAGATION_EXTRACT_FIRST = "trace.propagation.extract.first";
94+
public static final String TRACE_BAGGAGE_MAX_ITEMS = "trace.baggage.max.items";
95+
public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes";
9496

9597
public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled";
9698

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

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
package datadog.trace.core.propagation;
2+
3+
import static datadog.trace.api.TracePropagationStyle.BAGGAGE;
4+
5+
import datadog.trace.api.Config;
6+
import datadog.trace.api.TraceConfig;
7+
import datadog.trace.api.TracePropagationStyle;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
9+
import datadog.trace.bootstrap.instrumentation.api.TagContext;
10+
import datadog.trace.core.DDSpanContext;
11+
import java.nio.charset.StandardCharsets;
12+
import java.util.Collections;
13+
import java.util.HashMap;
14+
import java.util.Map;
15+
import java.util.function.Supplier;
16+
import org.slf4j.Logger;
17+
import org.slf4j.LoggerFactory;
18+
19+
/** A codec designed for HTTP transport via headers using Datadog headers */
20+
class BaggageHttpCodec {
21+
private static final Logger log = LoggerFactory.getLogger(BaggageHttpCodec.class);
22+
23+
static final String BAGGAGE_KEY = "baggage";
24+
private static final int MAX_CHARACTER_SIZE = 4;
25+
26+
private BaggageHttpCodec() {
27+
// This class should not be created. This also makes code coverage checks happy.
28+
}
29+
30+
public static HttpCodec.Injector newInjector(Map<String, String> invertedBaggageMapping) {
31+
return new Injector(invertedBaggageMapping);
32+
}
33+
34+
private static class Injector implements HttpCodec.Injector {
35+
36+
private final Map<String, String> invertedBaggageMapping;
37+
38+
public Injector(Map<String, String> invertedBaggageMapping) {
39+
assert invertedBaggageMapping != null;
40+
this.invertedBaggageMapping = invertedBaggageMapping;
41+
}
42+
43+
@Override
44+
public <C> void inject(
45+
final DDSpanContext context, final C carrier, final AgentPropagation.Setter<C> setter) {
46+
Config config = Config.get();
47+
48+
StringBuilder baggageText = new StringBuilder();
49+
int processedBaggage = 0;
50+
int currentBytes = 0;
51+
int maxItems = config.getTraceBaggageMaxItems();
52+
int maxBytes = config.getTraceBaggageMaxBytes();
53+
int currentCharacters = 0;
54+
int maxSafeCharacters = maxBytes / MAX_CHARACTER_SIZE;
55+
for (final Map.Entry<String, String> entry : context.baggageItems()) {
56+
if (processedBaggage >= maxItems) {
57+
break;
58+
}
59+
StringBuilder currentText = new StringBuilder();
60+
if (processedBaggage != 0) {
61+
currentText.append(',');
62+
}
63+
64+
currentText.append(HttpCodec.encodeBaggage(entry.getKey()));
65+
currentText.append('=');
66+
currentText.append(HttpCodec.encodeBaggage(entry.getValue()));
67+
68+
// worst case check
69+
if (currentCharacters + currentText.length() <= maxSafeCharacters) {
70+
currentCharacters += currentText.length();
71+
} else {
72+
if (currentBytes
73+
== 0) { // special case to calculate byte size after surpassing worst-case number of
74+
// characters
75+
currentBytes = baggageText.toString().getBytes(StandardCharsets.UTF_8).length;
76+
}
77+
int byteSize =
78+
currentText
79+
.toString()
80+
.getBytes(StandardCharsets.UTF_8)
81+
.length; // find largest possible byte size for UTF encoded characters and only do
82+
// size checking after we hit this worst case scenario
83+
if (byteSize + currentBytes > maxBytes) {
84+
break;
85+
}
86+
currentBytes += byteSize;
87+
}
88+
baggageText.append(currentText);
89+
processedBaggage++;
90+
}
91+
92+
setter.set(carrier, BAGGAGE_KEY, baggageText.toString());
93+
}
94+
}
95+
96+
public static HttpCodec.Extractor newExtractor(
97+
Config config, Supplier<TraceConfig> traceConfigSupplier) {
98+
return new TagContextExtractor(
99+
traceConfigSupplier, () -> new BaggageContextInterpreter(config));
100+
}
101+
102+
private static class BaggageContextInterpreter extends ContextInterpreter {
103+
104+
private BaggageContextInterpreter(Config config) {
105+
super(config);
106+
}
107+
108+
@Override
109+
public TracePropagationStyle style() {
110+
return BAGGAGE;
111+
}
112+
113+
private Map<String, String> parseBaggageHeaders(String input) {
114+
Map<String, String> baggage = new HashMap<>();
115+
char keyValueSeparator = '=';
116+
char pairSeparator = ',';
117+
int start = 0;
118+
119+
int pairSeparatorInd = input.indexOf(pairSeparator);
120+
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
121+
int kvSeparatorInd = input.indexOf(keyValueSeparator);
122+
while (kvSeparatorInd != -1) {
123+
int end = pairSeparatorInd;
124+
if (kvSeparatorInd > end) { // value is missing
125+
return Collections.emptyMap();
126+
}
127+
String key = HttpCodec.decode(input.substring(start, kvSeparatorInd).trim());
128+
String value = HttpCodec.decode(input.substring(kvSeparatorInd + 1, end).trim());
129+
if (key.isEmpty() || value.isEmpty()) {
130+
return Collections.emptyMap();
131+
}
132+
baggage.put(key, value);
133+
134+
kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1);
135+
pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1);
136+
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
137+
start = end + 1;
138+
}
139+
return baggage;
140+
}
141+
142+
@Override
143+
public boolean accept(String key, String value) {
144+
if (null == key || key.isEmpty()) {
145+
return true;
146+
}
147+
if (LOG_EXTRACT_HEADER_NAMES) {
148+
log.debug("Header: {}", key);
149+
}
150+
151+
if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage
152+
baggage = parseBaggageHeaders(value);
153+
}
154+
return true;
155+
}
156+
157+
@Override
158+
protected TagContext build() {
159+
return super.build();
160+
}
161+
}
162+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,9 @@ private static Map<TracePropagationStyle, Injector> createInjectors(
125125
case TRACECONTEXT:
126126
result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping));
127127
break;
128+
case BAGGAGE:
129+
result.put(style, BaggageHttpCodec.newInjector(reverseBaggageMapping));
130+
break;
128131
default:
129132
log.debug("No implementation found to inject propagation style: {}", style);
130133
break;
@@ -159,6 +162,8 @@ public static Extractor createExtractor(
159162
case TRACECONTEXT:
160163
extractors.add(W3CHttpCodec.newExtractor(config, traceConfigSupplier));
161164
break;
165+
case BAGGAGE:
166+
extractors.add(BaggageHttpCodec.newExtractor(config, traceConfigSupplier));
162167
default:
163168
log.debug("No implementation found to extract propagation style: {}", style);
164169
break;

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package datadog.trace.core.propagation
2+
3+
import datadog.trace.api.Config
4+
import datadog.trace.api.DynamicConfig
5+
import datadog.trace.bootstrap.instrumentation.api.TagContext
6+
import datadog.trace.bootstrap.instrumentation.api.ContextVisitors
7+
import datadog.trace.test.util.DDSpecification
8+
import static datadog.trace.core.propagation.BaggageHttpCodec.BAGGAGE_KEY
9+
10+
class BaggageHttpExtractorTest extends DDSpecification {
11+
12+
private DynamicConfig dynamicConfig
13+
private HttpCodec.Extractor _extractor
14+
15+
private HttpCodec.Extractor getExtractor() {
16+
_extractor ?: (_extractor = createExtractor(Config.get()))
17+
}
18+
19+
private HttpCodec.Extractor createExtractor(Config config) {
20+
BaggageHttpCodec.newExtractor(config, { dynamicConfig.captureTraceConfig() })
21+
}
22+
23+
void setup() {
24+
dynamicConfig = DynamicConfig.create()
25+
.apply()
26+
}
27+
28+
void cleanup() {
29+
extractor.cleanup()
30+
}
31+
32+
def "extract valid baggage headers"() {
33+
setup:
34+
def extractor = createExtractor(Config.get())
35+
def headers = [
36+
(BAGGAGE_KEY) : baggageHeader,
37+
]
38+
39+
when:
40+
final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap())
41+
42+
then:
43+
context.baggage == baggageMap
44+
45+
cleanup:
46+
extractor.cleanup()
47+
48+
where:
49+
baggageHeader | baggageMap
50+
"key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"]
51+
"%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\']
52+
}
53+
54+
def "extract invalid baggage headers"() {
55+
setup:
56+
def extractor = createExtractor(Config.get())
57+
def headers = [
58+
(BAGGAGE_KEY) : baggageHeader,
59+
]
60+
61+
when:
62+
final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap())
63+
64+
then:
65+
context == null
66+
67+
cleanup:
68+
extractor.cleanup()
69+
70+
where:
71+
baggageHeader | baggageMap
72+
"no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | []
73+
"foo=gets-dropped-because-subsequent-pair-is-malformed,=" | []
74+
"=no-key" | []
75+
"no-value=" | []
76+
}
77+
}

0 commit comments

Comments
 (0)