Skip to content

Commit a511d59

Browse files
mhliddPerfectSlayer
mhlidd
and
PerfectSlayer
authored
Enabling baggage cache to support limits and non-ascii characters (#8713)
* enabling baggage cache to support limits * updating limits to consider non-ascii characters * adding edge case * PR comments * feat(core): Simplify baggage config usage * feat(core): Add non ASCII character baggage extraction --------- Co-authored-by: Bruce Bujon <[email protected]>
1 parent 6c8da4d commit a511d59

File tree

3 files changed

+141
-39
lines changed

3 files changed

+141
-39
lines changed

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

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -24,42 +24,41 @@ public class BaggagePropagator implements Propagator {
2424
private static final Logger LOG = LoggerFactory.getLogger(BaggagePropagator.class);
2525
private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create();
2626
static final String BAGGAGE_KEY = "baggage";
27-
private final Config config;
2827
private final boolean injectBaggage;
2928
private final boolean extractBaggage;
29+
private final int maxItems;
30+
private final int maxBytes;
3031

3132
public BaggagePropagator(Config config) {
32-
this.injectBaggage = config.isBaggageInject();
33-
this.extractBaggage = config.isBaggageExtract();
34-
this.config = config;
33+
this(
34+
config.isBaggageInject(),
35+
config.isBaggageInject(),
36+
config.getTraceBaggageMaxItems(),
37+
config.getTraceBaggageMaxBytes());
3538
}
3639

3740
// use primarily for testing purposes
38-
public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) {
41+
BaggagePropagator(boolean injectBaggage, boolean extractBaggage, int maxItems, int maxBytes) {
3942
this.injectBaggage = injectBaggage;
4043
this.extractBaggage = extractBaggage;
41-
this.config = Config.get();
44+
this.maxItems = maxItems;
45+
this.maxBytes = maxBytes;
4246
}
4347

4448
@Override
4549
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
46-
int maxItems = this.config.getTraceBaggageMaxItems();
47-
int maxBytes = this.config.getTraceBaggageMaxBytes();
48-
//noinspection ConstantValue
50+
Baggage baggage;
4951
if (!this.injectBaggage
50-
|| maxItems == 0
51-
|| maxBytes == 0
52+
|| this.maxItems == 0
53+
|| this.maxBytes == 0
5254
|| context == null
5355
|| carrier == null
54-
|| setter == null) {
55-
return;
56-
}
57-
58-
Baggage baggage = Baggage.fromContext(context);
59-
if (baggage == null) {
56+
|| setter == null
57+
|| (baggage = Baggage.fromContext(context)) == null) {
6058
return;
6159
}
6260

61+
// Inject cached header if any as optimized path
6362
String headerValue = baggage.getW3cHeader();
6463
if (headerValue != null) {
6564
setter.set(carrier, BAGGAGE_KEY, headerValue);
@@ -86,25 +85,25 @@ public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
8685

8786
processedItems++;
8887
// reached the max number of baggage items allowed
89-
if (processedItems == maxItems) {
88+
if (processedItems == this.maxItems) {
9089
break;
9190
}
9291
// Drop newest k/v pair if adding it leads to exceeding the limit
93-
if (currentBytes + escapedKey.size + escapedVal.size + extraBytes > maxBytes) {
92+
if (currentBytes + escapedKey.size + escapedVal.size + extraBytes > this.maxBytes) {
9493
baggageText.setLength(currentBytes);
9594
break;
9695
}
9796
currentBytes += escapedKey.size + escapedVal.size + extraBytes;
9897
}
9998

10099
headerValue = baggageText.toString();
100+
// Save header as cache to re-inject it later if baggage did not change
101101
baggage.setW3cHeader(headerValue);
102102
setter.set(carrier, BAGGAGE_KEY, headerValue);
103103
}
104104

105105
@Override
106106
public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
107-
//noinspection ConstantValue
108107
if (!this.extractBaggage || context == null || carrier == null || visitor == null) {
109108
return context;
110109
}
@@ -113,12 +112,11 @@ public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
113112
return baggageExtractor.extracted == null ? context : context.with(baggageExtractor.extracted);
114113
}
115114

116-
private static class BaggageExtractor implements BiConsumer<String, String> {
115+
private class BaggageExtractor implements BiConsumer<String, String> {
117116
private static final char KEY_VALUE_SEPARATOR = '=';
118117
private static final char PAIR_SEPARATOR = ',';
119118
private Baggage extracted;
120-
121-
private BaggageExtractor() {}
119+
private String w3cHeader;
122120

123121
/** URL decode value */
124122
private String decode(final String value) {
@@ -134,6 +132,7 @@ private String decode(final String value) {
134132
private Map<String, String> parseBaggageHeaders(String input) {
135133
Map<String, String> baggage = new HashMap<>();
136134
int start = 0;
135+
boolean truncatedCache = false;
137136
int pairSeparatorInd = input.indexOf(PAIR_SEPARATOR);
138137
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
139138
int kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR);
@@ -152,11 +151,29 @@ private Map<String, String> parseBaggageHeaders(String input) {
152151
}
153152
baggage.put(key, value);
154153

154+
// need to percent-encode non-ascii headers we pass down
155+
if (UTF_ESCAPER.keyNeedsEncoding(key) || UTF_ESCAPER.valNeedsEncoding(value)) {
156+
truncatedCache = true;
157+
this.w3cHeader = null;
158+
} else if (!truncatedCache && (end > maxBytes || baggage.size() > maxItems)) {
159+
if (start == 0) { // if we go out of range after first k/v pair, there is no cache
160+
this.w3cHeader = null;
161+
} else {
162+
this.w3cHeader = input.substring(0, start - 1); // -1 to ignore the k/v separator
163+
}
164+
truncatedCache = true;
165+
}
166+
155167
kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR, pairSeparatorInd + 1);
156168
pairSeparatorInd = input.indexOf(PAIR_SEPARATOR, pairSeparatorInd + 1);
157169
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
158170
start = end + 1;
159171
}
172+
173+
if (!truncatedCache) {
174+
this.w3cHeader = input;
175+
}
176+
160177
return baggage;
161178
}
162179

@@ -166,7 +183,7 @@ public void accept(String key, String value) {
166183
if (BAGGAGE_KEY.equalsIgnoreCase(key)) {
167184
Map<String, String> baggage = parseBaggageHeaders(value);
168185
if (!baggage.isEmpty()) {
169-
this.extracted = Baggage.create(baggage, value);
186+
this.extracted = Baggage.create(baggage, this.w3cHeader);
170187
}
171188
}
172189
}

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

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,12 +115,38 @@ public Escaped escapeValue(String s) {
115115
return escape(s, unsafeValOctets);
116116
}
117117

118+
private boolean needsEncoding(char c, boolean[] unsafeOctets) {
119+
if (c > '~' || c <= ' ' || c < unsafeOctets.length && unsafeOctets[c]) {
120+
return true;
121+
}
122+
return false;
123+
}
124+
125+
private boolean needsEncoding(String key, boolean[] unsafeOctets) {
126+
int slen = key.length();
127+
for (int index = 0; index < slen; index++) {
128+
char c = key.charAt(index);
129+
if (needsEncoding(c, unsafeOctets)) {
130+
return true;
131+
}
132+
}
133+
return false;
134+
}
135+
136+
public boolean keyNeedsEncoding(String key) {
137+
return needsEncoding(key, unsafeKeyOctets);
138+
}
139+
140+
public boolean valNeedsEncoding(String val) {
141+
return needsEncoding(val, unsafeValOctets);
142+
}
143+
118144
/** Escape the provided String, using percent-style URL Encoding. */
119145
public Escaped escape(String s, boolean[] unsafeOctets) {
120146
int slen = s.length();
121147
for (int index = 0; index < slen; index++) {
122148
char c = s.charAt(index);
123-
if (c > '~' || c <= ' ' || c <= unsafeOctets.length && unsafeOctets[c]) {
149+
if (needsEncoding(c, unsafeOctets)) {
124150
return escapeSlow(s, index, unsafeOctets);
125151
}
126152
}

‎dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy

Lines changed: 73 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import datadog.trace.test.util.DDSpecification
99

1010
import java.util.function.BiConsumer
1111

12+
import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_BAGGAGE_MAX_BYTES
13+
import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_BAGGAGE_MAX_ITEMS
1214
import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY
1315

1416
class BaggagePropagatorTest extends DDSpecification {
@@ -33,7 +35,7 @@ class BaggagePropagatorTest extends DDSpecification {
3335
}
3436

3537
def setup() {
36-
this.propagator = new BaggagePropagator(true, true)
38+
this.propagator = new BaggagePropagator(true, true, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS, DEFAULT_TRACE_BAGGAGE_MAX_BYTES)
3739
this.setter = new MapCarrierAccessor()
3840
this.carrier = [:]
3941
this.context = Context.root()
@@ -61,10 +63,9 @@ class BaggagePropagatorTest extends DDSpecification {
6163
["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5"
6264
}
6365

64-
def "test baggage item limit"() {
66+
def "test baggage inject item limit"() {
6567
setup:
66-
injectSysConfig("trace.baggage.max.items", '2')
67-
propagator = new BaggagePropagator(true, true) //creating a new instance after injecting config
68+
propagator = new BaggagePropagator(true, true, 2, DEFAULT_TRACE_BAGGAGE_MAX_BYTES) //creating a new instance after injecting config
6869
context = Baggage.create(baggage).storeInto(context)
6970

7071
when:
@@ -79,10 +80,9 @@ class BaggagePropagatorTest extends DDSpecification {
7980
[key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2"
8081
}
8182

82-
def "test baggage bytes limit"() {
83+
def "test baggage inject bytes limit"() {
8384
setup:
84-
injectSysConfig("trace.baggage.max.bytes", '20')
85-
propagator = new BaggagePropagator(true, true) //creating a new instance after injecting config
85+
propagator = new BaggagePropagator(true, true, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS, 20) //creating a new instance after injecting config
8686
context = Baggage.create(baggage).storeInto(context)
8787

8888
when:
@@ -116,6 +116,30 @@ class BaggagePropagatorTest extends DDSpecification {
116116
"%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\']
117117
}
118118

119+
def "test extracting non ASCII headers"() {
120+
setup:
121+
def headers = [
122+
(BAGGAGE_KEY) : "key1=vallée,clé2=value",
123+
]
124+
125+
when:
126+
context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap())
127+
def baggage = Baggage.fromContext(context)
128+
129+
then: 'non ASCII values data are still accessible as part of the API'
130+
baggage != null
131+
baggage.asMap().get('key1') == 'vallée'
132+
baggage.asMap().get('clé2') == 'value'
133+
baggage.w3cHeader == null
134+
135+
136+
when:
137+
this.propagator.inject(Context.root().with(baggage), carrier, setter)
138+
139+
then: 'baggage are URL encoded if not valid, even if not modified'
140+
assert carrier[BAGGAGE_KEY] == 'key1=vall%C3%A9e,cl%C3%A92=value'
141+
}
142+
119143
def "extract invalid baggage headers"() {
120144
setup:
121145
def headers = [
@@ -139,8 +163,28 @@ class BaggagePropagatorTest extends DDSpecification {
139163
"=" | _
140164
}
141165

142-
def "testing baggage cache"(){
166+
def "test baggage cache"(){
167+
setup:
168+
def headers = [
169+
(BAGGAGE_KEY) : baggageHeader,
170+
]
171+
172+
when:
173+
context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap())
174+
175+
then:
176+
Baggage baggageContext = Baggage.fromContext(context)
177+
baggageContext.w3cHeader == cachedString
178+
179+
where:
180+
baggageHeader | cachedString
181+
"key1=val1,key2=val2,foo=bar" | "key1=val1,key2=val2,foo=bar"
182+
'";\\()/:<=>?@[]{}=";\\' | null
183+
}
184+
185+
def "test baggage cache items limit"(){
143186
setup:
187+
propagator = new BaggagePropagator(true, true, 2, DEFAULT_TRACE_BAGGAGE_MAX_BYTES) //creating a new instance after injecting config
144188
def headers = [
145189
(BAGGAGE_KEY) : baggageHeader,
146190
]
@@ -150,17 +194,32 @@ class BaggagePropagatorTest extends DDSpecification {
150194

151195
then:
152196
Baggage baggageContext = Baggage.fromContext(context)
153-
baggageContext.asMap() == baggageMap
197+
baggageContext.getW3cHeader() as String == cachedString
198+
199+
where:
200+
baggageHeader | cachedString
201+
"key1=val1,key2=val2" | "key1=val1,key2=val2"
202+
"key1=val1,key2=val2,key3=val3" | "key1=val1,key2=val2"
203+
"key1=val1,key2=val2,key3=val3,key4=val4" | "key1=val1,key2=val2"
204+
}
205+
206+
def "test baggage cache bytes limit"(){
207+
setup:
208+
propagator = new BaggagePropagator(true, true, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS, 20) //creating a new instance after injecting config
209+
def headers = [
210+
(BAGGAGE_KEY) : baggageHeader,
211+
]
154212

155213
when:
156-
this.propagator.inject(context, carrier, setter)
214+
context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap())
157215

158216
then:
159-
assert carrier[BAGGAGE_KEY] == baggageHeader
217+
Baggage baggageContext = Baggage.fromContext(context)
218+
baggageContext.getW3cHeader() as String == cachedString
160219

161220
where:
162-
baggageHeader | baggageMap
163-
"key1=val1,key2=val2,foo=bar" | ["key1": "val1", "key2": "val2", "foo": "bar"]
164-
"%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\']
221+
baggageHeader | cachedString
222+
"key1=val1,key2=val2" | "key1=val1,key2=val2"
223+
"key1=val1,key2=val2,key3=val3" | "key1=val1,key2=val2"
165224
}
166225
}

0 commit comments

Comments
 (0)