Skip to content

Commit f8c6b49

Browse files
PerfectSlayerPerfectSlayer
authored
Improve Baggage API (#8523)
* feat(internal-api): Improve Baggage API * feat(internal-api): Add Baggage unit tests
1 parent 1dff23c commit f8c6b49

File tree

8 files changed

+295
-192
lines changed

8 files changed

+295
-192
lines changed

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

+47-51
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
package datadog.trace.core.baggage;
22

3+
import static java.util.Collections.emptyMap;
4+
35
import datadog.context.Context;
46
import datadog.context.propagation.CarrierSetter;
57
import datadog.context.propagation.CarrierVisitor;
68
import datadog.context.propagation.Propagator;
79
import datadog.trace.api.Config;
8-
import datadog.trace.bootstrap.instrumentation.api.BaggageContext;
9-
import datadog.trace.core.util.EscapedData;
10+
import datadog.trace.bootstrap.instrumentation.api.Baggage;
1011
import datadog.trace.core.util.PercentEscaper;
12+
import datadog.trace.core.util.PercentEscaper.Escaped;
1113
import java.io.UnsupportedEncodingException;
1214
import java.net.URLDecoder;
13-
import java.util.Collections;
1415
import java.util.HashMap;
1516
import java.util.Map;
1617
import java.util.function.BiConsumer;
@@ -20,7 +21,7 @@
2021

2122
@ParametersAreNonnullByDefault
2223
public class BaggagePropagator implements Propagator {
23-
private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class);
24+
private static final Logger LOG = LoggerFactory.getLogger(BaggagePropagator.class);
2425
private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create();
2526
static final String BAGGAGE_KEY = "baggage";
2627
private final Config config;
@@ -37,13 +38,13 @@ public BaggagePropagator(Config config) {
3738
public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) {
3839
this.injectBaggage = injectBaggage;
3940
this.extractBaggage = extractBaggage;
40-
config = Config.get();
41+
this.config = Config.get();
4142
}
4243

4344
@Override
4445
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
45-
int maxItems = config.getTraceBaggageMaxItems();
46-
int maxBytes = config.getTraceBaggageMaxBytes();
46+
int maxItems = this.config.getTraceBaggageMaxItems();
47+
int maxBytes = this.config.getTraceBaggageMaxBytes();
4748
//noinspection ConstantValue
4849
if (!this.injectBaggage
4950
|| maxItems == 0
@@ -54,52 +55,52 @@ public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
5455
return;
5556
}
5657

57-
BaggageContext baggageContext = BaggageContext.fromContext(context);
58-
if (baggageContext == null) {
59-
log.debug("BaggageContext instance is missing from the following context {}", context);
58+
Baggage baggage = Baggage.fromContext(context);
59+
if (baggage == null) {
60+
LOG.debug("Baggage instance is missing from the following context {}", context);
6061
return;
6162
}
6263

63-
String baggageHeader = baggageContext.getW3cBaggageHeader();
64-
if (baggageHeader != null) {
65-
setter.set(carrier, BAGGAGE_KEY, baggageHeader);
64+
String headerValue = baggage.getW3cHeader();
65+
if (headerValue != null) {
66+
setter.set(carrier, BAGGAGE_KEY, headerValue);
6667
return;
6768
}
6869

69-
int processedBaggage = 0;
70+
int processedItems = 0;
7071
int currentBytes = 0;
7172
StringBuilder baggageText = new StringBuilder();
72-
for (final Map.Entry<String, String> entry : baggageContext.asMap().entrySet()) {
73+
for (final Map.Entry<String, String> entry : baggage.asMap().entrySet()) {
7374
// if there are already baggage items processed, add and allocate bytes for a comma
7475
int extraBytes = 1;
75-
if (processedBaggage != 0) {
76+
if (processedItems != 0) {
7677
baggageText.append(',');
7778
extraBytes++;
7879
}
7980

80-
EscapedData escapedKey = UTF_ESCAPER.escapeKey(entry.getKey());
81-
EscapedData escapedVal = UTF_ESCAPER.escapeValue(entry.getValue());
81+
Escaped escapedKey = UTF_ESCAPER.escapeKey(entry.getKey());
82+
Escaped escapedVal = UTF_ESCAPER.escapeValue(entry.getValue());
8283

83-
baggageText.append(escapedKey.getData());
84+
baggageText.append(escapedKey.data);
8485
baggageText.append('=');
85-
baggageText.append(escapedVal.getData());
86+
baggageText.append(escapedVal.data);
8687

87-
processedBaggage++;
88+
processedItems++;
8889
// reached the max number of baggage items allowed
89-
if (processedBaggage == maxItems) {
90+
if (processedItems == maxItems) {
9091
break;
9192
}
9293
// Drop newest k/v pair if adding it leads to exceeding the limit
93-
if (currentBytes + escapedKey.getSize() + escapedVal.getSize() + extraBytes > maxBytes) {
94+
if (currentBytes + escapedKey.size + escapedVal.size + extraBytes > maxBytes) {
9495
baggageText.setLength(currentBytes);
9596
break;
9697
}
97-
currentBytes += escapedKey.getSize() + escapedVal.getSize() + extraBytes;
98+
currentBytes += escapedKey.size + escapedVal.size + extraBytes;
9899
}
99100

100-
String baggageString = baggageText.toString();
101-
baggageContext.setW3cBaggageHeader(baggageString);
102-
setter.set(carrier, BAGGAGE_KEY, baggageString);
101+
headerValue = baggageText.toString();
102+
baggage.setW3cHeader(headerValue);
103+
setter.set(carrier, BAGGAGE_KEY, headerValue);
103104
}
104105

105106
@Override
@@ -108,57 +109,52 @@ public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
108109
if (!this.extractBaggage || context == null || carrier == null || visitor == null) {
109110
return context;
110111
}
111-
BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor();
112-
visitor.forEachKeyValue(carrier, baggageContextExtractor);
113-
BaggageContext extractedContext = baggageContextExtractor.extractedContext;
114-
if (extractedContext == null) {
115-
return context;
116-
}
117-
return extractedContext.storeInto(context);
112+
BaggageExtractor baggageExtractor = new BaggageExtractor();
113+
visitor.forEachKeyValue(carrier, baggageExtractor);
114+
return baggageExtractor.extracted == null ? context : context.with(baggageExtractor.extracted);
118115
}
119116

120-
public static class BaggageContextExtractor implements BiConsumer<String, String> {
121-
private BaggageContext extractedContext;
117+
private static class BaggageExtractor implements BiConsumer<String, String> {
118+
private static final char KEY_VALUE_SEPARATOR = '=';
119+
private static final char PAIR_SEPARATOR = ',';
120+
private Baggage extracted;
122121

123-
BaggageContextExtractor() {}
122+
private BaggageExtractor() {}
124123

125124
/** URL decode value */
126125
private String decode(final String value) {
127126
String decoded = value;
128127
try {
129128
decoded = URLDecoder.decode(value, "UTF-8");
130129
} catch (final UnsupportedEncodingException | IllegalArgumentException e) {
131-
log.debug("Failed to decode {}", value);
130+
LOG.debug("Failed to decode {}", value);
132131
}
133132
return decoded;
134133
}
135134

136135
private Map<String, String> parseBaggageHeaders(String input) {
137136
Map<String, String> baggage = new HashMap<>();
138-
char keyValueSeparator = '=';
139-
char pairSeparator = ',';
140137
int start = 0;
141-
142-
int pairSeparatorInd = input.indexOf(pairSeparator);
138+
int pairSeparatorInd = input.indexOf(PAIR_SEPARATOR);
143139
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
144-
int kvSeparatorInd = input.indexOf(keyValueSeparator);
140+
int kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR);
145141
while (kvSeparatorInd != -1) {
146142
int end = pairSeparatorInd;
147143
if (kvSeparatorInd > end) {
148-
log.debug(
144+
LOG.debug(
149145
"Dropping baggage headers due to key with no value {}", input.substring(start, end));
150-
return Collections.emptyMap();
146+
return emptyMap();
151147
}
152148
String key = decode(input.substring(start, kvSeparatorInd).trim());
153149
String value = decode(input.substring(kvSeparatorInd + 1, end).trim());
154150
if (key.isEmpty() || value.isEmpty()) {
155-
log.debug("Dropping baggage headers due to empty k/v {}:{}", key, value);
156-
return Collections.emptyMap();
151+
LOG.debug("Dropping baggage headers due to empty k/v {}:{}", key, value);
152+
return emptyMap();
157153
}
158154
baggage.put(key, value);
159155

160-
kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1);
161-
pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1);
156+
kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR, pairSeparatorInd + 1);
157+
pairSeparatorInd = input.indexOf(PAIR_SEPARATOR, pairSeparatorInd + 1);
162158
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
163159
start = end + 1;
164160
}
@@ -168,10 +164,10 @@ private Map<String, String> parseBaggageHeaders(String input) {
168164
@Override
169165
public void accept(String key, String value) {
170166
// Only process tags that are relevant to baggage
171-
if (key != null && key.equalsIgnoreCase(BAGGAGE_KEY)) {
167+
if (BAGGAGE_KEY.equalsIgnoreCase(key)) {
172168
Map<String, String> baggage = parseBaggageHeaders(value);
173169
if (!baggage.isEmpty()) {
174-
extractedContext = BaggageContext.create(baggage, value);
170+
this.extracted = Baggage.create(baggage, value);
175171
}
176172
}
177173
}

‎dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java

-36
This file was deleted.

‎dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java

+26-16
Original file line numberDiff line numberDiff line change
@@ -107,25 +107,24 @@ private static boolean[] createUnsafeOctets(String safeChars) {
107107
return octets;
108108
}
109109

110-
public EscapedData escapeKey(String s) {
110+
public Escaped escapeKey(String s) {
111111
return escape(s, unsafeKeyOctets);
112112
}
113113

114-
public EscapedData escapeValue(String s) {
114+
public Escaped escapeValue(String s) {
115115
return escape(s, unsafeValOctets);
116116
}
117117

118118
/** Escape the provided String, using percent-style URL Encoding. */
119-
public EscapedData escape(String s, boolean[] unsafeOctets) {
120-
int size = 0;
119+
public Escaped escape(String s, boolean[] unsafeOctets) {
121120
int slen = s.length();
122121
for (int index = 0; index < slen; index++) {
123122
char c = s.charAt(index);
124123
if (c > '~' || c <= ' ' || c <= unsafeOctets.length && unsafeOctets[c]) {
125124
return escapeSlow(s, index, unsafeOctets);
126125
}
127126
}
128-
return new EscapedData(s, slen);
127+
return new Escaped(s, slen);
129128
}
130129

131130
/*
@@ -147,14 +146,14 @@ public EscapedData escape(String s, boolean[] unsafeOctets) {
147146
* @throws NullPointerException if {@code string} is null
148147
* @throws IllegalArgumentException if invalid surrogate characters are encountered
149148
*/
150-
private static EscapedData escapeSlow(String s, int index, boolean[] unsafeOctets) {
149+
private static Escaped escapeSlow(String s, int index, boolean[] unsafeOctets) {
151150
int end = s.length();
152151

153152
// Get a destination buffer and setup some loop variables.
154153
char[] dest = new char[1024]; // 1024 from the original guava source
155154
int destIndex = 0;
156155
int unescapedChunkStart = 0;
157-
EscapedData data = new EscapedData("", index);
156+
Escaped result = new Escaped("", index);
158157

159158
while (index < end) {
160159
int cp = codePointAt(s, index, end);
@@ -164,7 +163,7 @@ private static EscapedData escapeSlow(String s, int index, boolean[] unsafeOctet
164163
// It is possible for this to return null because nextEscapeIndex() may
165164
// (for performance reasons) yield some false positives but it must never
166165
// give false negatives.
167-
char[] escaped = escape(cp, data, unsafeOctets);
166+
char[] escaped = escape(cp, result, unsafeOctets);
168167
int nextIndex = index + (Character.isSupplementaryCodePoint(cp) ? 2 : 1);
169168
if (escaped != null) {
170169
int charsSkipped = index - unescapedChunkStart;
@@ -202,10 +201,11 @@ private static EscapedData escapeSlow(String s, int index, boolean[] unsafeOctet
202201
s.getChars(unescapedChunkStart, end, dest, destIndex);
203202
destIndex = endIndex;
204203
}
205-
data.addSize(charsSkipped); // Adding characters in-between characters that want to be encoded
204+
// Adding characters in-between characters that want to be encoded
205+
result.size += charsSkipped;
206206

207-
data.setData(new String(dest, 0, destIndex));
208-
return data;
207+
result.data = new String(dest, 0, destIndex);
208+
return result;
209209
}
210210

211211
private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean[] unsafeOctets) {
@@ -221,7 +221,7 @@ private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean
221221
/** Escapes the given Unicode code point in UTF-8. */
222222
@CheckForNull
223223
@SuppressWarnings("UngroupedOverloads")
224-
private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
224+
private static char[] escape(int cp, Escaped escaped, boolean[] unsafeOctets) {
225225
// We should never get negative values here but if we do it will throw an
226226
// IndexOutOfBoundsException, so at least it will get spotted.
227227
if (cp < unsafeOctets.length && !unsafeOctets[cp]) {
@@ -233,7 +233,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
233233
dest[0] = '%';
234234
dest[2] = UPPER_HEX_DIGITS[cp & 0xF];
235235
dest[1] = UPPER_HEX_DIGITS[cp >>> 4];
236-
data.incrementSize();
236+
escaped.size++;
237237
return dest;
238238
} else if (cp <= 0x7ff) {
239239
// Two byte UTF-8 characters [cp >= 0x80 && cp <= 0x7ff]
@@ -248,7 +248,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
248248
dest[2] = UPPER_HEX_DIGITS[cp & 0xF];
249249
cp >>>= 4;
250250
dest[1] = UPPER_HEX_DIGITS[0xC | cp];
251-
data.addSize(2);
251+
escaped.size += 2;
252252
return dest;
253253
} else if (cp <= 0xffff) {
254254
// Three byte UTF-8 characters [cp >= 0x800 && cp <= 0xffff]
@@ -267,7 +267,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
267267
dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)];
268268
cp >>>= 2;
269269
dest[2] = UPPER_HEX_DIGITS[cp];
270-
data.addSize(3);
270+
escaped.size += 3;
271271
return dest;
272272
} else if (cp <= 0x10ffff) {
273273
char[] dest = new char[12];
@@ -291,7 +291,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
291291
dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)];
292292
cp >>>= 2;
293293
dest[2] = UPPER_HEX_DIGITS[cp & 0x7];
294-
data.addSize(4);
294+
escaped.size += 4;
295295
return dest;
296296
} else {
297297
// If this ever happens it is due to bug in UnicodeEscaper, not bad input.
@@ -386,4 +386,14 @@ private static char[] growBuffer(char[] dest, int index, int size) {
386386
}
387387
return copy;
388388
}
389+
390+
public static class Escaped {
391+
public String data;
392+
public int size;
393+
394+
public Escaped(String data, int size) {
395+
this.data = data;
396+
this.size = size;
397+
}
398+
}
389399
}

0 commit comments

Comments
 (0)