From 20ef57a61518f19d78054b789460b7c60b0c9d7e Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 7 Mar 2025 12:01:05 +0100 Subject: [PATCH 1/2] feat(internal-api): Improve Baggage API --- .../trace/core/baggage/BaggagePropagator.java | 98 +++++++------- .../datadog/trace/core/util/EscapedData.java | 36 ------ .../trace/core/util/PercentEscaper.java | 42 +++--- .../core/baggage/BaggagePropagatorTest.groovy | 25 ++-- .../instrumentation/api/Baggage.java | 122 ++++++++++++++++++ .../instrumentation/api/BaggageContext.java | 74 ----------- 6 files changed, 207 insertions(+), 190 deletions(-) delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java create mode 100644 internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Baggage.java delete mode 100644 internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java diff --git a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java index 67f8702291e..e203a1722db 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java @@ -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; @@ -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); private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create(); static final String BAGGAGE_KEY = "baggage"; private final Config config; @@ -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(); } @Override public void inject(Context context, C carrier, CarrierSetter 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 @@ -54,52 +55,52 @@ public void inject(Context context, C carrier, CarrierSetter 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 entry : baggageContext.asMap().entrySet()) { + for (final Map.Entry 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 @@ -108,19 +109,17 @@ public Context extract(Context context, C carrier, CarrierVisitor 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 { - private BaggageContext extractedContext; + private static class BaggageExtractor implements BiConsumer { + private static final char KEY_VALUE_SEPARATOR = '='; + private static final char PAIR_SEPARATOR = ','; + private Baggage extracted; - BaggageContextExtractor() {} + private BaggageExtractor() {} /** URL decode value */ private String decode(final String value) { @@ -128,37 +127,34 @@ private String decode(final String 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 parseBaggageHeaders(String input) { Map 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; } @@ -168,10 +164,10 @@ private Map 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)) { Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { - extractedContext = BaggageContext.create(baggage, value); + this.extracted = Baggage.create(baggage, value); } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java b/dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java deleted file mode 100644 index 0ff428d2f60..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java +++ /dev/null @@ -1,36 +0,0 @@ -package datadog.trace.core.util; - -public class EscapedData { - private String data; - private int size; - - public EscapedData(String data, int size) { - this.data = data; - this.size = size; - } - - public EscapedData() { - this.data = ""; - this.size = 0; - } - - public String getData() { - return data; - } - - public int getSize() { - return size; - } - - public void setData(String data) { - this.data = data; - } - - public void incrementSize() { - size++; - } - - public void addSize(int delta) { - size += delta; - } -} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java b/dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java index 8bd6a7b21d5..c32036713ad 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java @@ -107,17 +107,16 @@ 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); @@ -125,7 +124,7 @@ public EscapedData escape(String s, boolean[] unsafeOctets) { return escapeSlow(s, index, unsafeOctets); } } - return new EscapedData(s, slen); + return new Escaped(s, slen); } /* @@ -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); @@ -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; @@ -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) { @@ -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]) { @@ -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] @@ -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] @@ -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]; @@ -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. @@ -386,4 +386,14 @@ private static char[] growBuffer(char[] dest, int index, int size) { } return copy; } + + public static class Escaped { + public String data; + public int size; + + public Escaped(String data, int size) { + this.data = data; + this.size = size; + } + } } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy index 56a16cc376f..4910898c047 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy @@ -3,21 +3,20 @@ package datadog.trace.core.baggage import datadog.context.Context import datadog.context.propagation.CarrierSetter import datadog.context.propagation.CarrierVisitor -import datadog.trace.bootstrap.instrumentation.api.BaggageContext +import datadog.trace.bootstrap.instrumentation.api.Baggage import datadog.trace.bootstrap.instrumentation.api.ContextVisitors -import datadog.trace.core.test.DDCoreSpecification +import datadog.trace.test.util.DDSpecification import java.util.function.BiConsumer import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY -class BaggagePropagatorTest extends DDCoreSpecification { +class BaggagePropagatorTest extends DDSpecification { BaggagePropagator propagator CarrierSetter setter Map carrier Context context - static class MapCarrierAccessor implements CarrierSetter>, CarrierVisitor> { @Override @@ -35,14 +34,14 @@ class BaggagePropagatorTest extends DDCoreSpecification { def setup() { this.propagator = new BaggagePropagator(true, true) - setter = new MapCarrierAccessor() - carrier = [:] - context = Context.root() + this.setter = new MapCarrierAccessor() + this.carrier = [:] + this.context = Context.root() } def 'test baggage propagator context injection'() { setup: - context = BaggageContext.create(baggageMap).storeInto(context) + this.context = Baggage.create(baggageMap).storeInto(this.context) when: this.propagator.inject(context, carrier, setter) @@ -66,7 +65,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { setup: injectSysConfig("trace.baggage.max.items", '2') propagator = new BaggagePropagator(true, true) //creating a new instance after injecting config - context = BaggageContext.create(baggage).storeInto(context) + context = Baggage.create(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -84,7 +83,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { setup: injectSysConfig("trace.baggage.max.bytes", '20') propagator = new BaggagePropagator(true, true) //creating a new instance after injecting config - context = BaggageContext.create(baggage).storeInto(context) + context = Baggage.create(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -109,7 +108,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - BaggageContext.fromContext(context).asMap() == baggageMap + Baggage.fromContext(context).asMap() == baggageMap where: baggageHeader | baggageMap @@ -127,7 +126,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - BaggageContext.fromContext(context) == null + Baggage.fromContext(context) == null where: baggageHeader | _ @@ -150,7 +149,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - BaggageContext baggageContext = BaggageContext.fromContext(context) + Baggage baggageContext = Baggage.fromContext(context) baggageContext.asMap() == baggageMap when: diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Baggage.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Baggage.java new file mode 100644 index 00000000000..e6ce2f73751 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Baggage.java @@ -0,0 +1,122 @@ +package datadog.trace.bootstrap.instrumentation.api; + +import static datadog.context.ContextKey.named; +import static java.util.Collections.unmodifiableMap; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ImplicitContextKeyed; +import java.util.HashMap; +import java.util.Map; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** Baggage are key/value store which propagate alongside {@link Context}. */ +public class Baggage implements ImplicitContextKeyed { + private static final ContextKey CONTEXT_KEY = named("baggage-key"); + private final Map items; + /** + * The W3C Baggage header representation of the + * baggage instance, {@code null} if not in sync with the current baggage items. + */ + private String w3cHeader; + + /** + * Gets baggage from context. + * + * @param context The context to get baggage from. + * @return The baggage from the given context if any, {@code null} if none. + */ + public static @Nullable Baggage fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + /** + * Create empty baggage. + * + * @return Empty baggage. + */ + public static Baggage empty() { + return create(new HashMap<>(), ""); + } + + /** + * Create baggage from items. + * + * @param items The original baggage items. + * @return Baggage with the given items. + */ + public static Baggage create(Map items) { + return new Baggage(items, null); + } + + /** + * Create baggage from items. + * + * @param items The original baggage items. + * @param w3cHeader The W3C Baggage header representation. + * @return Baggage with the given items and its W3C header cached. + */ + public static Baggage create(Map items, String w3cHeader) { + return new Baggage(items, w3cHeader); + } + + private Baggage(Map items, String w3cHeader) { + this.items = items; + this.w3cHeader = w3cHeader; + } + + /** + * Adds a baggage item. + * + * @param key The item key. + * @param value The item value. + */ + public void addItem(String key, String value) { + this.items.put(key, value); + this.w3cHeader = null; + } + + /** + * Removes a baggage item. + * + * @param key The item key to remove. + */ + public void removeItem(String key) { + if (this.items.remove(key) != null) { + this.w3cHeader = null; + } + } + + /** + * Gets the W3C Baggage header representation. + * + * @return The header value, {@code null} if not in sync with the current baggage items. + */ + public @Nullable String getW3cHeader() { + return this.w3cHeader; + } + + /** + * Updates the W3C Baggage header representation. + * + * @param w3cHeader The header value. + */ + public void setW3cHeader(String w3cHeader) { + this.w3cHeader = w3cHeader; + } + + /** + * Gets a view of the baggage items. + * + * @return The read-only view of the baggage items. + */ + public Map asMap() { + return unmodifiableMap(this.items); + } + + @Override + public Context storeInto(@Nonnull Context context) { + return context.with(CONTEXT_KEY, this); + } +} diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java deleted file mode 100644 index 99048f7a107..00000000000 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java +++ /dev/null @@ -1,74 +0,0 @@ -package datadog.trace.bootstrap.instrumentation.api; - -import datadog.context.Context; -import datadog.context.ContextKey; -import datadog.context.ImplicitContextKeyed; -import java.util.HashMap; -import java.util.Map; - -public class BaggageContext implements ImplicitContextKeyed { - private static final ContextKey CONTEXT_KEY = ContextKey.named("baggage-key"); - - private final Map baggage; - private String baggageString; - private boolean updatedCache; - - public BaggageContext empty() { - return create(new HashMap<>(), ""); - } - - public static BaggageContext create(Map baggage) { - return new BaggageContext(baggage); - } - - private BaggageContext(Map baggage) { - this.baggage = baggage; - this.baggageString = ""; - updatedCache = false; - } - - public static BaggageContext create(Map baggage, String w3cHeader) { - return new BaggageContext(baggage, w3cHeader); - } - - private BaggageContext(Map baggage, String baggageString) { - this.baggage = baggage; - this.baggageString = baggageString; - updatedCache = true; - } - - public void addW3CBaggage(String key, String value) { - baggage.put(key, value); - updatedCache = false; - } - - public void removeW3CBaggage(String key) { - baggage.remove(key); - updatedCache = false; - } - - public void setW3cBaggageHeader(String w3cHeader) { - this.baggageString = w3cHeader; - updatedCache = true; - } - - public String getW3cBaggageHeader() { - if (updatedCache) { - return baggageString; - } - return null; - } - - public Map asMap() { - return new HashMap<>(baggage); - } - - public static BaggageContext fromContext(Context context) { - return context.get(CONTEXT_KEY); - } - - @Override - public Context storeInto(Context context) { - return context.with(CONTEXT_KEY, this); - } -} From 477151989ba4da60609204724d74b9c6e1b2472d Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Fri, 7 Mar 2025 13:34:47 +0100 Subject: [PATCH 2/2] feat(internal-api): Add Baggage unit tests --- internal-api/build.gradle | 2 - .../instrumentation/api/BaggageTest.groovy | 88 +++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) create mode 100644 internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/BaggageTest.groovy diff --git a/internal-api/build.gradle b/internal-api/build.gradle index 4731d37a255..001b0adfd4a 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -189,8 +189,6 @@ excludedClassesCoverage += [ "datadog.trace.api.iast.Taintable", "datadog.trace.api.Stateful", "datadog.trace.api.Stateful.1", - // BaggageContext class tested in BaggagePropagatorTest in dd-trace-core - 'datadog.trace.bootstrap.instrumentation.api.BaggageContext', // a stub "datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration", "datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration.NoOp", diff --git a/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/BaggageTest.groovy b/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/BaggageTest.groovy new file mode 100644 index 00000000000..7611a87b407 --- /dev/null +++ b/internal-api/src/test/groovy/datadog/trace/bootstrap/instrumentation/api/BaggageTest.groovy @@ -0,0 +1,88 @@ +package datadog.trace.bootstrap.instrumentation.api + +import datadog.context.Context +import spock.lang.Specification + +class BaggageTest extends Specification { + def 'test empty baggage'() { + when: + def baggage = Baggage.empty() + + then: + baggage.asMap().isEmpty() + baggage.w3cHeader != null + baggage.w3cHeader.isEmpty() + + when: + baggage.addItem('key', 'value') + + then: + baggage.w3cHeader == null + } + + def 'test baggage creation'() { + setup: + def items = ['key1': 'value1', 'key2': 'value2'] + def header = 'key1=value1,key2=value2' + + when: + def baggage = Baggage.create(items) + + then: + baggage.asMap() == items + baggage.w3cHeader == null + + when: + baggage = Baggage.create(items, header) + + then: + baggage.asMap() == items + baggage.w3cHeader == header + } + + def 'test baggage header'() { + setup: + def items = ['key1': 'value1', 'key2': 'value2'] + def header = 'key1=value1,key2=value2' + def baggage = Baggage.create(items, header) + + when: + baggage.removeItem('missingKey') + + then: 'header is preserved' + baggage.w3cHeader == header + + when: + baggage.removeItem('key2') + + then: 'header is out of sync' + baggage.w3cHeader == null + + when: + baggage.w3cHeader = 'key1=value1' + + then: 'header is forced' + baggage.w3cHeader == 'key1=value1' + + when: + baggage.addItem('key3', 'value3') + + then: 'header is out of sync' + baggage.w3cHeader == null + } + + def 'test context storage'() { + given: + def baggage = Baggage.empty() + def context = Context.root() + + expect: + Baggage.fromContext(context) == null + + when: + context = context.with(baggage) + + then: + Baggage.fromContext(context) == baggage + } +}