Skip to content

Improve Baggage API #8523

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 11, 2025
Merged

Improve Baggage API #8523

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
20ef57a
feat(internal-api): Improve Baggage API
PerfectSlayer Mar 7, 2025
4771519
feat(internal-api): Add Baggage unit tests
PerfectSlayer Mar 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 47 additions & 51 deletions dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
package datadog.trace.core.baggage;

import static java.util.Collections.emptyMap;

import datadog.context.Context;
import datadog.context.propagation.CarrierSetter;
import datadog.context.propagation.CarrierVisitor;
import datadog.context.propagation.Propagator;
import datadog.trace.api.Config;
import datadog.trace.bootstrap.instrumentation.api.BaggageContext;
import datadog.trace.core.util.EscapedData;
import datadog.trace.bootstrap.instrumentation.api.Baggage;
import datadog.trace.core.util.PercentEscaper;
import datadog.trace.core.util.PercentEscaper.Escaped;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.function.BiConsumer;
Expand All @@ -20,7 +21,7 @@

@ParametersAreNonnullByDefault
public class BaggagePropagator implements Propagator {
private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class);
private static final Logger LOG = LoggerFactory.getLogger(BaggagePropagator.class);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

static final vars are considered as constants and are uppercased by convention

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I admit to having mixed views on this - I understand the convention of marking static final's as upper-cased constants, but find using LOG.info(...) rather glaring in the code.

Do we want to codify this in the spotless config and update all the code to use this convention, or just follow it for new files?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

find using LOG.info(...) rather glaring in the code.

Interesting. What's the rationale behind your feeling?

LOG being a constant, it makes sense to me.
During reviews, I usually end up looking for the definition of anything I can't find in the immediate scope that is not a constant (UPPER_CASE) nor a field member (this.field).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Past habit of distinguishing between value constants (denoting a fixed simple value) and singletons (which are similar in that the reference doesn't change, but internally things vary.)

I'm in no way consistent about this though, having picked up the use of INSTANCE while working in this codebase :)

private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create();
static final String BAGGAGE_KEY = "baggage";
private final Config config;
Expand All @@ -37,13 +38,13 @@ public BaggagePropagator(Config config) {
public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) {
this.injectBaggage = injectBaggage;
this.extractBaggage = extractBaggage;
config = Config.get();
this.config = Config.get();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't worry about the this. additions, it's mostly my formatter at play here and there...

}

@Override
public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
int maxItems = config.getTraceBaggageMaxItems();
int maxBytes = config.getTraceBaggageMaxBytes();
int maxItems = this.config.getTraceBaggageMaxItems();
int maxBytes = this.config.getTraceBaggageMaxBytes();
//noinspection ConstantValue
if (!this.injectBaggage
|| maxItems == 0
Expand All @@ -54,52 +55,52 @@ public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
return;
}

BaggageContext baggageContext = BaggageContext.fromContext(context);
if (baggageContext == null) {
log.debug("BaggageContext instance is missing from the following context {}", context);
Baggage baggage = Baggage.fromContext(context);
if (baggage == null) {
LOG.debug("Baggage instance is missing from the following context {}", context);
return;
}

String baggageHeader = baggageContext.getW3cBaggageHeader();
if (baggageHeader != null) {
setter.set(carrier, BAGGAGE_KEY, baggageHeader);
String headerValue = baggage.getW3cHeader();
if (headerValue != null) {
setter.set(carrier, BAGGAGE_KEY, headerValue);
return;
}

int processedBaggage = 0;
int processedItems = 0;
int currentBytes = 0;
StringBuilder baggageText = new StringBuilder();
for (final Map.Entry<String, String> entry : baggageContext.asMap().entrySet()) {
for (final Map.Entry<String, String> entry : baggage.asMap().entrySet()) {
// if there are already baggage items processed, add and allocate bytes for a comma
int extraBytes = 1;
if (processedBaggage != 0) {
if (processedItems != 0) {
baggageText.append(',');
extraBytes++;
}

EscapedData escapedKey = UTF_ESCAPER.escapeKey(entry.getKey());
EscapedData escapedVal = UTF_ESCAPER.escapeValue(entry.getValue());
Escaped escapedKey = UTF_ESCAPER.escapeKey(entry.getKey());
Escaped escapedVal = UTF_ESCAPER.escapeValue(entry.getValue());

baggageText.append(escapedKey.getData());
baggageText.append(escapedKey.data);
baggageText.append('=');
baggageText.append(escapedVal.getData());
baggageText.append(escapedVal.data);

processedBaggage++;
processedItems++;
// reached the max number of baggage items allowed
if (processedBaggage == maxItems) {
if (processedItems == maxItems) {
break;
}
// Drop newest k/v pair if adding it leads to exceeding the limit
if (currentBytes + escapedKey.getSize() + escapedVal.getSize() + extraBytes > maxBytes) {
if (currentBytes + escapedKey.size + escapedVal.size + extraBytes > maxBytes) {
baggageText.setLength(currentBytes);
break;
}
currentBytes += escapedKey.getSize() + escapedVal.getSize() + extraBytes;
currentBytes += escapedKey.size + escapedVal.size + extraBytes;
}

String baggageString = baggageText.toString();
baggageContext.setW3cBaggageHeader(baggageString);
setter.set(carrier, BAGGAGE_KEY, baggageString);
headerValue = baggageText.toString();
baggage.setW3cHeader(headerValue);
setter.set(carrier, BAGGAGE_KEY, headerValue);
}

@Override
Expand All @@ -108,57 +109,52 @@ public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor
if (!this.extractBaggage || context == null || carrier == null || visitor == null) {
return context;
}
BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor();
visitor.forEachKeyValue(carrier, baggageContextExtractor);
BaggageContext extractedContext = baggageContextExtractor.extractedContext;
if (extractedContext == null) {
return context;
}
return extractedContext.storeInto(context);
BaggageExtractor baggageExtractor = new BaggageExtractor();
visitor.forEachKeyValue(carrier, baggageExtractor);
return baggageExtractor.extracted == null ? context : context.with(baggageExtractor.extracted);
}

public static class BaggageContextExtractor implements BiConsumer<String, String> {
private BaggageContext extractedContext;
private static class BaggageExtractor implements BiConsumer<String, String> {
private static final char KEY_VALUE_SEPARATOR = '=';
private static final char PAIR_SEPARATOR = ',';
Comment on lines +118 to +119
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extracted few constants as class constant for clarity, avoiding magic string within method body.

private Baggage extracted;

BaggageContextExtractor() {}
private BaggageExtractor() {}

/** URL decode value */
private String decode(final String value) {
String decoded = value;
try {
decoded = URLDecoder.decode(value, "UTF-8");
} catch (final UnsupportedEncodingException | IllegalArgumentException e) {
log.debug("Failed to decode {}", value);
LOG.debug("Failed to decode {}", value);
}
return decoded;
}

private Map<String, String> parseBaggageHeaders(String input) {
Map<String, String> baggage = new HashMap<>();
char keyValueSeparator = '=';
char pairSeparator = ',';
int start = 0;

int pairSeparatorInd = input.indexOf(pairSeparator);
int pairSeparatorInd = input.indexOf(PAIR_SEPARATOR);
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
int kvSeparatorInd = input.indexOf(keyValueSeparator);
int kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR);
while (kvSeparatorInd != -1) {
int end = pairSeparatorInd;
if (kvSeparatorInd > end) {
log.debug(
LOG.debug(
"Dropping baggage headers due to key with no value {}", input.substring(start, end));
return Collections.emptyMap();
return emptyMap();
}
String key = decode(input.substring(start, kvSeparatorInd).trim());
String value = decode(input.substring(kvSeparatorInd + 1, end).trim());
if (key.isEmpty() || value.isEmpty()) {
log.debug("Dropping baggage headers due to empty k/v {}:{}", key, value);
return Collections.emptyMap();
LOG.debug("Dropping baggage headers due to empty k/v {}:{}", key, value);
return emptyMap();
}
baggage.put(key, value);

kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1);
pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1);
kvSeparatorInd = input.indexOf(KEY_VALUE_SEPARATOR, pairSeparatorInd + 1);
pairSeparatorInd = input.indexOf(PAIR_SEPARATOR, pairSeparatorInd + 1);
pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd;
start = end + 1;
}
Expand All @@ -168,10 +164,10 @@ private Map<String, String> parseBaggageHeaders(String input) {
@Override
public void accept(String key, String value) {
// Only process tags that are relevant to baggage
if (key != null && key.equalsIgnoreCase(BAGGAGE_KEY)) {
if (BAGGAGE_KEY.equalsIgnoreCase(key)) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would not checking for null here potentially lead to a NullPointerException?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting question. This changes was made to simplify the checke and denotes a common rule for Java object equality check: if you have value V to test against a constant C, you would better write if (C.equals(V)) { rather than if (V.equals(C)) {.

While equals implementation contract enforces to have A.equals(B) and B.equals(A) to give the same result, they don't behave the same with null value. As equals is an instance method, it will only throw an NPE if the left operator (the instance object on which method is called) is null not the right operator (the call parameter).

So you end up always writing C.equals(V) because the constant is usually not null and save you from future NPE.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh interesting, I never knew that but it makes sense!

Map<String, String> baggage = parseBaggageHeaders(value);
if (!baggage.isEmpty()) {
extractedContext = BaggageContext.create(baggage, value);
this.extracted = Baggage.create(baggage, value);
}
}
}
Expand Down
36 changes: 0 additions & 36 deletions dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java
View file
Open in desktop

This file was deleted.

42 changes: 26 additions & 16 deletions dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,25 +107,24 @@ private static boolean[] createUnsafeOctets(String safeChars) {
return octets;
}

public EscapedData escapeKey(String s) {
public Escaped escapeKey(String s) {
return escape(s, unsafeKeyOctets);
}

public EscapedData escapeValue(String s) {
public Escaped escapeValue(String s) {
return escape(s, unsafeValOctets);
}

/** Escape the provided String, using percent-style URL Encoding. */
public EscapedData escape(String s, boolean[] unsafeOctets) {
int size = 0;
public Escaped escape(String s, boolean[] unsafeOctets) {
int slen = s.length();
for (int index = 0; index < slen; index++) {
char c = s.charAt(index);
if (c > '~' || c <= ' ' || c <= unsafeOctets.length && unsafeOctets[c]) {
return escapeSlow(s, index, unsafeOctets);
}
}
return new EscapedData(s, slen);
return new Escaped(s, slen);
}

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

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

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

data.setData(new String(dest, 0, destIndex));
return data;
result.data = new String(dest, 0, destIndex);
return result;
}

private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean[] unsafeOctets) {
Expand All @@ -221,7 +221,7 @@ private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean
/** Escapes the given Unicode code point in UTF-8. */
@CheckForNull
@SuppressWarnings("UngroupedOverloads")
private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
private static char[] escape(int cp, Escaped escaped, boolean[] unsafeOctets) {
// We should never get negative values here but if we do it will throw an
// IndexOutOfBoundsException, so at least it will get spotted.
if (cp < unsafeOctets.length && !unsafeOctets[cp]) {
Expand All @@ -233,7 +233,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
dest[0] = '%';
dest[2] = UPPER_HEX_DIGITS[cp & 0xF];
dest[1] = UPPER_HEX_DIGITS[cp >>> 4];
data.incrementSize();
escaped.size++;
return dest;
} else if (cp <= 0x7ff) {
// Two byte UTF-8 characters [cp >= 0x80 && cp <= 0x7ff]
Expand All @@ -248,7 +248,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
dest[2] = UPPER_HEX_DIGITS[cp & 0xF];
cp >>>= 4;
dest[1] = UPPER_HEX_DIGITS[0xC | cp];
data.addSize(2);
escaped.size += 2;
return dest;
} else if (cp <= 0xffff) {
// Three byte UTF-8 characters [cp >= 0x800 && cp <= 0xffff]
Expand All @@ -267,7 +267,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)];
cp >>>= 2;
dest[2] = UPPER_HEX_DIGITS[cp];
data.addSize(3);
escaped.size += 3;
return dest;
} else if (cp <= 0x10ffff) {
char[] dest = new char[12];
Expand All @@ -291,7 +291,7 @@ private static char[] escape(int cp, EscapedData data, boolean[] unsafeOctets) {
dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)];
cp >>>= 2;
dest[2] = UPPER_HEX_DIGITS[cp & 0x7];
data.addSize(4);
escaped.size += 4;
return dest;
} else {
// If this ever happens it is due to bug in UnicodeEscaper, not bad input.
Expand Down Expand Up @@ -386,4 +386,14 @@ private static char[] growBuffer(char[] dest, int index, int size) {
}
return copy;
}

public static class Escaped {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I went with a very basic structure like object.
As there is no logic here, all fields are opened.
Only left an helper constructor to create the object with the right value in a single line.

This whole object thing is a not great but a general mitigation of Java lang as not being able to return multiple values.
I also renamed it to avoid having call like data.data = xxx.

public String data;
public int size;

public Escaped(String data, int size) {
this.data = data;
this.size = size;
}
}
}
Loading