From 0192caff850b59381f9ce31c61bea6bc4cb73cc5 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Feb 2025 16:34:29 -0500 Subject: [PATCH 01/31] initial implementation and unit tests for inject/extract --- .../datadog/trace/api/ConfigDefaults.java | 5 +- .../trace/api/TracePropagationStyle.java | 3 + .../trace/api/config/TracerConfig.java | 2 + .../core/propagation/BaggageHttpCodec.java | 162 ++++++++++++++++++ .../trace/core/propagation/HttpCodec.java | 5 + .../BaggageHttpExtractorTest.groovy | 77 +++++++++ .../BaggageHttpInjectorTest.groovy | 156 +++++++++++++++++ .../main/java/datadog/trace/api/Config.java | 16 ++ 8 files changed, 425 insertions(+), 1 deletion(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java create mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy create mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index af1be8804ee..26678aca032 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java @@ -1,5 +1,6 @@ package datadog.trace.api; +import static datadog.trace.api.TracePropagationStyle.BAGGAGE; import static datadog.trace.api.TracePropagationStyle.DATADOG; import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT; import static java.util.Arrays.asList; @@ -78,9 +79,11 @@ public final class ConfigDefaults { static final int DEFAULT_PARTIAL_FLUSH_MIN_SPANS = 1000; static final boolean DEFAULT_PROPAGATION_EXTRACT_LOG_HEADER_NAMES_ENABLED = false; static final Set DEFAULT_TRACE_PROPAGATION_STYLE = - new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT)); + new LinkedHashSet<>(asList(DATADOG, TRACECONTEXT, BAGGAGE)); static final Set DEFAULT_PROPAGATION_STYLE = new LinkedHashSet<>(asList(PropagationStyle.DATADOG)); + static final int DEFAULT_TRACE_BAGGAGE_MAX_ITEMS = 64; + static final int DEFAULT_TRACE_BAGGAGE_MAX_BYTES = 8192; static final boolean DEFAULT_JMX_FETCH_ENABLED = true; static final boolean DEFAULT_TRACE_AGENT_V05_ENABLED = false; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java b/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java index 192978cc388..2e9e1cf3b79 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java @@ -21,6 +21,9 @@ public enum TracePropagationStyle { // W3C trace context propagation style // https://www.w3.org/TR/trace-context-1/ TRACECONTEXT, + // OTEL baggage support + // https://www.w3.org/TR/baggage/ + BAGGAGE, // None does not extract or inject NONE; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java index b3426d8d989..e4cebef3308 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java @@ -91,6 +91,8 @@ public final class TracerConfig { public static final String TRACE_PROPAGATION_STYLE_EXTRACT = "trace.propagation.style.extract"; public static final String TRACE_PROPAGATION_STYLE_INJECT = "trace.propagation.style.inject"; public static final String TRACE_PROPAGATION_EXTRACT_FIRST = "trace.propagation.extract.first"; + public static final String TRACE_BAGGAGE_MAX_ITEMS = "trace.baggage.max.items"; + public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes"; public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled"; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java new file mode 100644 index 00000000000..6856a94ac3a --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java @@ -0,0 +1,162 @@ +package datadog.trace.core.propagation; + +import static datadog.trace.api.TracePropagationStyle.BAGGAGE; + +import datadog.trace.api.Config; +import datadog.trace.api.TraceConfig; +import datadog.trace.api.TracePropagationStyle; +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +import datadog.trace.bootstrap.instrumentation.api.TagContext; +import datadog.trace.core.DDSpanContext; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Supplier; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** A codec designed for HTTP transport via headers using Datadog headers */ +class BaggageHttpCodec { + private static final Logger log = LoggerFactory.getLogger(BaggageHttpCodec.class); + + static final String BAGGAGE_KEY = "baggage"; + private static final int MAX_CHARACTER_SIZE = 4; + + private BaggageHttpCodec() { + // This class should not be created. This also makes code coverage checks happy. + } + + public static HttpCodec.Injector newInjector(Map invertedBaggageMapping) { + return new Injector(invertedBaggageMapping); + } + + private static class Injector implements HttpCodec.Injector { + + private final Map invertedBaggageMapping; + + public Injector(Map invertedBaggageMapping) { + assert invertedBaggageMapping != null; + this.invertedBaggageMapping = invertedBaggageMapping; + } + + @Override + public void inject( + final DDSpanContext context, final C carrier, final AgentPropagation.Setter setter) { + Config config = Config.get(); + + StringBuilder baggageText = new StringBuilder(); + int processedBaggage = 0; + int currentBytes = 0; + int maxItems = config.getTraceBaggageMaxItems(); + int maxBytes = config.getTraceBaggageMaxBytes(); + int currentCharacters = 0; + int maxSafeCharacters = maxBytes / MAX_CHARACTER_SIZE; + for (final Map.Entry entry : context.baggageItems()) { + if (processedBaggage >= maxItems) { + break; + } + StringBuilder currentText = new StringBuilder(); + if (processedBaggage != 0) { + currentText.append(','); + } + + currentText.append(HttpCodec.encodeBaggage(entry.getKey())); + currentText.append('='); + currentText.append(HttpCodec.encodeBaggage(entry.getValue())); + + // worst case check + if (currentCharacters + currentText.length() <= maxSafeCharacters) { + currentCharacters += currentText.length(); + } else { + if (currentBytes + == 0) { // special case to calculate byte size after surpassing worst-case number of + // characters + currentBytes = baggageText.toString().getBytes(StandardCharsets.UTF_8).length; + } + int byteSize = + currentText + .toString() + .getBytes(StandardCharsets.UTF_8) + .length; // find largest possible byte size for UTF encoded characters and only do + // size checking after we hit this worst case scenario + if (byteSize + currentBytes > maxBytes) { + break; + } + currentBytes += byteSize; + } + baggageText.append(currentText); + processedBaggage++; + } + + setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); + } + } + + public static HttpCodec.Extractor newExtractor( + Config config, Supplier traceConfigSupplier) { + return new TagContextExtractor( + traceConfigSupplier, () -> new BaggageContextInterpreter(config)); + } + + private static class BaggageContextInterpreter extends ContextInterpreter { + + private BaggageContextInterpreter(Config config) { + super(config); + } + + @Override + public TracePropagationStyle style() { + return BAGGAGE; + } + + private Map parseBaggageHeaders(String input) { + Map baggage = new HashMap<>(); + char keyValueSeparator = '='; + char pairSeparator = ','; + int start = 0; + + int pairSeparatorInd = input.indexOf(pairSeparator); + pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; + int kvSeparatorInd = input.indexOf(keyValueSeparator); + while (kvSeparatorInd != -1) { + int end = pairSeparatorInd; + if (kvSeparatorInd > end) { // value is missing + return Collections.emptyMap(); + } + String key = HttpCodec.decode(input.substring(start, kvSeparatorInd).trim()); + String value = HttpCodec.decode(input.substring(kvSeparatorInd + 1, end).trim()); + if (key.isEmpty() || value.isEmpty()) { + return Collections.emptyMap(); + } + baggage.put(key, value); + + kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1); + pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1); + pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; + start = end + 1; + } + return baggage; + } + + @Override + public boolean accept(String key, String value) { + if (null == key || key.isEmpty()) { + return true; + } + if (LOG_EXTRACT_HEADER_NAMES) { + log.debug("Header: {}", key); + } + + if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage + baggage = parseBaggageHeaders(value); + } + return true; + } + + @Override + protected TagContext build() { + return super.build(); + } + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 120d14f7ee5..a3a4398f03e 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -125,6 +125,9 @@ private static Map createInjectors( case TRACECONTEXT: result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping)); break; + case BAGGAGE: + result.put(style, BaggageHttpCodec.newInjector(reverseBaggageMapping)); + break; default: log.debug("No implementation found to inject propagation style: {}", style); break; @@ -159,6 +162,8 @@ public static Extractor createExtractor( case TRACECONTEXT: extractors.add(W3CHttpCodec.newExtractor(config, traceConfigSupplier)); break; + case BAGGAGE: + extractors.add(BaggageHttpCodec.newExtractor(config, traceConfigSupplier)); default: log.debug("No implementation found to extract propagation style: {}", style); break; diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy new file mode 100644 index 00000000000..b2541d2871f --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy @@ -0,0 +1,77 @@ +package datadog.trace.core.propagation + +import datadog.trace.api.Config +import datadog.trace.api.DynamicConfig +import datadog.trace.bootstrap.instrumentation.api.TagContext +import datadog.trace.bootstrap.instrumentation.api.ContextVisitors +import datadog.trace.test.util.DDSpecification +import static datadog.trace.core.propagation.BaggageHttpCodec.BAGGAGE_KEY + +class BaggageHttpExtractorTest extends DDSpecification { + + private DynamicConfig dynamicConfig + private HttpCodec.Extractor _extractor + + private HttpCodec.Extractor getExtractor() { + _extractor ?: (_extractor = createExtractor(Config.get())) + } + + private HttpCodec.Extractor createExtractor(Config config) { + BaggageHttpCodec.newExtractor(config, { dynamicConfig.captureTraceConfig() }) + } + + void setup() { + dynamicConfig = DynamicConfig.create() + .apply() + } + + void cleanup() { + extractor.cleanup() + } + + def "extract valid baggage headers"() { + setup: + def extractor = createExtractor(Config.get()) + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + + when: + final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap()) + + then: + context.baggage == baggageMap + + cleanup: + extractor.cleanup() + + where: + baggageHeader | baggageMap + "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] + "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] + } + + def "extract invalid baggage headers"() { + setup: + def extractor = createExtractor(Config.get()) + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + + when: + final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap()) + + then: + context == null + + cleanup: + extractor.cleanup() + + where: + baggageHeader | baggageMap + "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | [] + "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | [] + "=no-key" | [] + "no-value=" | [] + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy new file mode 100644 index 00000000000..8369d2b3808 --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy @@ -0,0 +1,156 @@ +package datadog.trace.core.propagation + +import datadog.trace.api.DDSpanId +import datadog.trace.api.DDTraceId +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopPathwayContext +import datadog.trace.common.writer.ListWriter +import datadog.trace.core.DDSpanContext +import datadog.trace.core.test.DDCoreSpecification + + +import static datadog.trace.api.sampling.PrioritySampling.* +import static datadog.trace.core.propagation.BaggageHttpCodec.* + + +class BaggageHttpInjectorTest extends DDCoreSpecification { + + HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"]) + + def "test baggage injection and encoding"() { + setup: + def writer = new ListWriter() + def tracer = tracerBuilder().writer(writer).build() + final DDSpanContext mockedContext = + new DDSpanContext( + DDTraceId.from("1"), + DDSpanId.from("2"), + DDSpanId.ZERO, + null, + "fakeService", + "fakeOperation", + "fakeResource", + UNSET, + "fakeOrigin", + baggage, + false, + "fakeType", + 0, + tracer.traceCollectorFactory.create(DDTraceId.ONE), + null, + null, + NoopPathwayContext.INSTANCE, + false, + PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) + + final Map carrier = Mock() + + when: + injector.inject(mockedContext, carrier, MapSetter.INSTANCE) + + then: + 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) + 0 * _ + + cleanup: + tracer.close() + + where: + baggage | baggageHeaders + [key1: "val1"] | "key1=val1" + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [serverNode: "DF 28"] | "serverNode=DF%2028" + [userId: "Amélie"] | "userId=Am%C3%A9lie" + ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" + ["user!d(me)": "false"] | "user!d%28me%29=false" //failing + } + + def "test baggage item limit"() { + setup: + injectSysConfig("trace.baggage.max.items", '2') + def writer = new ListWriter() + def tracer = tracerBuilder().writer(writer).build() + final DDSpanContext mockedContext = + new DDSpanContext( + DDTraceId.from("1"), + DDSpanId.from("2"), + DDSpanId.ZERO, + null, + "fakeService", + "fakeOperation", + "fakeResource", + UNSET, + "fakeOrigin", + baggage, + false, + "fakeType", + 0, + tracer.traceCollectorFactory.create(DDTraceId.ONE), + null, + null, + NoopPathwayContext.INSTANCE, + false, + PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) + + final Map carrier = Mock() + + when: + injector.inject(mockedContext, carrier, MapSetter.INSTANCE) + + then: + 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) + 0 * _ + + cleanup: + tracer.close() + + where: + baggage | baggageHeaders + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + } + + def "test baggage bytes limit"() { + setup: + injectSysConfig("trace.baggage.max.bytes", '20') + def writer = new ListWriter() + def tracer = tracerBuilder().writer(writer).build() + final DDSpanContext mockedContext = + new DDSpanContext( + DDTraceId.from("1"), + DDSpanId.from("2"), + DDSpanId.ZERO, + null, + "fakeService", + "fakeOperation", + "fakeResource", + UNSET, + "fakeOrigin", + baggage, + false, + "fakeType", + 0, + tracer.traceCollectorFactory.create(DDTraceId.ONE), + null, + null, + NoopPathwayContext.INSTANCE, + false, + PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) + + final Map carrier = Mock() + + when: + injector.inject(mockedContext, carrier, MapSetter.INSTANCE) + + then: + 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) + 0 * _ + + cleanup: + tracer.close() + + where: + baggage | baggageHeaders + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index d758f675c32..df3b4d9f371 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -190,6 +190,8 @@ public static String getHostName() { private final Set tracePropagationStylesToExtract; private final Set tracePropagationStylesToInject; private final boolean tracePropagationExtractFirst; + private final int traceBaggageMaxItems; + private final int traceBaggageMaxBytes; private final int clockSyncPeriod; private final boolean logsInjectionEnabled; @@ -991,6 +993,12 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins tracePropagationStylesToExtract = extract.isEmpty() ? DEFAULT_TRACE_PROPAGATION_STYLE : extract; tracePropagationStylesToInject = inject.isEmpty() ? DEFAULT_TRACE_PROPAGATION_STYLE : inject; + + traceBaggageMaxItems = + configProvider.getInteger(TRACE_BAGGAGE_MAX_ITEMS, DEFAULT_TRACE_BAGGAGE_MAX_ITEMS); + traceBaggageMaxBytes = + configProvider.getInteger(TRACE_BAGGAGE_MAX_BYTES, DEFAULT_TRACE_BAGGAGE_MAX_BYTES); + // These setting are here for backwards compatibility until they can be removed in a major // release of the tracer propagationStylesToExtract = @@ -2248,6 +2256,14 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public int getTraceBaggageMaxItems() { + return traceBaggageMaxItems; + } + + public int getTraceBaggageMaxBytes() { + return traceBaggageMaxBytes; + } + public int getClockSyncPeriod() { return clockSyncPeriod; } From fd89905596ab7bd0f5608d081f0ee17ed71f6b19 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Feb 2025 16:55:30 -0500 Subject: [PATCH 02/31] adding break to switch statement[C --- .../src/main/java/datadog/trace/core/propagation/HttpCodec.java | 1 + 1 file changed, 1 insertion(+) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index a3a4398f03e..726bb9f8d14 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -164,6 +164,7 @@ public static Extractor createExtractor( break; case BAGGAGE: extractors.add(BaggageHttpCodec.newExtractor(config, traceConfigSupplier)); + break; default: log.debug("No implementation found to extract propagation style: {}", style); break; From fd19e75ffff026a49863587de8eef42d061c732d Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 4 Feb 2025 15:26:23 -0500 Subject: [PATCH 03/31] better check for max bytes --- .../core/propagation/BaggageHttpCodec.java | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java index 6856a94ac3a..084c20f8459 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java @@ -56,18 +56,16 @@ public void inject( if (processedBaggage >= maxItems) { break; } - StringBuilder currentText = new StringBuilder(); + int additionalCharacters = 1; // accounting for potential comma and colon if (processedBaggage != 0) { - currentText.append(','); + additionalCharacters = 2; // allocating space for comma } - - currentText.append(HttpCodec.encodeBaggage(entry.getKey())); - currentText.append('='); - currentText.append(HttpCodec.encodeBaggage(entry.getValue())); + int currentPairSize = + entry.getKey().length() + entry.getValue().length() + additionalCharacters; // worst case check - if (currentCharacters + currentText.length() <= maxSafeCharacters) { - currentCharacters += currentText.length(); + if (currentCharacters + currentPairSize <= maxSafeCharacters) { + currentCharacters += currentPairSize; } else { if (currentBytes == 0) { // special case to calculate byte size after surpassing worst-case number of @@ -75,17 +73,24 @@ public void inject( currentBytes = baggageText.toString().getBytes(StandardCharsets.UTF_8).length; } int byteSize = - currentText - .toString() - .getBytes(StandardCharsets.UTF_8) - .length; // find largest possible byte size for UTF encoded characters and only do + entry.getKey().getBytes(StandardCharsets.UTF_8).length + + entry.getValue().getBytes(StandardCharsets.UTF_8).length + + additionalCharacters; // find largest possible byte size for UTF encoded + // characters and only do // size checking after we hit this worst case scenario if (byteSize + currentBytes > maxBytes) { break; } currentBytes += byteSize; } - baggageText.append(currentText); + + if (additionalCharacters == 2) { + baggageText.append(','); + } + baggageText + .append(HttpCodec.encodeBaggage(entry.getKey())) + .append('=') + .append(HttpCodec.encodeBaggage(entry.getValue())); processedBaggage++; } From c6d3a0e272c3835c0b6f870a9f1bfc5499601347 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Wed, 5 Feb 2025 15:53:42 -0500 Subject: [PATCH 04/31] updating key encoder and adding test case --- .../core/propagation/BaggageHttpCodec.java | 31 ++++++++++++++++--- .../BaggageHttpInjectorTest.groovy | 4 ++- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java index 084c20f8459..fc30f3b8f3a 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java @@ -9,9 +9,13 @@ import datadog.trace.bootstrap.instrumentation.api.TagContext; import datadog.trace.core.DDSpanContext; import java.nio.charset.StandardCharsets; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; +import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -22,6 +26,10 @@ class BaggageHttpCodec { static final String BAGGAGE_KEY = "baggage"; private static final int MAX_CHARACTER_SIZE = 4; + private static final List UNSAFE_CHARACTERS = + Arrays.asList( + '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', '}'); + private static final Set UNSAFE_CHARACTERS_SET = new HashSet<>(UNSAFE_CHARACTERS); private BaggageHttpCodec() { // This class should not be created. This also makes code coverage checks happy. @@ -40,6 +48,22 @@ public Injector(Map invertedBaggageMapping) { this.invertedBaggageMapping = invertedBaggageMapping; } + private void encodeKey(String key, StringBuilder buffer) { + for (int i = 0; i < key.length(); i++) { + char c = key.charAt(i); + if (UNSAFE_CHARACTERS_SET.contains(c) || c > 126 || c <= 20) { // encode character + byte[] bytes = + Character.toString(c) + .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does + for (byte b : bytes) { + buffer.append(String.format("%%%02X", b & 0xFF)); + } + } else { + buffer.append(c); + } + } + } + @Override public void inject( final DDSpanContext context, final C carrier, final AgentPropagation.Setter setter) { @@ -87,10 +111,9 @@ public void inject( if (additionalCharacters == 2) { baggageText.append(','); } - baggageText - .append(HttpCodec.encodeBaggage(entry.getKey())) - .append('=') - .append(HttpCodec.encodeBaggage(entry.getValue())); + + encodeKey(entry.getKey(), baggageText); + baggageText.append('=').append(HttpCodec.encodeBaggage(entry.getValue())); processedBaggage++; } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy index 8369d2b3808..cbc883fe708 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy @@ -61,7 +61,8 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { [serverNode: "DF 28"] | "serverNode=DF%2028" [userId: "Amélie"] | "userId=Am%C3%A9lie" ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" - ["user!d(me)": "false"] | "user!d%28me%29=false" //failing + ["user!d(me)": "false"] | "user!d%28me%29=false" + ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" } def "test baggage item limit"() { @@ -152,5 +153,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { baggage | baggageHeaders [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + ["abcdefg": "hijklmnopq♥"] | "" } } From 7aa0d4b89c0ef255c86b44d80b9e2a3b77f737d2 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Wed, 5 Feb 2025 16:03:40 -0500 Subject: [PATCH 05/31] cleanup --- .../BaggageHttpInjectorTest.groovy | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy index cbc883fe708..efee85608f8 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy @@ -2,8 +2,10 @@ package datadog.trace.core.propagation import datadog.trace.api.DDSpanId import datadog.trace.api.DDTraceId +import datadog.trace.api.DynamicConfig import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopPathwayContext import datadog.trace.common.writer.ListWriter +import datadog.trace.core.CoreTracer import datadog.trace.core.DDSpanContext import datadog.trace.core.test.DDCoreSpecification @@ -14,12 +16,19 @@ import static datadog.trace.core.propagation.BaggageHttpCodec.* class BaggageHttpInjectorTest extends DDCoreSpecification { - HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"]) + private HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"]) + private ListWriter writer + private CoreTracer tracer + private Map carrier + + void setup() { + writer = new ListWriter() + tracer = tracerBuilder().writer(writer).build() + carrier = Mock() + } def "test baggage injection and encoding"() { setup: - def writer = new ListWriter() - def tracer = tracerBuilder().writer(writer).build() final DDSpanContext mockedContext = new DDSpanContext( DDTraceId.from("1"), @@ -41,9 +50,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { NoopPathwayContext.INSTANCE, false, PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - - final Map carrier = Mock() - when: injector.inject(mockedContext, carrier, MapSetter.INSTANCE) @@ -68,8 +74,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { def "test baggage item limit"() { setup: injectSysConfig("trace.baggage.max.items", '2') - def writer = new ListWriter() - def tracer = tracerBuilder().writer(writer).build() final DDSpanContext mockedContext = new DDSpanContext( DDTraceId.from("1"), @@ -92,8 +96,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { false, PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - final Map carrier = Mock() - when: injector.inject(mockedContext, carrier, MapSetter.INSTANCE) @@ -113,8 +115,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { def "test baggage bytes limit"() { setup: injectSysConfig("trace.baggage.max.bytes", '20') - def writer = new ListWriter() - def tracer = tracerBuilder().writer(writer).build() final DDSpanContext mockedContext = new DDSpanContext( DDTraceId.from("1"), @@ -137,8 +137,6 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { false, PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - final Map carrier = Mock() - when: injector.inject(mockedContext, carrier, MapSetter.INSTANCE) From 2ef1887fc0b2ca1a902d5bc67edfb70439e22393 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 6 Feb 2025 11:04:02 -0500 Subject: [PATCH 06/31] updating encoder and byte checks --- .../core/propagation/BaggageHttpCodec.java | 73 ++++++++++--------- .../BaggageHttpInjectorTest.groovy | 2 +- 2 files changed, 38 insertions(+), 37 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java index fc30f3b8f3a..81c38319e61 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java @@ -13,7 +13,6 @@ import java.util.Collections; import java.util.HashMap; import java.util.HashSet; -import java.util.List; import java.util.Map; import java.util.Set; import java.util.function.Supplier; @@ -26,10 +25,14 @@ class BaggageHttpCodec { static final String BAGGAGE_KEY = "baggage"; private static final int MAX_CHARACTER_SIZE = 4; - private static final List UNSAFE_CHARACTERS = - Arrays.asList( - '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', '}'); - private static final Set UNSAFE_CHARACTERS_SET = new HashSet<>(UNSAFE_CHARACTERS); + + private static final Set UNSAFE_CHARACTERS_KEY = + new HashSet<>( + Arrays.asList( + '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', + '}')); + private static final Set UNSAFE_CHARACTERS_VALUE = + new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); private BaggageHttpCodec() { // This class should not be created. This also makes code coverage checks happy. @@ -48,20 +51,32 @@ public Injector(Map invertedBaggageMapping) { this.invertedBaggageMapping = invertedBaggageMapping; } - private void encodeKey(String key, StringBuilder buffer) { - for (int i = 0; i < key.length(); i++) { - char c = key.charAt(i); - if (UNSAFE_CHARACTERS_SET.contains(c) || c > 126 || c <= 20) { // encode character + private int encodeKey(String key, StringBuilder builder) { + return encode(key, builder, UNSAFE_CHARACTERS_KEY); + } + + private int encodeValue(String key, StringBuilder builder) { + return encode(key, builder, UNSAFE_CHARACTERS_VALUE); + } + + private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { + int size = 0; + for (int i = 0; i < input.length(); i++) { + char c = input.charAt(i); + if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character byte[] bytes = Character.toString(c) .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does for (byte b : bytes) { - buffer.append(String.format("%%%02X", b & 0xFF)); + builder.append(String.format("%%%02X", b & 0xFF)); + size += 1; } } else { - buffer.append(c); + builder.append(c); + size += 1; } } + return size; } @Override @@ -84,36 +99,22 @@ public void inject( if (processedBaggage != 0) { additionalCharacters = 2; // allocating space for comma } - int currentPairSize = - entry.getKey().length() + entry.getValue().length() + additionalCharacters; - - // worst case check - if (currentCharacters + currentPairSize <= maxSafeCharacters) { - currentCharacters += currentPairSize; - } else { - if (currentBytes - == 0) { // special case to calculate byte size after surpassing worst-case number of - // characters - currentBytes = baggageText.toString().getBytes(StandardCharsets.UTF_8).length; - } - int byteSize = - entry.getKey().getBytes(StandardCharsets.UTF_8).length - + entry.getValue().getBytes(StandardCharsets.UTF_8).length - + additionalCharacters; // find largest possible byte size for UTF encoded - // characters and only do - // size checking after we hit this worst case scenario - if (byteSize + currentBytes > maxBytes) { - break; - } - currentBytes += byteSize; - } + int byteSize = 1; // default include size of '=' if (additionalCharacters == 2) { baggageText.append(','); + byteSize += 1; } - encodeKey(entry.getKey(), baggageText); - baggageText.append('=').append(HttpCodec.encodeBaggage(entry.getValue())); + byteSize += encodeKey(entry.getKey(), baggageText); + baggageText.append('='); + byteSize += encodeValue(entry.getValue(), baggageText); + + if (currentBytes + byteSize > maxBytes) { + baggageText.setLength(currentBytes); + break; + } + currentBytes += byteSize; processedBaggage++; } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy index efee85608f8..adbac85489a 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy @@ -149,7 +149,7 @@ class BaggageHttpInjectorTest extends DDCoreSpecification { where: baggage | baggageHeaders - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + // [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" ["abcdefg": "hijklmnopq♥"] | "" } From 1987e9dcd01fd6111aa100fff2a258269d0b082f Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 6 Feb 2025 13:19:55 -0500 Subject: [PATCH 07/31] cleanup --- .../java/datadog/trace/core/propagation/BaggageHttpCodec.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java index 81c38319e61..4ded8742ad8 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java @@ -89,8 +89,6 @@ public void inject( int currentBytes = 0; int maxItems = config.getTraceBaggageMaxItems(); int maxBytes = config.getTraceBaggageMaxBytes(); - int currentCharacters = 0; - int maxSafeCharacters = maxBytes / MAX_CHARACTER_SIZE; for (final Map.Entry entry : context.baggageItems()) { if (processedBaggage >= maxItems) { break; From 11d1350bd12105d4c5739bee2a82fe24a7884e51 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Fri, 7 Feb 2025 11:10:46 -0500 Subject: [PATCH 08/31] introducing baggagecontext --- .../instrumentation/api/BaggageContext.java | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java 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 new file mode 100644 index 00000000000..00e44365628 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java @@ -0,0 +1,33 @@ +package datadog.trace.bootstrap.instrumentation.api; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ImplicitContextKeyed; +import java.util.Map; + +public class BaggageContext implements ImplicitContextKeyed { + private static final ContextKey CONTEXT_KEY = ContextKey.named("baggage-key"); + + private Map baggage; + + public static BaggageContext create(Map baggage) { + return new BaggageContext(baggage); + } + + private BaggageContext(Map baggage) { + this.baggage = baggage; + } + + public static BaggageContext fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + public Map getBaggage() { + return baggage; + } + + @Override + public Context storeInto(Context context) { + return context.with(CONTEXT_KEY, this); + } +} From 64b130349b1fc3983717776b577512b1a026382a Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Sat, 8 Feb 2025 22:06:41 -0500 Subject: [PATCH 09/31] initial migration to contextAPI rebasing --- .../java/datadog/trace/core/CoreTracer.java | 3 + .../trace/core/baggage/BaggageHttpCodec.java | 194 ++++++++++++++++++ .../trace/core/baggage/BaggagePropagator.java | 178 ++++++++++++++++ .../core/propagation/BaggageHttpCodec.java | 189 ----------------- .../trace/core/propagation/HttpCodec.java | 6 - .../core/baggage/BaggagePropagatorTest.groovy | 152 ++++++++++++++ .../BaggageHttpExtractorTest.groovy | 3 +- .../BaggageHttpInjectorTest.groovy | 3 +- .../instrumentation/api/AgentPropagation.java | 2 +- 9 files changed, 531 insertions(+), 199 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java create mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index b9a3d5f07da..bd8329352c6 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -7,6 +7,7 @@ import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.STANDALONE_ASM_CONCERN; +import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.BAGGAGE_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN; import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator; @@ -79,6 +80,7 @@ import datadog.trace.common.writer.WriterFactory; import datadog.trace.common.writer.ddintake.DDIntakeTraceInterceptor; import datadog.trace.context.TraceScope; +import datadog.trace.core.baggage.BaggagePropagator; import datadog.trace.core.datastreams.DataStreamsMonitoring; import datadog.trace.core.datastreams.DefaultDataStreamsMonitoring; import datadog.trace.core.flare.TracerFlarePoller; @@ -735,6 +737,7 @@ private CoreTracer( if (dsm) { Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator()); } + Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator()); this.tagInterceptor = null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java new file mode 100644 index 00000000000..402da477aaa --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java @@ -0,0 +1,194 @@ +// package datadog.trace.core.baggage; +// +// import static datadog.trace.api.TracePropagationStyle.BAGGAGE; +// +// import datadog.trace.api.Config; +// import datadog.trace.api.TraceConfig; +// import datadog.trace.api.TracePropagationStyle; +// import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; +// import datadog.trace.bootstrap.instrumentation.api.TagContext; +// import datadog.trace.core.DDSpanContext; +// import java.nio.charset.StandardCharsets; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.HashMap; +// import java.util.HashSet; +// import java.util.Map; +// import java.util.Set; +// import java.util.function.Supplier; +// +// import datadog.trace.core.propagation.ContextInterpreter; +// import datadog.trace.core.propagation.HttpCodec; +// import datadog.trace.core.propagation.TagContextExtractor; +// import org.slf4j.Logger; +// import org.slf4j.LoggerFactory; +// +/// ** A codec designed for HTTP transport via headers using Datadog headers */ +// class BaggageHttpCodec { +// private static final Logger log = LoggerFactory.getLogger(BaggageHttpCodec.class); +// +// static final String BAGGAGE_KEY = "baggage"; +// private static final int MAX_CHARACTER_SIZE = 4; +// +// private static final Set UNSAFE_CHARACTERS_KEY = +// new HashSet<>( +// Arrays.asList( +// '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', +// '}')); +// private static final Set UNSAFE_CHARACTERS_VALUE = +// new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); +// +// private BaggageHttpCodec() { +// // This class should not be created. This also makes code coverage checks happy. +// } +// +// public static HttpCodec.Injector newInjector(Map invertedBaggageMapping) { +// return new Injector(invertedBaggageMapping); +// } +// +// private static class Injector implements HttpCodec.Injector { +// +// private final Map invertedBaggageMapping; +// +// public Injector(Map invertedBaggageMapping) { +// assert invertedBaggageMapping != null; +// this.invertedBaggageMapping = invertedBaggageMapping; +// } +// +// private int encodeKey(String key, StringBuilder builder) { +// return encode(key, builder, UNSAFE_CHARACTERS_KEY); +// } +// +// private int encodeValue(String key, StringBuilder builder) { +// return encode(key, builder, UNSAFE_CHARACTERS_VALUE); +// } +// +// private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { +// int size = 0; +// for (int i = 0; i < input.length(); i++) { +// char c = input.charAt(i); +// if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character +// byte[] bytes = +// Character.toString(c) +// .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder +// does +// for (byte b : bytes) { +// builder.append(String.format("%%%02X", b & 0xFF)); +// size += 1; +// } +// } else { +// builder.append(c); +// size += 1; +// } +// } +// return size; +// } +// +// @Override +// public void inject( +// final DDSpanContext context, final C carrier, final AgentPropagation.Setter setter) { +// Config config = Config.get(); +// +// StringBuilder baggageText = new StringBuilder(); +// int processedBaggage = 0; +// int currentBytes = 0; +// int maxItems = config.getTraceBaggageMaxItems(); +// int maxBytes = config.getTraceBaggageMaxBytes(); +// for (final Map.Entry entry : context.baggageItems()) { +// if (processedBaggage >= maxItems) { +// break; +// } +// int additionalCharacters = 1; // accounting for potential comma and colon +// if (processedBaggage != 0) { +// additionalCharacters = 2; // allocating space for comma +// } +// +// int byteSize = 1; // default include size of '=' +// if (additionalCharacters == 2) { +// baggageText.append(','); +// byteSize += 1; +// } +// +// byteSize += encodeKey(entry.getKey(), baggageText); +// baggageText.append('='); +// byteSize += encodeValue(entry.getValue(), baggageText); +// +// if (currentBytes + byteSize > maxBytes) { +// baggageText.setLength(currentBytes); +// break; +// } +// currentBytes += byteSize; +// processedBaggage++; +// } +// +// setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); +// } +// } +// +// public static HttpCodec.Extractor newExtractor( +// Config config, Supplier traceConfigSupplier) { +// return new TagContextExtractor( +// traceConfigSupplier, () -> new BaggageContextInterpreter(config)); +// } +// +// private static class BaggageContextInterpreter extends ContextInterpreter { +// +// private BaggageContextInterpreter(Config config) { +// super(config); +// } +// +// @Override +// public TracePropagationStyle style() { +// return BAGGAGE; +// } +// +// private Map parseBaggageHeaders(String input) { +// Map baggage = new HashMap<>(); +// char keyValueSeparator = '='; +// char pairSeparator = ','; +// int start = 0; +// +// int pairSeparatorInd = input.indexOf(pairSeparator); +// pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; +// int kvSeparatorInd = input.indexOf(keyValueSeparator); +// while (kvSeparatorInd != -1) { +// int end = pairSeparatorInd; +// if (kvSeparatorInd > end) { // value is missing +// return Collections.emptyMap(); +// } +// String key = HttpCodec.decode(input.substring(start, kvSeparatorInd).trim()); +// String value = HttpCodec.decode(input.substring(kvSeparatorInd + 1, end).trim()); +// if (key.isEmpty() || value.isEmpty()) { +// return Collections.emptyMap(); +// } +// baggage.put(key, value); +// +// kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1); +// pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1); +// pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; +// start = end + 1; +// } +// return baggage; +// } +// +// @Override +// public boolean accept(String key, String value) { +// if (null == key || key.isEmpty()) { +// return true; +// } +// if (LOG_EXTRACT_HEADER_NAMES) { +// log.debug("Header: {}", key); +// } +// +// if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage +// baggage = parseBaggageHeaders(value); +// } +// return true; +// } +// +// @Override +// protected TagContext build() { +// return super.build(); +// } +// } +// } 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 new file mode 100644 index 00000000000..147c4031125 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggagePropagator.java @@ -0,0 +1,178 @@ +package datadog.trace.core.baggage; + +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 java.io.UnsupportedEncodingException; +import java.net.URLDecoder; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; +import javax.annotation.ParametersAreNonnullByDefault; + +/** Propagator for tracing concern. */ +@ParametersAreNonnullByDefault +public class BaggagePropagator implements Propagator { + static final String BAGGAGE_KEY = "baggage"; + + private static final Set UNSAFE_CHARACTERS_KEY = + new HashSet<>( + Arrays.asList( + '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', + '}')); + private static final Set UNSAFE_CHARACTERS_VALUE = + new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); + + public BaggagePropagator() {} + + private int encodeKey(String key, StringBuilder builder) { + return encode(key, builder, UNSAFE_CHARACTERS_KEY); + } + + private int encodeValue(String key, StringBuilder builder) { + return encode(key, builder, UNSAFE_CHARACTERS_VALUE); + } + + private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { + int size = 0; + for (int i = 0; i < input.length(); i++) { + char c = input.charAt(i); + if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character + byte[] bytes = + Character.toString(c) + .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does + for (byte b : bytes) { + builder.append(String.format("%%%02X", b & 0xFF)); + size += 1; + } + } else { + builder.append(c); + size += 1; + } + } + return size; + } + + @Override + public void inject(Context context, C carrier, CarrierSetter setter) { + Config config = Config.get(); + + StringBuilder baggageText = new StringBuilder(); + int processedBaggage = 0; + int currentBytes = 0; + int maxItems = config.getTraceBaggageMaxItems(); + int maxBytes = config.getTraceBaggageMaxBytes(); + BaggageContext baggageContext = BaggageContext.fromContext(context); + for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { + if (processedBaggage >= maxItems) { + break; + } + int additionalCharacters = 1; // accounting for potential comma and colon + if (processedBaggage != 0) { + additionalCharacters = 2; // allocating space for comma + } + + int byteSize = 1; // default include size of '=' + if (additionalCharacters == 2) { + baggageText.append(','); + byteSize += 1; + } + + byteSize += encodeKey(entry.getKey(), baggageText); + baggageText.append('='); + byteSize += encodeValue(entry.getValue(), baggageText); + + if (currentBytes + byteSize > maxBytes) { + baggageText.setLength(currentBytes); + break; + } + currentBytes += byteSize; + processedBaggage++; + } + + setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); + } + + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + //noinspection ConstantValue + if (context == null || carrier == null || visitor == null) { + return context; + } + BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor(); + visitor.forEachKeyValue( + carrier, baggageContextExtractor); // If the extraction fails, return the original context + BaggageContext extractedContext = baggageContextExtractor.extractedContext; + if (extractedContext == null) { + return context; + } + return extractedContext.storeInto(context); // does this make sense? + } + + public static class BaggageContextExtractor implements BiConsumer { + private BaggageContext extractedContext; + + BaggageContextExtractor() {} + + /** URL decode value */ + private String decode(final String value) { + String decoded = value; + try { + decoded = URLDecoder.decode(value, "UTF-8"); + } catch (final UnsupportedEncodingException | IllegalArgumentException e) { + } + return decoded; + } + + private Map parseBaggageHeaders(String input) { + Map baggage = new HashMap<>(); + char keyValueSeparator = '='; + char pairSeparator = ','; + int start = 0; + + int pairSeparatorInd = input.indexOf(pairSeparator); + pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; + int kvSeparatorInd = input.indexOf(keyValueSeparator); + while (kvSeparatorInd != -1) { + int end = pairSeparatorInd; + if (kvSeparatorInd > end) { // value is missing + return Collections.emptyMap(); + } + String key = decode(input.substring(start, kvSeparatorInd).trim()); + String value = decode(input.substring(kvSeparatorInd + 1, end).trim()); + if (key.isEmpty() || value.isEmpty()) { + return Collections.emptyMap(); + } + baggage.put(key, value); + + kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1); + pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1); + pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; + start = end + 1; + } + return baggage; + } + + @Override + public void accept(String key, String value) { + if (null == key || key.isEmpty()) { + return; + } + if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage + extractedContext = + BaggageContext.create( + parseBaggageHeaders( + value)); // is this correct to assume that one instance will be created for each + // propagator? + } + } + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java deleted file mode 100644 index 4ded8742ad8..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/BaggageHttpCodec.java +++ /dev/null @@ -1,189 +0,0 @@ -package datadog.trace.core.propagation; - -import static datadog.trace.api.TracePropagationStyle.BAGGAGE; - -import datadog.trace.api.Config; -import datadog.trace.api.TraceConfig; -import datadog.trace.api.TracePropagationStyle; -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -import datadog.trace.bootstrap.instrumentation.api.TagContext; -import datadog.trace.core.DDSpanContext; -import java.nio.charset.StandardCharsets; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Map; -import java.util.Set; -import java.util.function.Supplier; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** A codec designed for HTTP transport via headers using Datadog headers */ -class BaggageHttpCodec { - private static final Logger log = LoggerFactory.getLogger(BaggageHttpCodec.class); - - static final String BAGGAGE_KEY = "baggage"; - private static final int MAX_CHARACTER_SIZE = 4; - - private static final Set UNSAFE_CHARACTERS_KEY = - new HashSet<>( - Arrays.asList( - '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', - '}')); - private static final Set UNSAFE_CHARACTERS_VALUE = - new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); - - private BaggageHttpCodec() { - // This class should not be created. This also makes code coverage checks happy. - } - - public static HttpCodec.Injector newInjector(Map invertedBaggageMapping) { - return new Injector(invertedBaggageMapping); - } - - private static class Injector implements HttpCodec.Injector { - - private final Map invertedBaggageMapping; - - public Injector(Map invertedBaggageMapping) { - assert invertedBaggageMapping != null; - this.invertedBaggageMapping = invertedBaggageMapping; - } - - private int encodeKey(String key, StringBuilder builder) { - return encode(key, builder, UNSAFE_CHARACTERS_KEY); - } - - private int encodeValue(String key, StringBuilder builder) { - return encode(key, builder, UNSAFE_CHARACTERS_VALUE); - } - - private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { - int size = 0; - for (int i = 0; i < input.length(); i++) { - char c = input.charAt(i); - if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character - byte[] bytes = - Character.toString(c) - .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does - for (byte b : bytes) { - builder.append(String.format("%%%02X", b & 0xFF)); - size += 1; - } - } else { - builder.append(c); - size += 1; - } - } - return size; - } - - @Override - public void inject( - final DDSpanContext context, final C carrier, final AgentPropagation.Setter setter) { - Config config = Config.get(); - - StringBuilder baggageText = new StringBuilder(); - int processedBaggage = 0; - int currentBytes = 0; - int maxItems = config.getTraceBaggageMaxItems(); - int maxBytes = config.getTraceBaggageMaxBytes(); - for (final Map.Entry entry : context.baggageItems()) { - if (processedBaggage >= maxItems) { - break; - } - int additionalCharacters = 1; // accounting for potential comma and colon - if (processedBaggage != 0) { - additionalCharacters = 2; // allocating space for comma - } - - int byteSize = 1; // default include size of '=' - if (additionalCharacters == 2) { - baggageText.append(','); - byteSize += 1; - } - - byteSize += encodeKey(entry.getKey(), baggageText); - baggageText.append('='); - byteSize += encodeValue(entry.getValue(), baggageText); - - if (currentBytes + byteSize > maxBytes) { - baggageText.setLength(currentBytes); - break; - } - currentBytes += byteSize; - processedBaggage++; - } - - setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); - } - } - - public static HttpCodec.Extractor newExtractor( - Config config, Supplier traceConfigSupplier) { - return new TagContextExtractor( - traceConfigSupplier, () -> new BaggageContextInterpreter(config)); - } - - private static class BaggageContextInterpreter extends ContextInterpreter { - - private BaggageContextInterpreter(Config config) { - super(config); - } - - @Override - public TracePropagationStyle style() { - return BAGGAGE; - } - - private Map parseBaggageHeaders(String input) { - Map baggage = new HashMap<>(); - char keyValueSeparator = '='; - char pairSeparator = ','; - int start = 0; - - int pairSeparatorInd = input.indexOf(pairSeparator); - pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; - int kvSeparatorInd = input.indexOf(keyValueSeparator); - while (kvSeparatorInd != -1) { - int end = pairSeparatorInd; - if (kvSeparatorInd > end) { // value is missing - return Collections.emptyMap(); - } - String key = HttpCodec.decode(input.substring(start, kvSeparatorInd).trim()); - String value = HttpCodec.decode(input.substring(kvSeparatorInd + 1, end).trim()); - if (key.isEmpty() || value.isEmpty()) { - return Collections.emptyMap(); - } - baggage.put(key, value); - - kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1); - pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1); - pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; - start = end + 1; - } - return baggage; - } - - @Override - public boolean accept(String key, String value) { - if (null == key || key.isEmpty()) { - return true; - } - if (LOG_EXTRACT_HEADER_NAMES) { - log.debug("Header: {}", key); - } - - if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage - baggage = parseBaggageHeaders(value); - } - return true; - } - - @Override - protected TagContext build() { - return super.build(); - } - } -} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 726bb9f8d14..120d14f7ee5 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -125,9 +125,6 @@ private static Map createInjectors( case TRACECONTEXT: result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping)); break; - case BAGGAGE: - result.put(style, BaggageHttpCodec.newInjector(reverseBaggageMapping)); - break; default: log.debug("No implementation found to inject propagation style: {}", style); break; @@ -162,9 +159,6 @@ public static Extractor createExtractor( case TRACECONTEXT: extractors.add(W3CHttpCodec.newExtractor(config, traceConfigSupplier)); break; - case BAGGAGE: - extractors.add(BaggageHttpCodec.newExtractor(config, traceConfigSupplier)); - break; default: log.debug("No implementation found to extract propagation style: {}", style); break; 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 new file mode 100644 index 00000000000..808837bcfca --- /dev/null +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/baggage/BaggagePropagatorTest.groovy @@ -0,0 +1,152 @@ +package datadog.trace.core.baggage + +import datadog.context.Context +import datadog.context.propagation.CarrierSetter +import datadog.context.propagation.Propagators +import datadog.trace.api.sampling.PrioritySampling +import datadog.trace.bootstrap.instrumentation.api.AgentPropagation +import datadog.trace.common.writer.LoggingWriter +import datadog.trace.core.ControllableSampler +import datadog.trace.core.test.DDCoreSpecification + +import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP +import static datadog.trace.api.sampling.PrioritySampling.USER_DROP +import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY + +class BaggagePropagatorTest extends DDCoreSpecification { + BaggagePropagator propagator + + def setup() { + this.propagator = new BaggagePropagator() + } + + def 'test tracing propagator context injection'() { + setup: + def tracer = tracerBuilder().build() + def span = tracer.buildSpan('test', 'operation').start() + def setter = Mock(CarrierSetter) + def carrier = new Object() + + when: + this.propagator.inject(span, carrier, setter) + + then: + 1 * injector.inject(span.context(), carrier, _) + + cleanup: + span.finish() + tracer.close() + } + + def 'test tracing propagator context extractor'() { + setup: + def context = Context.root() + // TODO Use ContextVisitor mock as getter once extractor API is refactored + def getter = Mock(AgentPropagation.ContextVisitor) + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + // BaggagePropagator.BaggageContextExtractor baggageContextExtractor = new BaggagePropagator.BaggageContextExtractor() + + + when: + this.propagator.extract(context, headers, getter) + // getter.forEachKeyValue(headers, baggageContextExtractor) + + then: + 1*_ + 1*_ + 1*_ + 1*_ + 1*_ + + where: + baggageHeader | baggageMap + "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] + "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] + + // 1 * extractor.extract(carrier, _) + } + + def 'span priority set when injecting'() { + given: + injectSysConfig('writer.type', 'LoggingWriter') + def tracer = tracerBuilder().build() + def setter = Mock(CarrierSetter) + def carrier = new Object() + + when: + def root = tracer.buildSpan('test', 'parent').start() + def child = tracer.buildSpan('test', 'child').asChildOf(root).start() + Propagators.defaultPropagator().inject(child, carrier, setter) + + then: + root.getSamplingPriority() == SAMPLER_KEEP as int + child.getSamplingPriority() == root.getSamplingPriority() + 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) + + cleanup: + child.finish() + root.finish() + tracer.close() + } + + def 'span priority only set after first injection'() { + given: + def sampler = new ControllableSampler() + def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build() + def setter = Mock(AgentPropagation.Setter) + def carrier = new Object() + + when: + def root = tracer.buildSpan('test', 'parent').start() + def child = tracer.buildSpan('test', 'child').asChildOf(root).start() + Propagators.defaultPropagator().inject(child, carrier, setter) + + then: + root.getSamplingPriority() == SAMPLER_KEEP as int + child.getSamplingPriority() == root.getSamplingPriority() + 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) + + when: + sampler.nextSamplingPriority = PrioritySampling.SAMPLER_DROP as int + def child2 = tracer.buildSpan('test', 'child2').asChildOf(root).start() + Propagators.defaultPropagator().inject(child2, carrier, setter) + + then: + root.getSamplingPriority() == SAMPLER_KEEP as int + child.getSamplingPriority() == root.getSamplingPriority() + child2.getSamplingPriority() == root.getSamplingPriority() + 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) + + cleanup: + child.finish() + child2.finish() + root.finish() + tracer.close() + } + + def 'injection does not override set priority'() { + given: + def sampler = new ControllableSampler() + def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build() + def setter = Mock(AgentPropagation.Setter) + def carrier = new Object() + + when: + def root = tracer.buildSpan('test', 'root').start() + def child = tracer.buildSpan('test', 'child').asChildOf(root).start() + child.setSamplingPriority(USER_DROP) + Propagators.defaultPropagator().inject(child, carrier, setter) + + then: + root.getSamplingPriority() == USER_DROP as int + child.getSamplingPriority() == root.getSamplingPriority() + 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(USER_DROP)) + + cleanup: + child.finish() + root.finish() + tracer.close() + } +} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy index b2541d2871f..7ecddc426ce 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy @@ -4,8 +4,9 @@ import datadog.trace.api.Config import datadog.trace.api.DynamicConfig import datadog.trace.bootstrap.instrumentation.api.TagContext import datadog.trace.bootstrap.instrumentation.api.ContextVisitors +//import datadog.trace.core.baggage.BaggageHttpCodec import datadog.trace.test.util.DDSpecification -import static datadog.trace.core.propagation.BaggageHttpCodec.BAGGAGE_KEY +//import static datadog.trace.core.baggage.BaggageHttpCodec.BAGGAGE_KEY class BaggageHttpExtractorTest extends DDSpecification { diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy index adbac85489a..34b56a917df 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy @@ -2,7 +2,6 @@ package datadog.trace.core.propagation import datadog.trace.api.DDSpanId import datadog.trace.api.DDTraceId -import datadog.trace.api.DynamicConfig import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopPathwayContext import datadog.trace.common.writer.ListWriter import datadog.trace.core.CoreTracer @@ -11,7 +10,7 @@ import datadog.trace.core.test.DDCoreSpecification import static datadog.trace.api.sampling.PrioritySampling.* -import static datadog.trace.core.propagation.BaggageHttpCodec.* +//import static datadog.trace.core.baggage.BaggageHttpCodec.* class BaggageHttpInjectorTest extends DDCoreSpecification { diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java index 0fd8cbca94f..3496a641686 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java @@ -14,12 +14,12 @@ public interface AgentPropagation { Concern TRACING_CONCERN = named("tracing"); + Concern BAGGAGE_CONCERN = Concern.named("baggage"); Concern XRAY_TRACING_CONCERN = named("tracing-xray"); Concern STANDALONE_ASM_CONCERN = named("asm-standalone"); // TODO DSM propagator should run after the other propagators as it stores the pathway context // TODO into the span context for now. Remove priority after the migration is complete. Concern DSM_CONCERN = withPriority("data-stream-monitoring", 110); - // The input tags should be sorted. void injectPathwayContext( AgentSpan span, C carrier, Setter setter, LinkedHashMap sortedTags); From 73e038446d504345424a89f39b39557a0eb0fea8 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 10 Feb 2025 15:41:44 -0500 Subject: [PATCH 10/31] updating tests for inject/extract rebasing --- .../java/datadog/trace/core/CoreTracer.java | 5 +- .../trace/core/baggage/BaggagePropagator.java | 2 +- .../core/baggage/BaggagePropagatorTest.groovy | 121 ++---------------- 3 files changed, 18 insertions(+), 110 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index bd8329352c6..cfa5c7b2158 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -737,7 +737,10 @@ private CoreTracer( if (dsm) { Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator()); } - Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator()); + if (config.getTracePropagationStylesToExtract().contains(TracePropagationStyle.BAGGAGE) + && config.getTracePropagationStylesToInject().contains(TracePropagationStyle.BAGGAGE)) { + Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator()); + } this.tagInterceptor = null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor; 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 147c4031125..d82383d55cb 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 @@ -114,7 +114,7 @@ public Context extract(Context context, C carrier, CarrierVisitor visitor if (extractedContext == null) { return context; } - return extractedContext.storeInto(context); // does this make sense? + return extractedContext.storeInto(context); } public static class BaggageContextExtractor implements BiConsumer { 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 808837bcfca..c8334b7aeb1 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 @@ -2,15 +2,9 @@ package datadog.trace.core.baggage import datadog.context.Context import datadog.context.propagation.CarrierSetter -import datadog.context.propagation.Propagators -import datadog.trace.api.sampling.PrioritySampling -import datadog.trace.bootstrap.instrumentation.api.AgentPropagation -import datadog.trace.common.writer.LoggingWriter -import datadog.trace.core.ControllableSampler +import datadog.trace.bootstrap.instrumentation.api.BaggageContext +import datadog.trace.bootstrap.instrumentation.api.ContextVisitors import datadog.trace.core.test.DDCoreSpecification - -import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP -import static datadog.trace.api.sampling.PrioritySampling.USER_DROP import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY class BaggagePropagatorTest extends DDCoreSpecification { @@ -22,131 +16,42 @@ class BaggagePropagatorTest extends DDCoreSpecification { def 'test tracing propagator context injection'() { setup: - def tracer = tracerBuilder().build() - def span = tracer.buildSpan('test', 'operation').start() def setter = Mock(CarrierSetter) def carrier = new Object() + def context = Context.root() + context = BaggageContext.create(baggageMap).storeInto(context) + when: - this.propagator.inject(span, carrier, setter) + this.propagator.inject(context, carrier, setter) then: - 1 * injector.inject(span.context(), carrier, _) + 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) - cleanup: - span.finish() - tracer.close() + where: + baggageMap | baggageHeader + ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y" + ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" } def 'test tracing propagator context extractor'() { setup: def context = Context.root() // TODO Use ContextVisitor mock as getter once extractor API is refactored - def getter = Mock(AgentPropagation.ContextVisitor) def headers = [ (BAGGAGE_KEY) : baggageHeader, ] - // BaggagePropagator.BaggageContextExtractor baggageContextExtractor = new BaggagePropagator.BaggageContextExtractor() - when: - this.propagator.extract(context, headers, getter) - // getter.forEachKeyValue(headers, baggageContextExtractor) + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - 1*_ - 1*_ - 1*_ - 1*_ - 1*_ + BaggageContext.fromContext(context).getBaggage() == baggageMap where: baggageHeader | baggageMap "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] - // 1 * extractor.extract(carrier, _) - } - - def 'span priority set when injecting'() { - given: - injectSysConfig('writer.type', 'LoggingWriter') - def tracer = tracerBuilder().build() - def setter = Mock(CarrierSetter) - def carrier = new Object() - - when: - def root = tracer.buildSpan('test', 'parent').start() - def child = tracer.buildSpan('test', 'child').asChildOf(root).start() - Propagators.defaultPropagator().inject(child, carrier, setter) - - then: - root.getSamplingPriority() == SAMPLER_KEEP as int - child.getSamplingPriority() == root.getSamplingPriority() - 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) - - cleanup: - child.finish() - root.finish() - tracer.close() - } - - def 'span priority only set after first injection'() { - given: - def sampler = new ControllableSampler() - def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build() - def setter = Mock(AgentPropagation.Setter) - def carrier = new Object() - - when: - def root = tracer.buildSpan('test', 'parent').start() - def child = tracer.buildSpan('test', 'child').asChildOf(root).start() - Propagators.defaultPropagator().inject(child, carrier, setter) - - then: - root.getSamplingPriority() == SAMPLER_KEEP as int - child.getSamplingPriority() == root.getSamplingPriority() - 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) - - when: - sampler.nextSamplingPriority = PrioritySampling.SAMPLER_DROP as int - def child2 = tracer.buildSpan('test', 'child2').asChildOf(root).start() - Propagators.defaultPropagator().inject(child2, carrier, setter) - - then: - root.getSamplingPriority() == SAMPLER_KEEP as int - child.getSamplingPriority() == root.getSamplingPriority() - child2.getSamplingPriority() == root.getSamplingPriority() - 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP)) - - cleanup: - child.finish() - child2.finish() - root.finish() - tracer.close() - } - - def 'injection does not override set priority'() { - given: - def sampler = new ControllableSampler() - def tracer = tracerBuilder().writer(new LoggingWriter()).sampler(sampler).build() - def setter = Mock(AgentPropagation.Setter) - def carrier = new Object() - - when: - def root = tracer.buildSpan('test', 'root').start() - def child = tracer.buildSpan('test', 'child').asChildOf(root).start() - child.setSamplingPriority(USER_DROP) - Propagators.defaultPropagator().inject(child, carrier, setter) - - then: - root.getSamplingPriority() == USER_DROP as int - child.getSamplingPriority() == root.getSamplingPriority() - 1 * setter.set(carrier, DatadogHttpCodec.SAMPLING_PRIORITY_KEY, String.valueOf(USER_DROP)) - - cleanup: - child.finish() - root.finish() - tracer.close() } } From 5ef93574b4187287a311fc19f636d7cc21951403 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 10 Feb 2025 16:09:41 -0500 Subject: [PATCH 11/31] finishing migration to baggage propagator --- .../trace/core/baggage/BaggageHttpCodec.java | 194 ------------------ .../trace/core/baggage/BaggagePropagator.java | 12 +- .../core/baggage/BaggagePropagatorTest.groovy | 69 ++++++- .../BaggageHttpExtractorTest.groovy | 78 ------- .../BaggageHttpInjectorTest.groovy | 155 -------------- 5 files changed, 75 insertions(+), 433 deletions(-) delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java delete mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy delete mode 100644 dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy diff --git a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java deleted file mode 100644 index 402da477aaa..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/baggage/BaggageHttpCodec.java +++ /dev/null @@ -1,194 +0,0 @@ -// package datadog.trace.core.baggage; -// -// import static datadog.trace.api.TracePropagationStyle.BAGGAGE; -// -// import datadog.trace.api.Config; -// import datadog.trace.api.TraceConfig; -// import datadog.trace.api.TracePropagationStyle; -// import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; -// import datadog.trace.bootstrap.instrumentation.api.TagContext; -// import datadog.trace.core.DDSpanContext; -// import java.nio.charset.StandardCharsets; -// import java.util.Arrays; -// import java.util.Collections; -// import java.util.HashMap; -// import java.util.HashSet; -// import java.util.Map; -// import java.util.Set; -// import java.util.function.Supplier; -// -// import datadog.trace.core.propagation.ContextInterpreter; -// import datadog.trace.core.propagation.HttpCodec; -// import datadog.trace.core.propagation.TagContextExtractor; -// import org.slf4j.Logger; -// import org.slf4j.LoggerFactory; -// -/// ** A codec designed for HTTP transport via headers using Datadog headers */ -// class BaggageHttpCodec { -// private static final Logger log = LoggerFactory.getLogger(BaggageHttpCodec.class); -// -// static final String BAGGAGE_KEY = "baggage"; -// private static final int MAX_CHARACTER_SIZE = 4; -// -// private static final Set UNSAFE_CHARACTERS_KEY = -// new HashSet<>( -// Arrays.asList( -// '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', -// '}')); -// private static final Set UNSAFE_CHARACTERS_VALUE = -// new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); -// -// private BaggageHttpCodec() { -// // This class should not be created. This also makes code coverage checks happy. -// } -// -// public static HttpCodec.Injector newInjector(Map invertedBaggageMapping) { -// return new Injector(invertedBaggageMapping); -// } -// -// private static class Injector implements HttpCodec.Injector { -// -// private final Map invertedBaggageMapping; -// -// public Injector(Map invertedBaggageMapping) { -// assert invertedBaggageMapping != null; -// this.invertedBaggageMapping = invertedBaggageMapping; -// } -// -// private int encodeKey(String key, StringBuilder builder) { -// return encode(key, builder, UNSAFE_CHARACTERS_KEY); -// } -// -// private int encodeValue(String key, StringBuilder builder) { -// return encode(key, builder, UNSAFE_CHARACTERS_VALUE); -// } -// -// private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { -// int size = 0; -// for (int i = 0; i < input.length(); i++) { -// char c = input.charAt(i); -// if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character -// byte[] bytes = -// Character.toString(c) -// .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder -// does -// for (byte b : bytes) { -// builder.append(String.format("%%%02X", b & 0xFF)); -// size += 1; -// } -// } else { -// builder.append(c); -// size += 1; -// } -// } -// return size; -// } -// -// @Override -// public void inject( -// final DDSpanContext context, final C carrier, final AgentPropagation.Setter setter) { -// Config config = Config.get(); -// -// StringBuilder baggageText = new StringBuilder(); -// int processedBaggage = 0; -// int currentBytes = 0; -// int maxItems = config.getTraceBaggageMaxItems(); -// int maxBytes = config.getTraceBaggageMaxBytes(); -// for (final Map.Entry entry : context.baggageItems()) { -// if (processedBaggage >= maxItems) { -// break; -// } -// int additionalCharacters = 1; // accounting for potential comma and colon -// if (processedBaggage != 0) { -// additionalCharacters = 2; // allocating space for comma -// } -// -// int byteSize = 1; // default include size of '=' -// if (additionalCharacters == 2) { -// baggageText.append(','); -// byteSize += 1; -// } -// -// byteSize += encodeKey(entry.getKey(), baggageText); -// baggageText.append('='); -// byteSize += encodeValue(entry.getValue(), baggageText); -// -// if (currentBytes + byteSize > maxBytes) { -// baggageText.setLength(currentBytes); -// break; -// } -// currentBytes += byteSize; -// processedBaggage++; -// } -// -// setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); -// } -// } -// -// public static HttpCodec.Extractor newExtractor( -// Config config, Supplier traceConfigSupplier) { -// return new TagContextExtractor( -// traceConfigSupplier, () -> new BaggageContextInterpreter(config)); -// } -// -// private static class BaggageContextInterpreter extends ContextInterpreter { -// -// private BaggageContextInterpreter(Config config) { -// super(config); -// } -// -// @Override -// public TracePropagationStyle style() { -// return BAGGAGE; -// } -// -// private Map parseBaggageHeaders(String input) { -// Map baggage = new HashMap<>(); -// char keyValueSeparator = '='; -// char pairSeparator = ','; -// int start = 0; -// -// int pairSeparatorInd = input.indexOf(pairSeparator); -// pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; -// int kvSeparatorInd = input.indexOf(keyValueSeparator); -// while (kvSeparatorInd != -1) { -// int end = pairSeparatorInd; -// if (kvSeparatorInd > end) { // value is missing -// return Collections.emptyMap(); -// } -// String key = HttpCodec.decode(input.substring(start, kvSeparatorInd).trim()); -// String value = HttpCodec.decode(input.substring(kvSeparatorInd + 1, end).trim()); -// if (key.isEmpty() || value.isEmpty()) { -// return Collections.emptyMap(); -// } -// baggage.put(key, value); -// -// kvSeparatorInd = input.indexOf(keyValueSeparator, pairSeparatorInd + 1); -// pairSeparatorInd = input.indexOf(pairSeparator, pairSeparatorInd + 1); -// pairSeparatorInd = pairSeparatorInd == -1 ? input.length() : pairSeparatorInd; -// start = end + 1; -// } -// return baggage; -// } -// -// @Override -// public boolean accept(String key, String value) { -// if (null == key || key.isEmpty()) { -// return true; -// } -// if (LOG_EXTRACT_HEADER_NAMES) { -// log.debug("Header: {}", key); -// } -// -// if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage -// baggage = parseBaggageHeaders(value); -// } -// return true; -// } -// -// @Override -// protected TagContext build() { -// return super.build(); -// } -// } -// } 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 d82383d55cb..7a1c158bc9c 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 @@ -167,11 +167,13 @@ public void accept(String key, String value) { return; } if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage - extractedContext = - BaggageContext.create( - parseBaggageHeaders( - value)); // is this correct to assume that one instance will be created for each - // propagator? + Map baggage = parseBaggageHeaders(value); + if (!baggage.isEmpty()) { + extractedContext = + BaggageContext.create( + baggage); // is this correct to assume that one instance will be created for each + // propagator? + } } } } 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 c8334b7aeb1..2bada620692 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 @@ -1,6 +1,7 @@ package datadog.trace.core.baggage import datadog.context.Context +import datadog.context.EmptyContext import datadog.context.propagation.CarrierSetter import datadog.trace.bootstrap.instrumentation.api.BaggageContext import datadog.trace.bootstrap.instrumentation.api.ContextVisitors @@ -21,7 +22,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { def context = Context.root() context = BaggageContext.create(baggageMap).storeInto(context) - when: this.propagator.inject(context, carrier, setter) @@ -32,6 +32,53 @@ class BaggagePropagatorTest extends DDCoreSpecification { baggageMap | baggageHeader ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y" ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" + [key1: "val1"] | "key1=val1" + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [serverNode: "DF 28"] | "serverNode=DF%2028" + [userId: "Amélie"] | "userId=Am%C3%A9lie" + ["user!d(me)": "false"] | "user!d%28me%29=false" + ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" + } + + def "test baggage item limit"() { + setup: + injectSysConfig("trace.baggage.max.items", '2') + def setter = Mock(CarrierSetter) + def carrier = new Object() + def context = Context.root() + context = BaggageContext.create(baggage).storeInto(context) + + when: + this.propagator.inject(context, carrier, setter) + + then: + 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) + + where: + baggage | baggageHeader + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + } + + def "test baggage bytes limit"() { + setup: + injectSysConfig("trace.baggage.max.bytes", '20') + def setter = Mock(CarrierSetter) + def carrier = new Object() + def context = Context.root() + context = BaggageContext.create(baggage).storeInto(context) + + when: + this.propagator.inject(context, carrier, setter) + + then: + 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) + + where: + baggage | baggageHeader + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + ["abcdefg": "hijklmnopq♥"] | "" } def 'test tracing propagator context extractor'() { @@ -52,6 +99,26 @@ class BaggagePropagatorTest extends DDCoreSpecification { baggageHeader | baggageMap "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] + } + def "extract invalid baggage headers"() { + setup: + def context = Context.root() + def headers = [ + (BAGGAGE_KEY) : baggageHeader, + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: + context instanceof EmptyContext + + where: + baggageHeader | baggageMap + "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | [] + "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | [] + "=no-key" | [] + "no-value=" | [] } } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy deleted file mode 100644 index 7ecddc426ce..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpExtractorTest.groovy +++ /dev/null @@ -1,78 +0,0 @@ -package datadog.trace.core.propagation - -import datadog.trace.api.Config -import datadog.trace.api.DynamicConfig -import datadog.trace.bootstrap.instrumentation.api.TagContext -import datadog.trace.bootstrap.instrumentation.api.ContextVisitors -//import datadog.trace.core.baggage.BaggageHttpCodec -import datadog.trace.test.util.DDSpecification -//import static datadog.trace.core.baggage.BaggageHttpCodec.BAGGAGE_KEY - -class BaggageHttpExtractorTest extends DDSpecification { - - private DynamicConfig dynamicConfig - private HttpCodec.Extractor _extractor - - private HttpCodec.Extractor getExtractor() { - _extractor ?: (_extractor = createExtractor(Config.get())) - } - - private HttpCodec.Extractor createExtractor(Config config) { - BaggageHttpCodec.newExtractor(config, { dynamicConfig.captureTraceConfig() }) - } - - void setup() { - dynamicConfig = DynamicConfig.create() - .apply() - } - - void cleanup() { - extractor.cleanup() - } - - def "extract valid baggage headers"() { - setup: - def extractor = createExtractor(Config.get()) - def headers = [ - (BAGGAGE_KEY) : baggageHeader, - ] - - when: - final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap()) - - then: - context.baggage == baggageMap - - cleanup: - extractor.cleanup() - - where: - baggageHeader | baggageMap - "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] - "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] - } - - def "extract invalid baggage headers"() { - setup: - def extractor = createExtractor(Config.get()) - def headers = [ - (BAGGAGE_KEY) : baggageHeader, - ] - - when: - final TagContext context = extractor.extract(headers, ContextVisitors.stringValuesMap()) - - then: - context == null - - cleanup: - extractor.cleanup() - - where: - baggageHeader | baggageMap - "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | [] - "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | [] - "=no-key" | [] - "no-value=" | [] - } -} diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy deleted file mode 100644 index 34b56a917df..00000000000 --- a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/BaggageHttpInjectorTest.groovy +++ /dev/null @@ -1,155 +0,0 @@ -package datadog.trace.core.propagation - -import datadog.trace.api.DDSpanId -import datadog.trace.api.DDTraceId -import datadog.trace.bootstrap.instrumentation.api.AgentTracer.NoopPathwayContext -import datadog.trace.common.writer.ListWriter -import datadog.trace.core.CoreTracer -import datadog.trace.core.DDSpanContext -import datadog.trace.core.test.DDCoreSpecification - - -import static datadog.trace.api.sampling.PrioritySampling.* -//import static datadog.trace.core.baggage.BaggageHttpCodec.* - - -class BaggageHttpInjectorTest extends DDCoreSpecification { - - private HttpCodec.Injector injector = newInjector(["some-baggage-key":"SOME_CUSTOM_HEADER"]) - private ListWriter writer - private CoreTracer tracer - private Map carrier - - void setup() { - writer = new ListWriter() - tracer = tracerBuilder().writer(writer).build() - carrier = Mock() - } - - def "test baggage injection and encoding"() { - setup: - final DDSpanContext mockedContext = - new DDSpanContext( - DDTraceId.from("1"), - DDSpanId.from("2"), - DDSpanId.ZERO, - null, - "fakeService", - "fakeOperation", - "fakeResource", - UNSET, - "fakeOrigin", - baggage, - false, - "fakeType", - 0, - tracer.traceCollectorFactory.create(DDTraceId.ONE), - null, - null, - NoopPathwayContext.INSTANCE, - false, - PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - when: - injector.inject(mockedContext, carrier, MapSetter.INSTANCE) - - then: - 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) - 0 * _ - - cleanup: - tracer.close() - - where: - baggage | baggageHeaders - [key1: "val1"] | "key1=val1" - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [serverNode: "DF 28"] | "serverNode=DF%2028" - [userId: "Amélie"] | "userId=Am%C3%A9lie" - ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" - ["user!d(me)": "false"] | "user!d%28me%29=false" - ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" - } - - def "test baggage item limit"() { - setup: - injectSysConfig("trace.baggage.max.items", '2') - final DDSpanContext mockedContext = - new DDSpanContext( - DDTraceId.from("1"), - DDSpanId.from("2"), - DDSpanId.ZERO, - null, - "fakeService", - "fakeOperation", - "fakeResource", - UNSET, - "fakeOrigin", - baggage, - false, - "fakeType", - 0, - tracer.traceCollectorFactory.create(DDTraceId.ONE), - null, - null, - NoopPathwayContext.INSTANCE, - false, - PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - - when: - injector.inject(mockedContext, carrier, MapSetter.INSTANCE) - - then: - 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) - 0 * _ - - cleanup: - tracer.close() - - where: - baggage | baggageHeaders - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" - } - - def "test baggage bytes limit"() { - setup: - injectSysConfig("trace.baggage.max.bytes", '20') - final DDSpanContext mockedContext = - new DDSpanContext( - DDTraceId.from("1"), - DDSpanId.from("2"), - DDSpanId.ZERO, - null, - "fakeService", - "fakeOperation", - "fakeResource", - UNSET, - "fakeOrigin", - baggage, - false, - "fakeType", - 0, - tracer.traceCollectorFactory.create(DDTraceId.ONE), - null, - null, - NoopPathwayContext.INSTANCE, - false, - PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=-4,_dd.p.anytag=value")) - - when: - injector.inject(mockedContext, carrier, MapSetter.INSTANCE) - - then: - 1 * carrier.put(BAGGAGE_KEY, baggageHeaders) - 0 * _ - - cleanup: - tracer.close() - - where: - baggage | baggageHeaders - // [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" - ["abcdefg": "hijklmnopq♥"] | "" - } -} From b53e449d9de6f11bd280ce69fc89d7d33297c8ca Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 11 Feb 2025 10:45:15 -0800 Subject: [PATCH 12/31] adding support for injector only or extractor only rebasing --- .../main/java/datadog/trace/core/CoreTracer.java | 9 +++++---- .../trace/core/baggage/BaggagePropagator.java | 14 +++++++++++--- .../core/baggage/BaggagePropagatorTest.groovy | 2 +- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index cfa5c7b2158..e0435b339be 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -737,10 +737,11 @@ private CoreTracer( if (dsm) { Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator()); } - if (config.getTracePropagationStylesToExtract().contains(TracePropagationStyle.BAGGAGE) - && config.getTracePropagationStylesToInject().contains(TracePropagationStyle.BAGGAGE)) { - Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator()); - } + Propagators.register( + BAGGAGE_CONCERN, + new BaggagePropagator( + config.getTracePropagationStylesToInject().contains(TracePropagationStyle.BAGGAGE), + config.getTracePropagationStylesToExtract().contains(TracePropagationStyle.BAGGAGE))); this.tagInterceptor = null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor; 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 7a1c158bc9c..4c61df33b11 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 @@ -22,7 +22,8 @@ @ParametersAreNonnullByDefault public class BaggagePropagator implements Propagator { static final String BAGGAGE_KEY = "baggage"; - + private final boolean injectBaggage; + private final boolean extractBaggage; private static final Set UNSAFE_CHARACTERS_KEY = new HashSet<>( Arrays.asList( @@ -31,7 +32,10 @@ public class BaggagePropagator implements Propagator { private static final Set UNSAFE_CHARACTERS_VALUE = new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); - public BaggagePropagator() {} + public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) { + this.injectBaggage = injectBaggage; + this.extractBaggage = extractBaggage; + } private int encodeKey(String key, StringBuilder builder) { return encode(key, builder, UNSAFE_CHARACTERS_KEY); @@ -63,6 +67,10 @@ private int encode(String input, StringBuilder builder, Set UNSAFE_CH @Override public void inject(Context context, C carrier, CarrierSetter setter) { + if (!this.injectBaggage) { + return; + } + Config config = Config.get(); StringBuilder baggageText = new StringBuilder(); @@ -104,7 +112,7 @@ public void inject(Context context, C carrier, CarrierSetter setter) { @Override public Context extract(Context context, C carrier, CarrierVisitor visitor) { //noinspection ConstantValue - if (context == null || carrier == null || visitor == null) { + if (!this.extractBaggage || context == null || carrier == null || visitor == null) { return context; } BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor(); 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 2bada620692..4d17447d29e 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 @@ -12,7 +12,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { BaggagePropagator propagator def setup() { - this.propagator = new BaggagePropagator() + this.propagator = new BaggagePropagator(true, true) } def 'test tracing propagator context injection'() { From 22d8a31c0cc03c377455ad7e83ea503cbe71cf23 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 11 Feb 2025 12:14:04 -0800 Subject: [PATCH 13/31] updating PR comments --- .../trace/core/baggage/BaggagePropagator.java | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) 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 4c61df33b11..b1a143ec4ee 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 @@ -54,50 +54,53 @@ private int encode(String input, StringBuilder builder, Set UNSAFE_CH Character.toString(c) .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does for (byte b : bytes) { - builder.append(String.format("%%%02X", b & 0xFF)); - size += 1; + builder.append('%'); + builder.append(encodeChar((b >> 4) & 0xF)); + builder.append(encodeChar(b & 0xF)); + size++; } } else { builder.append(c); - size += 1; + size++; } } return size; } + private char encodeChar(int ascii) { + return (char) (ascii < 10 ? '0' + ascii : 'A' + (ascii - 10)); + } + @Override public void inject(Context context, C carrier, CarrierSetter setter) { - if (!this.injectBaggage) { - return; - } - Config config = Config.get(); StringBuilder baggageText = new StringBuilder(); - int processedBaggage = 0; - int currentBytes = 0; int maxItems = config.getTraceBaggageMaxItems(); int maxBytes = config.getTraceBaggageMaxBytes(); + if (!this.injectBaggage || maxItems == 0 || maxBytes == 0) { + return; + } + + int processedBaggage = 0; + int currentBytes = 0; BaggageContext baggageContext = BaggageContext.fromContext(context); for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { - if (processedBaggage >= maxItems) { - break; - } - int additionalCharacters = 1; // accounting for potential comma and colon - if (processedBaggage != 0) { - additionalCharacters = 2; // allocating space for comma - } - int byteSize = 1; // default include size of '=' - if (additionalCharacters == 2) { + if (processedBaggage + != 0) { // if there are already baggage items processed, add and allocate bytes for a + // comma baggageText.append(','); - byteSize += 1; + byteSize++; } byteSize += encodeKey(entry.getKey(), baggageText); baggageText.append('='); byteSize += encodeValue(entry.getValue(), baggageText); + if (processedBaggage >= maxItems) { // reached the max number of baggage items allowed + break; + } if (currentBytes + byteSize > maxBytes) { baggageText.setLength(currentBytes); break; From a73bc872c8a7dee78dde9d90455e0aa652d47b0b Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 11 Feb 2025 13:47:46 -0800 Subject: [PATCH 14/31] cleanup --- .../trace/core/baggage/BaggagePropagator.java | 23 ++++++++----------- .../core/baggage/BaggagePropagatorTest.groovy | 18 +++++---------- 2 files changed, 16 insertions(+), 25 deletions(-) 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 b1a143ec4ee..3d681874d44 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 @@ -17,10 +17,12 @@ import java.util.Set; import java.util.function.BiConsumer; import javax.annotation.ParametersAreNonnullByDefault; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -/** Propagator for tracing concern. */ @ParametersAreNonnullByDefault public class BaggagePropagator implements Propagator { + private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class); static final String BAGGAGE_KEY = "baggage"; private final boolean injectBaggage; private final boolean extractBaggage; @@ -50,9 +52,7 @@ private int encode(String input, StringBuilder builder, Set UNSAFE_CH for (int i = 0; i < input.length(); i++) { char c = input.charAt(i); if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character - byte[] bytes = - Character.toString(c) - .getBytes(StandardCharsets.UTF_8); // not most efficient but what URLEncoder does + byte[] bytes = Character.toString(c).getBytes(StandardCharsets.UTF_8); for (byte b : bytes) { builder.append('%'); builder.append(encodeChar((b >> 4) & 0xF)); @@ -75,7 +75,6 @@ private char encodeChar(int ascii) { public void inject(Context context, C carrier, CarrierSetter setter) { Config config = Config.get(); - StringBuilder baggageText = new StringBuilder(); int maxItems = config.getTraceBaggageMaxItems(); int maxBytes = config.getTraceBaggageMaxBytes(); if (!this.injectBaggage || maxItems == 0 || maxBytes == 0) { @@ -84,6 +83,7 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int processedBaggage = 0; int currentBytes = 0; + StringBuilder baggageText = new StringBuilder(); BaggageContext baggageContext = BaggageContext.fromContext(context); for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { int byteSize = 1; // default include size of '=' @@ -98,7 +98,8 @@ public void inject(Context context, C carrier, CarrierSetter setter) { baggageText.append('='); byteSize += encodeValue(entry.getValue(), baggageText); - if (processedBaggage >= maxItems) { // reached the max number of baggage items allowed + processedBaggage++; + if (processedBaggage == maxItems) { // reached the max number of baggage items allowed break; } if (currentBytes + byteSize > maxBytes) { @@ -106,7 +107,6 @@ public void inject(Context context, C carrier, CarrierSetter setter) { break; } currentBytes += byteSize; - processedBaggage++; } setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); @@ -119,8 +119,7 @@ public Context extract(Context context, C carrier, CarrierVisitor visitor return context; } BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor(); - visitor.forEachKeyValue( - carrier, baggageContextExtractor); // If the extraction fails, return the original context + visitor.forEachKeyValue(carrier, baggageContextExtractor); BaggageContext extractedContext = baggageContextExtractor.extractedContext; if (extractedContext == null) { return context; @@ -139,6 +138,7 @@ private String decode(final String value) { try { decoded = URLDecoder.decode(value, "UTF-8"); } catch (final UnsupportedEncodingException | IllegalArgumentException e) { + log.debug("Failed to decode {}", value); } return decoded; } @@ -180,10 +180,7 @@ public void accept(String key, String value) { if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { - extractedContext = - BaggageContext.create( - baggage); // is this correct to assume that one instance will be created for each - // propagator? + extractedContext = BaggageContext.create(baggage); } } } 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 4d17447d29e..7d22431189f 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 @@ -10,16 +10,19 @@ import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY class BaggagePropagatorTest extends DDCoreSpecification { BaggagePropagator propagator + CarrierSetter setter + Object carrier + Context context def setup() { this.propagator = new BaggagePropagator(true, true) + setter = Mock(CarrierSetter) + carrier = new Object() + context = Context.root() } def 'test tracing propagator context injection'() { setup: - def setter = Mock(CarrierSetter) - def carrier = new Object() - def context = Context.root() context = BaggageContext.create(baggageMap).storeInto(context) when: @@ -43,9 +46,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { def "test baggage item limit"() { setup: injectSysConfig("trace.baggage.max.items", '2') - def setter = Mock(CarrierSetter) - def carrier = new Object() - def context = Context.root() context = BaggageContext.create(baggage).storeInto(context) when: @@ -63,9 +63,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { def "test baggage bytes limit"() { setup: injectSysConfig("trace.baggage.max.bytes", '20') - def setter = Mock(CarrierSetter) - def carrier = new Object() - def context = Context.root() context = BaggageContext.create(baggage).storeInto(context) when: @@ -83,8 +80,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { def 'test tracing propagator context extractor'() { setup: - def context = Context.root() - // TODO Use ContextVisitor mock as getter once extractor API is refactored def headers = [ (BAGGAGE_KEY) : baggageHeader, ] @@ -103,7 +98,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { def "extract invalid baggage headers"() { setup: - def context = Context.root() def headers = [ (BAGGAGE_KEY) : baggageHeader, ] From b37dcbd19fc235ecab6d93988bdc60d676c9610c Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 11 Feb 2025 14:36:47 -0800 Subject: [PATCH 15/31] adding null context check --- .../java/datadog/trace/core/baggage/BaggagePropagator.java | 4 ++++ 1 file changed, 4 insertions(+) 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 3d681874d44..f15d5e443a7 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 @@ -85,6 +85,10 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); BaggageContext baggageContext = BaggageContext.fromContext(context); + if (baggageContext == null) { + log.debug("BaggageContext instance is null in the following context {}", context); + return; + } for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { int byteSize = 1; // default include size of '=' if (processedBaggage From 36940e8a2873de8b5f32d09d5e865141688ace8e Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Wed, 12 Feb 2025 16:52:20 -0800 Subject: [PATCH 16/31] updating PR comments rebasing --- .../trace/api/TracePropagationStyle.java | 2 +- .../java/datadog/trace/core/CoreTracer.java | 8 ++--- .../trace/core/baggage/BaggagePropagator.java | 33 +++++++++++-------- .../core/baggage/BaggagePropagatorTest.groovy | 11 +++++-- .../main/java/datadog/trace/api/Config.java | 12 +++++++ 5 files changed, 44 insertions(+), 22 deletions(-) diff --git a/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java b/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java index 2e9e1cf3b79..af850871650 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationStyle.java @@ -21,7 +21,7 @@ public enum TracePropagationStyle { // W3C trace context propagation style // https://www.w3.org/TR/trace-context-1/ TRACECONTEXT, - // OTEL baggage support + // W3C baggage support // https://www.w3.org/TR/baggage/ BAGGAGE, // None does not extract or inject diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index e0435b339be..4254bfdedeb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -737,11 +737,9 @@ private CoreTracer( if (dsm) { Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator()); } - Propagators.register( - BAGGAGE_CONCERN, - new BaggagePropagator( - config.getTracePropagationStylesToInject().contains(TracePropagationStyle.BAGGAGE), - config.getTracePropagationStylesToExtract().contains(TracePropagationStyle.BAGGAGE))); + if (config.isBaggagePropagationEnabled()) { + Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator(config)); + } this.tagInterceptor = null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor; 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 f15d5e443a7..5a430602907 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 @@ -24,19 +24,28 @@ public class BaggagePropagator implements Propagator { private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class); static final String BAGGAGE_KEY = "baggage"; + private Config config; private final boolean injectBaggage; private final boolean extractBaggage; private static final Set UNSAFE_CHARACTERS_KEY = new HashSet<>( Arrays.asList( - '\"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', + '"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', '}')); private static final Set UNSAFE_CHARACTERS_VALUE = new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); + public BaggagePropagator(Config config) { + this.injectBaggage = config.isBaggageInject(); + this.extractBaggage = config.isBaggageExtract(); + this.config = config; + } + + // use primarily for testing purposes public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) { this.injectBaggage = injectBaggage; this.extractBaggage = extractBaggage; + config = Config.get(); } private int encodeKey(String key, StringBuilder builder) { @@ -51,7 +60,7 @@ private int encode(String input, StringBuilder builder, Set UNSAFE_CH int size = 0; for (int i = 0; i < input.length(); i++) { char c = input.charAt(i); - if (UNSAFE_CHARACTERS.contains(c) || c > 126 || c <= 32) { // encode character + if (UNSAFE_CHARACTERS.contains(c) || c > '~' || c <= ' ') { // encode character byte[] bytes = Character.toString(c).getBytes(StandardCharsets.UTF_8); for (byte b : bytes) { builder.append('%'); @@ -73,7 +82,11 @@ private char encodeChar(int ascii) { @Override public void inject(Context context, C carrier, CarrierSetter setter) { - Config config = Config.get(); + BaggageContext baggageContext = BaggageContext.fromContext(context); + if (baggageContext == null) { + log.debug("BaggageContext instance is missing from the following context {}", context); + return; + } int maxItems = config.getTraceBaggageMaxItems(); int maxBytes = config.getTraceBaggageMaxBytes(); @@ -84,11 +97,6 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int processedBaggage = 0; int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); - BaggageContext baggageContext = BaggageContext.fromContext(context); - if (baggageContext == null) { - log.debug("BaggageContext instance is null in the following context {}", context); - return; - } for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { int byteSize = 1; // default include size of '=' if (processedBaggage @@ -106,7 +114,8 @@ public void inject(Context context, C carrier, CarrierSetter setter) { if (processedBaggage == maxItems) { // reached the max number of baggage items allowed break; } - if (currentBytes + byteSize > maxBytes) { + if (currentBytes + byteSize + > maxBytes) { // Drop newest k/v pair if adding it leads to exceeding the limit baggageText.setLength(currentBytes); break; } @@ -178,10 +187,8 @@ private Map parseBaggageHeaders(String input) { @Override public void accept(String key, String value) { - if (null == key || key.isEmpty()) { - return; - } - if (key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage + if (key != null + && key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { extractedContext = BaggageContext.create(baggage); 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 7d22431189f..fcd701f0752 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 @@ -11,17 +11,17 @@ import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY class BaggagePropagatorTest extends DDCoreSpecification { BaggagePropagator propagator CarrierSetter setter - Object carrier + Map carrier Context context def setup() { this.propagator = new BaggagePropagator(true, true) setter = Mock(CarrierSetter) - carrier = new Object() + carrier = new HashMap<>() context = Context.root() } - def 'test tracing propagator context injection'() { + def 'test baggage propagator context injection'() { setup: context = BaggageContext.create(baggageMap).storeInto(context) @@ -46,6 +46,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { def "test baggage item limit"() { 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) when: @@ -63,6 +64,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { def "test baggage bytes limit"() { 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) when: @@ -114,5 +116,8 @@ class BaggagePropagatorTest extends DDCoreSpecification { "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | [] "=no-key" | [] "no-value=" | [] + "" | [] + ",," | [] + "=" | [] } } diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index df3b4d9f371..171b7ab8092 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -2256,6 +2256,18 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public boolean isBaggageExtract() { + return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); + } + + public boolean isBaggageInject() { + return tracePropagationStylesToInject.contains(TracePropagationStyle.BAGGAGE); + } + + public boolean isBaggagePropagationEnabled() { + return isBaggageInject() || isBaggageExtract(); + } + public int getTraceBaggageMaxItems() { return traceBaggageMaxItems; } From 188e82b0708efde2dfd8f109bbba6ff6ecd2a8f5 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 13 Feb 2025 15:10:13 -0800 Subject: [PATCH 17/31] updating tests --- .../trace/core/baggage/BaggagePropagator.java | 5 +- .../trace/core/propagation/HttpCodec.java | 2 +- .../core/baggage/BaggagePropagatorTest.groovy | 51 ++++++++++++++----- 3 files changed, 42 insertions(+), 16 deletions(-) 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 5a430602907..370a71f0b4c 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 @@ -167,12 +167,15 @@ private Map parseBaggageHeaders(String input) { int kvSeparatorInd = input.indexOf(keyValueSeparator); while (kvSeparatorInd != -1) { int end = pairSeparatorInd; - if (kvSeparatorInd > end) { // value is missing + if (kvSeparatorInd > end) { + log.debug( + "Dropping baggage headers due to key with no value {}", input.substring(start, end)); return Collections.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(); } baggage.put(key, value); diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 120d14f7ee5..925abafb439 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -261,7 +261,7 @@ else if (extracted != null && partialContext == null) { log.debug("Extract complete context {}", context); return context; } else if (partialContext != null) { - log.debug("Extract incomplete context {}", partialContext); + log.debug("Extract incomplete context f", partialContext); return partialContext; } else { log.debug("Extract no context"); 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 fcd701f0752..e57f68fd74e 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,9 +3,13 @@ package datadog.trace.core.baggage import datadog.context.Context import datadog.context.EmptyContext import datadog.context.propagation.CarrierSetter +import datadog.context.propagation.CarrierVisitor import datadog.trace.bootstrap.instrumentation.api.BaggageContext import datadog.trace.bootstrap.instrumentation.api.ContextVisitors import datadog.trace.core.test.DDCoreSpecification + +import java.util.function.BiConsumer + import static datadog.trace.core.baggage.BaggagePropagator.BAGGAGE_KEY class BaggagePropagatorTest extends DDCoreSpecification { @@ -14,10 +18,26 @@ class BaggagePropagatorTest extends DDCoreSpecification { Map carrier Context context + + static class MapCarrierAccessor + implements CarrierSetter>, CarrierVisitor> { + @Override + void set(Map carrier, String key, String value) { + if (carrier != null && key != null && value != null) { + carrier.put(key, value); + } + } + + @Override + void forEachKeyValue(Map carrier, BiConsumer visitor) { + carrier.forEach(visitor); + } + } + def setup() { this.propagator = new BaggagePropagator(true, true) - setter = Mock(CarrierSetter) - carrier = new HashMap<>() + setter = new MapCarrierAccessor() + carrier = [:] context = Context.root() } @@ -29,7 +49,8 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: - 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) +// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) + assert carrier[BAGGAGE_KEY] == baggageHeader where: baggageMap | baggageHeader @@ -53,7 +74,8 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: - 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) +// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) + assert carrier[BAGGAGE_KEY] == baggageHeader where: baggage | baggageHeader @@ -71,7 +93,8 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: - 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) +// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) + assert carrier[BAGGAGE_KEY] == baggageHeader where: baggage | baggageHeader @@ -108,16 +131,16 @@ class BaggagePropagatorTest extends DDCoreSpecification { context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - context instanceof EmptyContext + BaggageContext.fromContext(context) == null where: - baggageHeader | baggageMap - "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | [] - "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | [] - "=no-key" | [] - "no-value=" | [] - "" | [] - ",," | [] - "=" | [] + baggageHeader | _ + "no-equal-sign,foo=gets-dropped-because-previous-pair-is-malformed" | _ + "foo=gets-dropped-because-subsequent-pair-is-malformed,=" | _ + "=no-key" | _ + "no-value=" | _ + "" | _ + ",," | _ + "=" | _ } } From 58234e48d4252e4f1ae6501599eecad946791c4d Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 18 Feb 2025 16:10:41 -0800 Subject: [PATCH 18/31] adding baggagecache --- .../trace/core/baggage/BaggagePropagator.java | 48 +++++++++++++++++++ .../core/baggage/BaggagePropagatorTest.groovy | 33 ++++++++++--- 2 files changed, 75 insertions(+), 6 deletions(-) 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 370a71f0b4c..2cde7def8fb 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 @@ -9,10 +9,12 @@ import java.io.UnsupportedEncodingException; import java.net.URLDecoder; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.function.BiConsumer; @@ -27,6 +29,7 @@ public class BaggagePropagator implements Propagator { private Config config; private final boolean injectBaggage; private final boolean extractBaggage; + private BaggageCache baggageCache; private static final Set UNSAFE_CHARACTERS_KEY = new HashSet<>( Arrays.asList( @@ -39,6 +42,7 @@ public BaggagePropagator(Config config) { this.injectBaggage = config.isBaggageInject(); this.extractBaggage = config.isBaggageExtract(); this.config = config; + baggageCache = null; } // use primarily for testing purposes @@ -46,6 +50,7 @@ public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) { this.injectBaggage = injectBaggage; this.extractBaggage = extractBaggage; config = Config.get(); + baggageCache = null; } private int encodeKey(String key, StringBuilder builder) { @@ -88,6 +93,10 @@ public void inject(Context context, C carrier, CarrierSetter setter) { return; } + if (baggageCache != null) { + setter.set(carrier, BAGGAGE_KEY, baggageCache.getBaggage()); + } + int maxItems = config.getTraceBaggageMaxItems(); int maxBytes = config.getTraceBaggageMaxBytes(); if (!this.injectBaggage || maxItems == 0 || maxBytes == 0) { @@ -132,6 +141,7 @@ public Context extract(Context context, C carrier, CarrierVisitor visitor return context; } BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor(); + baggageCache = new BaggageCache<>(carrier, visitor); visitor.forEachKeyValue(carrier, baggageContextExtractor); BaggageContext extractedContext = baggageContextExtractor.extractedContext; if (extractedContext == null) { @@ -199,4 +209,42 @@ public void accept(String key, String value) { } } } + + private static class BaggageCache + implements BiConsumer, CarrierVisitor> { + /** Cached context key-values (even indexes are header names, odd indexes are header values). */ + private final List keysAndValues; + + private String baggage; + + public BaggageCache(C carrier, CarrierVisitor getter) { + this.keysAndValues = new ArrayList<>(32); + getter.forEachKeyValue(carrier, this); + } + + @Override + public void accept(String key, String value) { + keysAndValues.add(key); + keysAndValues.add(value); + if (BAGGAGE_KEY.equalsIgnoreCase(key)) { + try { + // Parse numeric header value to format it as 16 hexadecimal character format + this.baggage = value; + } catch (NumberFormatException ignored) { + } + } + } + + private String getBaggage() { + return this.baggage; + } + + @Override + public void forEachKeyValue(BaggageCache carrier, BiConsumer visitor) { + List keysAndValues = carrier.keysAndValues; + for (int i = 0; i < keysAndValues.size(); i += 2) { + visitor.accept(keysAndValues.get(i), keysAndValues.get(i + 1)); + } + } + } } 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 e57f68fd74e..813d23c6207 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 @@ -20,17 +20,17 @@ class BaggagePropagatorTest extends DDCoreSpecification { static class MapCarrierAccessor - implements CarrierSetter>, CarrierVisitor> { + implements CarrierSetter>, CarrierVisitor> { @Override void set(Map carrier, String key, String value) { if (carrier != null && key != null && value != null) { - carrier.put(key, value); + carrier.put(key, value) } } @Override void forEachKeyValue(Map carrier, BiConsumer visitor) { - carrier.forEach(visitor); + carrier.forEach(visitor) } } @@ -49,7 +49,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: -// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) assert carrier[BAGGAGE_KEY] == baggageHeader where: @@ -74,7 +73,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: -// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) assert carrier[BAGGAGE_KEY] == baggageHeader where: @@ -93,7 +91,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { this.propagator.inject(context, carrier, setter) then: -// 1 * setter.set(carrier, BAGGAGE_KEY, baggageHeader) assert carrier[BAGGAGE_KEY] == baggageHeader where: @@ -143,4 +140,28 @@ class BaggagePropagatorTest extends DDCoreSpecification { ",," | _ "=" | _ } + + def 'test baggage cache'() { + setup: + def headers = [ + (BAGGAGE_KEY): baggageHeader, + ] + + when: + context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) + + then: + BaggageContext.fromContext(context).getBaggage() == baggageMap + + when: + this.propagator.inject(context, carrier, setter) + + then: + 1*_ + + where: + baggageHeader | baggageMap + "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] + "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] + } } From 660529d3553486e0d1d96bdab12ddf6f0871725a Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Wed, 19 Feb 2025 15:29:40 -0800 Subject: [PATCH 19/31] updating baggagecache --- .../trace/core/baggage/BaggagePropagator.java | 56 +++---------------- .../core/baggage/BaggagePropagatorTest.groovy | 50 +++++++++-------- .../instrumentation/api/BaggageContext.java | 39 +++++++++++-- 3 files changed, 68 insertions(+), 77 deletions(-) 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 2cde7def8fb..766dd0f5e80 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 @@ -9,12 +9,10 @@ import java.io.UnsupportedEncodingException; import java.net.URLDecoder; import java.nio.charset.StandardCharsets; -import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; -import java.util.List; import java.util.Map; import java.util.Set; import java.util.function.BiConsumer; @@ -29,7 +27,6 @@ public class BaggagePropagator implements Propagator { private Config config; private final boolean injectBaggage; private final boolean extractBaggage; - private BaggageCache baggageCache; private static final Set UNSAFE_CHARACTERS_KEY = new HashSet<>( Arrays.asList( @@ -42,7 +39,6 @@ public BaggagePropagator(Config config) { this.injectBaggage = config.isBaggageInject(); this.extractBaggage = config.isBaggageExtract(); this.config = config; - baggageCache = null; } // use primarily for testing purposes @@ -50,7 +46,6 @@ public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) { this.injectBaggage = injectBaggage; this.extractBaggage = extractBaggage; config = Config.get(); - baggageCache = null; } private int encodeKey(String key, StringBuilder builder) { @@ -93,8 +88,9 @@ public void inject(Context context, C carrier, CarrierSetter setter) { return; } - if (baggageCache != null) { - setter.set(carrier, BAGGAGE_KEY, baggageCache.getBaggage()); + if (baggageContext.isUpdatedCache()) { + setter.set(carrier, BAGGAGE_KEY, baggageContext.getBaggageString()); + return; } int maxItems = config.getTraceBaggageMaxItems(); @@ -131,7 +127,10 @@ public void inject(Context context, C carrier, CarrierSetter setter) { currentBytes += byteSize; } - setter.set(carrier, BAGGAGE_KEY, baggageText.toString()); + String baggageString = baggageText.toString(); + ; + baggageContext.setBaggageString(baggageString); + setter.set(carrier, BAGGAGE_KEY, baggageString); } @Override @@ -141,7 +140,6 @@ public Context extract(Context context, C carrier, CarrierVisitor visitor return context; } BaggageContextExtractor baggageContextExtractor = new BaggageContextExtractor(); - baggageCache = new BaggageCache<>(carrier, visitor); visitor.forEachKeyValue(carrier, baggageContextExtractor); BaggageContext extractedContext = baggageContextExtractor.extractedContext; if (extractedContext == null) { @@ -204,47 +202,9 @@ public void accept(String key, String value) { && key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { - extractedContext = BaggageContext.create(baggage); - } - } - } - } - - private static class BaggageCache - implements BiConsumer, CarrierVisitor> { - /** Cached context key-values (even indexes are header names, odd indexes are header values). */ - private final List keysAndValues; - - private String baggage; - - public BaggageCache(C carrier, CarrierVisitor getter) { - this.keysAndValues = new ArrayList<>(32); - getter.forEachKeyValue(carrier, this); - } - - @Override - public void accept(String key, String value) { - keysAndValues.add(key); - keysAndValues.add(value); - if (BAGGAGE_KEY.equalsIgnoreCase(key)) { - try { - // Parse numeric header value to format it as 16 hexadecimal character format - this.baggage = value; - } catch (NumberFormatException ignored) { + extractedContext = BaggageContext.create(baggage, value); } } } - - private String getBaggage() { - return this.baggage; - } - - @Override - public void forEachKeyValue(BaggageCache carrier, BiConsumer visitor) { - List keysAndValues = carrier.keysAndValues; - for (int i = 0; i < keysAndValues.size(); i += 2) { - visitor.accept(keysAndValues.get(i), keysAndValues.get(i + 1)); - } - } } } 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 813d23c6207..4e583c4d9f9 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 @@ -1,7 +1,6 @@ package datadog.trace.core.baggage import datadog.context.Context -import datadog.context.EmptyContext import datadog.context.propagation.CarrierSetter import datadog.context.propagation.CarrierVisitor import datadog.trace.bootstrap.instrumentation.api.BaggageContext @@ -43,7 +42,8 @@ class BaggagePropagatorTest extends DDCoreSpecification { def 'test baggage propagator context injection'() { setup: - context = BaggageContext.create(baggageMap).storeInto(context) + context = BaggageContext.create(baggageMap, "").storeInto(context) + BaggageContext.fromContext(context).addBaggage("cache", "off") //force the cache to be not up to date when: this.propagator.inject(context, carrier, setter) @@ -53,21 +53,22 @@ class BaggagePropagatorTest extends DDCoreSpecification { where: baggageMap | baggageHeader - ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y" - ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" - [key1: "val1"] | "key1=val1" - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [serverNode: "DF 28"] | "serverNode=DF%2028" - [userId: "Amélie"] | "userId=Am%C3%A9lie" - ["user!d(me)": "false"] | "user!d%28me%29=false" - ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" + ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y,cache=off" + ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C,cache=off" + [key1: "val1"] | "key1=val1,cache=off" + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2,cache=off" + [serverNode: "DF 28"] | "serverNode=DF%2028,cache=off" + [userId: "Amélie"] | "userId=Am%C3%A9lie,cache=off" + ["user!d(me)": "false"] | "user!d%28me%29=false,cache=off" + ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5,cache=off" } def "test baggage item limit"() { 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 = BaggageContext.create(baggage, "").storeInto(context) + BaggageContext.fromContext(context).addBaggage("key2", "val2") //force the cache to be not up to date when: this.propagator.inject(context, carrier, setter) @@ -76,16 +77,17 @@ class BaggagePropagatorTest extends DDCoreSpecification { assert carrier[BAGGAGE_KEY] == baggageHeader where: - baggage | baggageHeader - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + baggage | baggageHeader + [key1: "val1"] | "key1=val1,key2=val2" + [key1: "val1", key3: "val3"] | "key1=val1,key3=val3" } def "test baggage bytes limit"() { 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 = BaggageContext.create(baggage, "").storeInto(context) + BaggageContext.fromContext(context).addBaggage("key2", "val2") //force the cache to be not up to date when: this.propagator.inject(context, carrier, setter) @@ -94,10 +96,10 @@ class BaggagePropagatorTest extends DDCoreSpecification { assert carrier[BAGGAGE_KEY] == baggageHeader where: - baggage | baggageHeader - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" - ["abcdefg": "hijklmnopq♥"] | "" + baggage | baggageHeader + [key1: "val1"] | "key1=val1,key2=val2" + [key1: "val1", key3: "val3"] | "key1=val1,key3=val3" + ["abcdefg": "hijklmnopq♥"] | "" } def 'test tracing propagator context extractor'() { @@ -141,23 +143,25 @@ class BaggagePropagatorTest extends DDCoreSpecification { "=" | _ } - def 'test baggage cache'() { + def "testing baggage cache"(){ setup: def headers = [ - (BAGGAGE_KEY): baggageHeader, + (BAGGAGE_KEY) : baggageHeader, ] when: context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - BaggageContext.fromContext(context).getBaggage() == baggageMap + BaggageContext baggageContext = BaggageContext.fromContext(context) + baggageContext.getBaggage() == baggageMap + baggageContext.isUpdatedCache() //ensure that the cache is filled when: this.propagator.inject(context, carrier, setter) then: - 1*_ + assert carrier[BAGGAGE_KEY] == baggageHeader where: baggageHeader | baggageMap 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 index 00e44365628..ecfc633a314 100644 --- 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 @@ -8,24 +8,51 @@ public class BaggageContext implements ImplicitContextKeyed { private static final ContextKey CONTEXT_KEY = ContextKey.named("baggage-key"); - private Map baggage; + private final Map baggage; + private String baggageString; + private boolean updatedCache; - public static BaggageContext create(Map baggage) { - return new BaggageContext(baggage); + public static BaggageContext create(Map baggage, String baggageString) { + return new BaggageContext(baggage, baggageString); } - private BaggageContext(Map baggage) { + private BaggageContext(Map baggage, String baggageString) { this.baggage = baggage; + this.baggageString = baggageString; + updatedCache = true; } - public static BaggageContext fromContext(Context context) { - return context.get(CONTEXT_KEY); + public void addBaggage(String key, String value) { + baggage.put(key, value); + updatedCache = false; + } + + public void removeBaggage(String key) { + baggage.remove(key); + updatedCache = false; + } + + public void setBaggageString(String baggageString) { + this.baggageString = baggageString; + updatedCache = true; + } + + public boolean isUpdatedCache() { + return updatedCache; + } + + public String getBaggageString() { + return baggageString; } public Map getBaggage() { return 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 a09abe16c5a982649233315525335f39fc4a1b0c Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Wed, 19 Feb 2025 15:52:03 -0800 Subject: [PATCH 20/31] cleanup --- .../src/main/java/datadog/trace/core/propagation/HttpCodec.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 925abafb439..120d14f7ee5 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -261,7 +261,7 @@ else if (extracted != null && partialContext == null) { log.debug("Extract complete context {}", context); return context; } else if (partialContext != null) { - log.debug("Extract incomplete context f", partialContext); + log.debug("Extract incomplete context {}", partialContext); return partialContext; } else { log.debug("Extract no context"); From a3dda56730ee6690281e875113ac1e5930e47f85 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 24 Feb 2025 14:48:16 -0500 Subject: [PATCH 21/31] updating pr comments --- .../java/datadog/trace/core/CoreTracer.java | 2 +- .../trace/core/baggage/BaggagePropagator.java | 32 +++++++++------ .../core/baggage/BaggagePropagatorTest.groovy | 40 +++++++++---------- .../instrumentation/api/BaggageContext.java | 27 ++++++++----- 4 files changed, 55 insertions(+), 46 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index c44f7a155b0..a14d0a56ed8 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -5,9 +5,9 @@ import static datadog.trace.api.DDTags.DJM_ENABLED; import static datadog.trace.api.DDTags.DSM_ENABLED; import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE; +import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.BAGGAGE_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.STANDALONE_ASM_CONCERN; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.BAGGAGE_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN; import static datadog.trace.common.metrics.MetricsAggregatorFactory.createMetricsAggregator; 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 766dd0f5e80..16c3d748d35 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 @@ -33,7 +33,7 @@ public class BaggagePropagator implements Propagator { '"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', '}')); private static final Set UNSAFE_CHARACTERS_VALUE = - new HashSet<>(Arrays.asList('\"', ',', ';', '\\')); + new HashSet<>(Arrays.asList('"', ',', ';', '\\')); public BaggagePropagator(Config config) { this.injectBaggage = config.isBaggageInject(); @@ -60,7 +60,7 @@ private int encode(String input, StringBuilder builder, Set UNSAFE_CH int size = 0; for (int i = 0; i < input.length(); i++) { char c = input.charAt(i); - if (UNSAFE_CHARACTERS.contains(c) || c > '~' || c <= ' ') { // encode character + if (c > '~' || c <= ' ' || UNSAFE_CHARACTERS.contains(c)) { // encode character byte[] bytes = Character.toString(c).getBytes(StandardCharsets.UTF_8); for (byte b : bytes) { builder.append('%'); @@ -82,20 +82,27 @@ private char encodeChar(int ascii) { @Override public void inject(Context context, C carrier, CarrierSetter setter) { - BaggageContext baggageContext = BaggageContext.fromContext(context); - if (baggageContext == null) { - log.debug("BaggageContext instance is missing from the following context {}", context); + int maxItems = config.getTraceBaggageMaxItems(); + int maxBytes = config.getTraceBaggageMaxBytes(); + //noinspection ConstantValue + if (!this.injectBaggage + || maxItems == 0 + || maxBytes == 0 + || context == null + || carrier == null + || setter == null) { return; } - if (baggageContext.isUpdatedCache()) { - setter.set(carrier, BAGGAGE_KEY, baggageContext.getBaggageString()); + BaggageContext baggageContext = BaggageContext.fromContext(context); + if (baggageContext == null) { + log.debug("BaggageContext instance is missing from the following context {}", context); return; } - int maxItems = config.getTraceBaggageMaxItems(); - int maxBytes = config.getTraceBaggageMaxBytes(); - if (!this.injectBaggage || maxItems == 0 || maxBytes == 0) { + String baggageHeader = baggageContext.getW3cBaggageHeader(); + if (baggageHeader != null) { + setter.set(carrier, BAGGAGE_KEY, baggageHeader); return; } @@ -128,8 +135,7 @@ public void inject(Context context, C carrier, CarrierSetter setter) { } String baggageString = baggageText.toString(); - ; - baggageContext.setBaggageString(baggageString); + baggageContext.setW3cBaggageHeader(baggageString); setter.set(carrier, BAGGAGE_KEY, baggageString); } @@ -202,7 +208,7 @@ public void accept(String key, String value) { && key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { - extractedContext = BaggageContext.create(baggage, value); + extractedContext = BaggageContext.createW3CBaggageContext(baggage, value); } } } 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 4e583c4d9f9..e980a68c821 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 @@ -42,8 +42,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { def 'test baggage propagator context injection'() { setup: - context = BaggageContext.create(baggageMap, "").storeInto(context) - BaggageContext.fromContext(context).addBaggage("cache", "off") //force the cache to be not up to date + context = new BaggageContext(baggageMap).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -53,22 +52,21 @@ class BaggagePropagatorTest extends DDCoreSpecification { where: baggageMap | baggageHeader - ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y,cache=off" - ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C,cache=off" - [key1: "val1"] | "key1=val1,cache=off" - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2,cache=off" - [serverNode: "DF 28"] | "serverNode=DF%2028,cache=off" - [userId: "Amélie"] | "userId=Am%C3%A9lie,cache=off" - ["user!d(me)": "false"] | "user!d%28me%29=false,cache=off" - ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5,cache=off" + ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y" + ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" + [key1: "val1"] | "key1=val1" + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [serverNode: "DF 28"] | "serverNode=DF%2028" + [userId: "Amélie"] | "userId=Am%C3%A9lie" + ["user!d(me)": "false"] | "user!d%28me%29=false" + ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" } def "test baggage item limit"() { 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) - BaggageContext.fromContext(context).addBaggage("key2", "val2") //force the cache to be not up to date + context = new BaggageContext(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -77,17 +75,16 @@ class BaggagePropagatorTest extends DDCoreSpecification { assert carrier[BAGGAGE_KEY] == baggageHeader where: - baggage | baggageHeader - [key1: "val1"] | "key1=val1,key2=val2" - [key1: "val1", key3: "val3"] | "key1=val1,key3=val3" + baggage | baggageHeader + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" } def "test baggage bytes limit"() { 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) - BaggageContext.fromContext(context).addBaggage("key2", "val2") //force the cache to be not up to date + context = new BaggageContext(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -96,10 +93,10 @@ class BaggagePropagatorTest extends DDCoreSpecification { assert carrier[BAGGAGE_KEY] == baggageHeader where: - baggage | baggageHeader - [key1: "val1"] | "key1=val1,key2=val2" - [key1: "val1", key3: "val3"] | "key1=val1,key3=val3" - ["abcdefg": "hijklmnopq♥"] | "" + baggage | baggageHeader + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [key1: "val1", key2: "val2", key3: "val3"] | "key1=val1,key2=val2" + ["abcdefg": "hijklmnopq♥"] | "" } def 'test tracing propagator context extractor'() { @@ -155,7 +152,6 @@ class BaggagePropagatorTest extends DDCoreSpecification { then: BaggageContext baggageContext = BaggageContext.fromContext(context) baggageContext.getBaggage() == baggageMap - baggageContext.isUpdatedCache() //ensure that the cache is filled when: this.propagator.inject(context, carrier, setter) 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 index ecfc633a314..77c51f65ebc 100644 --- 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 @@ -12,7 +12,15 @@ public class BaggageContext implements ImplicitContextKeyed { private String baggageString; private boolean updatedCache; - public static BaggageContext create(Map baggage, String baggageString) { + // Exclusively used for testing purposes + public BaggageContext(Map baggage) { + this.baggage = baggage; + this.baggageString = ""; + updatedCache = false; + } + + public static BaggageContext createW3CBaggageContext( + Map baggage, String baggageString) { return new BaggageContext(baggage, baggageString); } @@ -22,27 +30,26 @@ private BaggageContext(Map baggage, String baggageString) { updatedCache = true; } - public void addBaggage(String key, String value) { + public void addW3CBaggage(String key, String value) { baggage.put(key, value); updatedCache = false; } - public void removeBaggage(String key) { + public void removeW3CBaggage(String key) { baggage.remove(key); updatedCache = false; } - public void setBaggageString(String baggageString) { + public void setW3cBaggageHeader(String baggageString) { this.baggageString = baggageString; updatedCache = true; } - public boolean isUpdatedCache() { - return updatedCache; - } - - public String getBaggageString() { - return baggageString; + public String getW3cBaggageHeader() { + if (updatedCache) { + return baggageString; + } + return null; } public Map getBaggage() { From bfe69fc1fa1027c4ad490bf8fe5eba179c7bc71e Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 25 Feb 2025 14:59:54 -0500 Subject: [PATCH 22/31] updating PR comments and fixing CI --- .../trace/core/baggage/BaggagePropagator.java | 4 +-- .../core/baggage/BaggagePropagatorTest.groovy | 32 ++++++++--------- .../instrumentation/api/BaggageContext.java | 25 ++++++++----- .../datadog/trace/api/ConfigTest.groovy | 35 ++++++++++--------- 4 files changed, 52 insertions(+), 44 deletions(-) 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 16c3d748d35..fc7a2ec14fc 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 @@ -109,7 +109,7 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int processedBaggage = 0; int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); - for (final Map.Entry entry : baggageContext.getBaggage().entrySet()) { + for (final Map.Entry entry : baggageContext.asMap().entrySet()) { int byteSize = 1; // default include size of '=' if (processedBaggage != 0) { // if there are already baggage items processed, add and allocate bytes for a @@ -208,7 +208,7 @@ public void accept(String key, String value) { && key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { - extractedContext = BaggageContext.createW3CBaggageContext(baggage, value); + extractedContext = BaggageContext.create(baggage, value); } } } 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 e980a68c821..56a16cc376f 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 @@ -42,7 +42,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { def 'test baggage propagator context injection'() { setup: - context = new BaggageContext(baggageMap).storeInto(context) + context = BaggageContext.create(baggageMap).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -51,22 +51,22 @@ class BaggagePropagatorTest extends DDCoreSpecification { assert carrier[BAGGAGE_KEY] == baggageHeader where: - baggageMap | baggageHeader - ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] | "key1=val1,key2=val2,foo=bar,x=y" - ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" - [key1: "val1"] | "key1=val1" - [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" - [serverNode: "DF 28"] | "serverNode=DF%2028" - [userId: "Amélie"] | "userId=Am%C3%A9lie" - ["user!d(me)": "false"] | "user!d%28me%29=false" - ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" + baggageMap | baggageHeader + ["key1": "val1", "key2": "val2", "foo": "bar"] | "key1=val1,key2=val2,foo=bar" + ['",;\\()/:<=>?@[]{}': '",;\\'] | "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" + [key1: "val1"] | "key1=val1" + [key1: "val1", key2: "val2"] | "key1=val1,key2=val2" + [serverNode: "DF 28"] | "serverNode=DF%2028" + [userId: "Amélie"] | "userId=Am%C3%A9lie" + ["user!d(me)": "false"] | "user!d%28me%29=false" + ["abcdefg": "hijklmnopq♥"] | "abcdefg=hijklmnopq%E2%99%A5" } def "test baggage item limit"() { setup: injectSysConfig("trace.baggage.max.items", '2') propagator = new BaggagePropagator(true, true) //creating a new instance after injecting config - context = new BaggageContext(baggage).storeInto(context) + context = BaggageContext.create(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -84,7 +84,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 = new BaggageContext(baggage).storeInto(context) + context = BaggageContext.create(baggage).storeInto(context) when: this.propagator.inject(context, carrier, setter) @@ -109,11 +109,11 @@ class BaggagePropagatorTest extends DDCoreSpecification { context = this.propagator.extract(context, headers, ContextVisitors.stringValuesMap()) then: - BaggageContext.fromContext(context).getBaggage() == baggageMap + BaggageContext.fromContext(context).asMap() == baggageMap where: baggageHeader | baggageMap - "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] + "key1=val1,key2=val2,foo=bar" | ["key1": "val1", "key2": "val2", "foo": "bar"] "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] } @@ -151,7 +151,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { then: BaggageContext baggageContext = BaggageContext.fromContext(context) - baggageContext.getBaggage() == baggageMap + baggageContext.asMap() == baggageMap when: this.propagator.inject(context, carrier, setter) @@ -161,7 +161,7 @@ class BaggagePropagatorTest extends DDCoreSpecification { where: baggageHeader | baggageMap - "key1=val1,key2=val2,foo=bar,x=y" | ["key1": "val1", "key2": "val2", "foo": "bar", "x": "y"] + "key1=val1,key2=val2,foo=bar" | ["key1": "val1", "key2": "val2", "foo": "bar"] "%22%2C%3B%5C%28%29%2F%3A%3C%3D%3E%3F%40%5B%5D%7B%7D=%22%2C%3B%5C" | ['",;\\()/:<=>?@[]{}': '",;\\'] } } 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 index 77c51f65ebc..99048f7a107 100644 --- 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 @@ -3,6 +3,7 @@ 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 { @@ -12,16 +13,22 @@ public class BaggageContext implements ImplicitContextKeyed { private String baggageString; private boolean updatedCache; - // Exclusively used for testing purposes - public BaggageContext(Map baggage) { + 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 createW3CBaggageContext( - Map baggage, String baggageString) { - return new BaggageContext(baggage, baggageString); + public static BaggageContext create(Map baggage, String w3cHeader) { + return new BaggageContext(baggage, w3cHeader); } private BaggageContext(Map baggage, String baggageString) { @@ -40,8 +47,8 @@ public void removeW3CBaggage(String key) { updatedCache = false; } - public void setW3cBaggageHeader(String baggageString) { - this.baggageString = baggageString; + public void setW3cBaggageHeader(String w3cHeader) { + this.baggageString = w3cHeader; updatedCache = true; } @@ -52,8 +59,8 @@ public String getW3cBaggageHeader() { return null; } - public Map getBaggage() { - return baggage; + public Map asMap() { + return new HashMap<>(baggage); } public static BaggageContext fromContext(Context context) { diff --git a/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy index b3e2a3d9b09..df5509b13ef 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy +++ b/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy @@ -26,6 +26,7 @@ import static datadog.trace.api.TracePropagationStyle.B3SINGLE import static datadog.trace.api.TracePropagationStyle.DATADOG import static datadog.trace.api.TracePropagationStyle.HAYSTACK import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT +import static datadog.trace.api.TracePropagationStyle.BAGGAGE import static datadog.trace.api.config.CiVisibilityConfig.CIVISIBILITY_AGENTLESS_ENABLED import static datadog.trace.api.config.CiVisibilityConfig.CIVISIBILITY_ENABLED import static datadog.trace.api.config.DebuggerConfig.DYNAMIC_INSTRUMENTATION_CLASSFILE_DUMP_ENABLED @@ -645,8 +646,8 @@ class ConfigTest extends DDSpecification { config.splitByTags == [].toSet() config.propagationStylesToExtract.toList() == [PropagationStyle.DATADOG] config.propagationStylesToInject.toList() == [PropagationStyle.DATADOG] - config.tracePropagationStylesToExtract.toList() == [DATADOG, TRACECONTEXT] - config.tracePropagationStylesToInject.toList() == [DATADOG, TRACECONTEXT] + config.tracePropagationStylesToExtract.toList() == [DATADOG, TRACECONTEXT, BAGGAGE] + config.tracePropagationStylesToInject.toList() == [DATADOG, TRACECONTEXT, BAGGAGE] config.longRunningTraceEnabled == false } @@ -2403,21 +2404,21 @@ class ConfigTest extends DDSpecification { where: // spotless:off - pse | psi | tps | tpse | tpsi | ePSE | ePSI | eTPSE | eTPSI - PropagationStyle.DATADOG | PropagationStyle.B3 | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.B3] | [DATADOG] | [B3SINGLE, B3MULTI] - PropagationStyle.B3 | PropagationStyle.DATADOG | null | null | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE, B3MULTI] | [DATADOG] - PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | null | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [HAYSTACK] | [HAYSTACK] - PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | B3SINGLE | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [HAYSTACK] - PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | null | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [HAYSTACK] | [B3MULTI] - PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | B3SINGLE | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] - PropagationStyle.B3 | PropagationStyle.DATADOG | null | B3SINGLE | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] - null | null | HAYSTACK | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [HAYSTACK] | [HAYSTACK] - null | null | HAYSTACK | B3SINGLE | B3MULTI | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] - null | null | null | B3SINGLE | B3MULTI | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] - null | null | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [DATADOG, TRACECONTEXT] | [DATADOG, TRACECONTEXT] - null | null | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [DATADOG, TRACECONTEXT] | [DATADOG, TRACECONTEXT] - null | null | null | "b3 single header" | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [DATADOG, TRACECONTEXT] - null | null | null | "b3" | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3MULTI] | [DATADOG, TRACECONTEXT] + pse | psi | tps | tpse | tpsi | ePSE | ePSI | eTPSE | eTPSI + PropagationStyle.DATADOG | PropagationStyle.B3 | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.B3] | [DATADOG] | [B3SINGLE, B3MULTI] + PropagationStyle.B3 | PropagationStyle.DATADOG | null | null | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE, B3MULTI] | [DATADOG] + PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | null | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [HAYSTACK] | [HAYSTACK] + PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | B3SINGLE | null | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [HAYSTACK] + PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | null | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [HAYSTACK] | [B3MULTI] + PropagationStyle.B3 | PropagationStyle.DATADOG | HAYSTACK | B3SINGLE | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] + PropagationStyle.B3 | PropagationStyle.DATADOG | null | B3SINGLE | B3MULTI | [PropagationStyle.B3] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] + null | null | HAYSTACK | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [HAYSTACK] | [HAYSTACK] + null | null | HAYSTACK | B3SINGLE | B3MULTI | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] + null | null | null | B3SINGLE | B3MULTI | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [B3MULTI] + null | null | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [DATADOG, TRACECONTEXT, BAGGAGE] | [DATADOG, TRACECONTEXT, BAGGAGE] + null | null | null | null | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [DATADOG, TRACECONTEXT, BAGGAGE] | [DATADOG, TRACECONTEXT, BAGGAGE] + null | null | null | "b3 single header" | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3SINGLE] | [DATADOG, TRACECONTEXT, BAGGAGE] + null | null | null | "b3" | null | [PropagationStyle.DATADOG] | [PropagationStyle.DATADOG] | [B3MULTI] | [DATADOG, TRACECONTEXT, BAGGAGE] // spotless:on } From a0df83ae553161f5c5af6f3ec64b4d41ba75a752 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 27 Feb 2025 15:10:18 -0500 Subject: [PATCH 23/31] adding OTEL PercentEncoder and utilizing it in BaggagePropagator --- .../trace/core/baggage/BaggagePropagator.java | 74 +--- .../trace/core/util/PercentEscaper.java | 396 ++++++++++++++++++ 2 files changed, 414 insertions(+), 56 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.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 fc7a2ec14fc..c772c44130f 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 @@ -6,15 +6,13 @@ import datadog.context.propagation.Propagator; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.BaggageContext; +import datadog.trace.core.util.PercentEscaper; import java.io.UnsupportedEncodingException; import java.net.URLDecoder; -import java.nio.charset.StandardCharsets; -import java.util.Arrays; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.Map; -import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import javax.annotation.ParametersAreNonnullByDefault; import org.slf4j.Logger; @@ -23,17 +21,11 @@ @ParametersAreNonnullByDefault public class BaggagePropagator implements Propagator { private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class); + private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create(); static final String BAGGAGE_KEY = "baggage"; private Config config; private final boolean injectBaggage; private final boolean extractBaggage; - private static final Set UNSAFE_CHARACTERS_KEY = - new HashSet<>( - Arrays.asList( - '"', ',', ';', '\\', '(', ')', '/', ':', '<', '=', '>', '?', '@', '[', ']', '{', - '}')); - private static final Set UNSAFE_CHARACTERS_VALUE = - new HashSet<>(Arrays.asList('"', ',', ';', '\\')); public BaggagePropagator(Config config) { this.injectBaggage = config.isBaggageInject(); @@ -48,38 +40,6 @@ public BaggagePropagator(boolean injectBaggage, boolean extractBaggage) { config = Config.get(); } - private int encodeKey(String key, StringBuilder builder) { - return encode(key, builder, UNSAFE_CHARACTERS_KEY); - } - - private int encodeValue(String key, StringBuilder builder) { - return encode(key, builder, UNSAFE_CHARACTERS_VALUE); - } - - private int encode(String input, StringBuilder builder, Set UNSAFE_CHARACTERS) { - int size = 0; - for (int i = 0; i < input.length(); i++) { - char c = input.charAt(i); - if (c > '~' || c <= ' ' || UNSAFE_CHARACTERS.contains(c)) { // encode character - byte[] bytes = Character.toString(c).getBytes(StandardCharsets.UTF_8); - for (byte b : bytes) { - builder.append('%'); - builder.append(encodeChar((b >> 4) & 0xF)); - builder.append(encodeChar(b & 0xF)); - size++; - } - } else { - builder.append(c); - size++; - } - } - return size; - } - - private char encodeChar(int ascii) { - return (char) (ascii < 10 ? '0' + ascii : 'A' + (ascii - 10)); - } - @Override public void inject(Context context, C carrier, CarrierSetter setter) { int maxItems = config.getTraceBaggageMaxItems(); @@ -110,28 +70,30 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); for (final Map.Entry entry : baggageContext.asMap().entrySet()) { - int byteSize = 1; // default include size of '=' - if (processedBaggage - != 0) { // if there are already baggage items processed, add and allocate bytes for a - // comma + AtomicInteger pairSize = new AtomicInteger(1); + // if there are already baggage items processed, add and allocate bytes for a comma + if (processedBaggage != 0) { baggageText.append(','); - byteSize++; + pairSize.incrementAndGet(); } - byteSize += encodeKey(entry.getKey(), baggageText); + baggageText.append( + UTF_ESCAPER.escape(entry.getKey(), pairSize, UTF_ESCAPER.getUnsafeKeyOctets())); baggageText.append('='); - byteSize += encodeValue(entry.getValue(), baggageText); + baggageText.append( + UTF_ESCAPER.escape(entry.getValue(), pairSize, UTF_ESCAPER.getUnsafeValOctets())); processedBaggage++; - if (processedBaggage == maxItems) { // reached the max number of baggage items allowed + // reached the max number of baggage items allowed + if (processedBaggage == maxItems) { break; } - if (currentBytes + byteSize - > maxBytes) { // Drop newest k/v pair if adding it leads to exceeding the limit + // Drop newest k/v pair if adding it leads to exceeding the limit + if (currentBytes + pairSize.get() > maxBytes) { baggageText.setLength(currentBytes); break; } - currentBytes += byteSize; + currentBytes += pairSize.get(); } String baggageString = baggageText.toString(); @@ -204,8 +166,8 @@ private Map parseBaggageHeaders(String input) { @Override public void accept(String key, String value) { - if (key != null - && key.equalsIgnoreCase(BAGGAGE_KEY)) { // Only process tags that are relevant to baggage + // Only process tags that are relevant to baggage + if (key != null && key.equalsIgnoreCase(BAGGAGE_KEY)) { Map baggage = parseBaggageHeaders(value); if (!baggage.isEmpty()) { extractedContext = BaggageContext.create(baggage, value); 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 new file mode 100644 index 00000000000..0745ca72da1 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/util/PercentEscaper.java @@ -0,0 +1,396 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +// Includes work from: +/* + * Copyright (C) 2008 The Guava Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except + * in compliance with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under the License + * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express + * or implied. See the License for the specific language governing permissions and limitations under + * the License. + */ + +package datadog.trace.core.util; + +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.CheckForNull; + +/** + * Note: This class is based on code from guava. It is comprised of code from three classes: + * + * + * + *

Escapes some set of Java characters using a UTF-8 based percent encoding scheme. The set of + * safe characters (those which remain unescaped) can be specified on construction. + * + *

This class is primarily used for creating URI escapers in {@code UrlEscapers} but can be used + * directly if required. While URI escapers impose specific semantics on which characters are + * considered 'safe', this class has a minimal set of restrictions. + * + *

When escaping a String, the following rules apply: + * + *

    + *
  • All specified safe characters remain unchanged. + *
  • If {@code plusForSpace} was specified, the space character " " is converted into a plus + * sign {@code "+"}. + *
  • All other characters are converted into one or more bytes using UTF-8 encoding and each + * byte is then represented by the 3-character string "%XX", where "XX" is the two-digit, + * uppercase, hexadecimal representation of the byte value. + *
+ * + *

For performance reasons the only currently supported character encoding of this class is + * UTF-8. + * + *

Note: This escaper produces uppercase hexadecimal sequences. + * + *

This class is internal and is hence not for public use. Its APIs are unstable and can change + * at any time. + * + * @author David Beaumont + * @since 15.0 + */ +public final class PercentEscaper { + + /** The amount of padding (chars) to use when growing the escape buffer. */ + private static final int DEST_PAD = 32; + + private static final String SAFE_CHARS = + "-._~" // Unreserved characters. + + "!$'()*" + // + ",;=" // baggage delimiters, so let's escape them + + "&" // The subdelim characters. + + "@:" // The gendelim characters permitted in paths. + + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; + + private static final String UNSAFE_CHARACTERS_KEY = "\",;\\()/:<=>?@[]{} "; + private static final String UNSAFE_CHARACTERS_VALUE = "\",;\\ "; + + // Percent escapers output upper case hex digits (uri escapers require this). + private static final char[] UPPER_HEX_DIGITS = "0123456789ABCDEF".toCharArray(); + + /** + * An array of flags where for any {@code char c} if {@code safeOctets[c]} is true then {@code c} + * should remain unmodified in the output. If {@code c >= safeOctets.length} then it should be + * escaped. + */ + // private static final boolean[] safeOctets = createSafeOctets(SAFE_CHARS); + private static final boolean[] unsafeKeyOctets = createUnsafeOctets(UNSAFE_CHARACTERS_KEY); + + private static final boolean[] unsafeValOctets = createUnsafeOctets(UNSAFE_CHARACTERS_VALUE); + + /** The default {@link PercentEscaper} which will *not* replace spaces with plus signs. */ + public static PercentEscaper create() { + return new PercentEscaper(); + } + + /** + * Creates a boolean array with entries corresponding to the character values specified in + * safeChars set to true. The array is as small as is required to hold the given character + * information. + */ + private static boolean[] createUnsafeOctets(String safeChars) { + int maxChar = -1; + char[] safeCharArray = safeChars.toCharArray(); + for (char c : safeCharArray) { + maxChar = Math.max(c, maxChar); + } + boolean[] octets = new boolean[maxChar + 1]; + for (char c : safeCharArray) { + octets[c] = true; + } + return octets; + } + + public boolean[] getUnsafeKeyOctets() { + return unsafeKeyOctets; + } + + public boolean[] getUnsafeValOctets() { + return unsafeValOctets; + } + + /** Escape the provided String, using percent-style URL Encoding. */ + public String escape(String s, AtomicInteger size, 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, size, unsafeOctets); + } + size.incrementAndGet(); + } + return s; + } + + /* + * Overridden for performance. For unescaped strings this improved the performance of the uri + * escaper from ~760ns to ~400ns as measured by {@link CharEscapersBenchmark}. + */ + /** + * Returns the escaped form of a given literal string, starting at the given index. This method is + * called by the {@link #escape(String, boolean[])} method when it discovers that escaping is + * required. It is protected to allow subclasses to override the fastpath escaping function to + * inline their escaping test. + * + *

This method is not reentrant and may only be invoked by the top level {@link #escape(String, + * boolean[])} method. + * + * @param s the literal string to be escaped + * @param index the index to start escaping from + * @return the escaped form of {@code string} + * @throws NullPointerException if {@code string} is null + * @throws IllegalArgumentException if invalid surrogate characters are encountered + */ + private static String escapeSlow( + String s, int index, AtomicInteger size, 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; + + while (index < end) { + int cp = codePointAt(s, index, end); + if (cp < 0) { + throw new IllegalArgumentException("Trailing high surrogate at end of input"); + } + // 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, size, unsafeOctets); + int nextIndex = index + (Character.isSupplementaryCodePoint(cp) ? 2 : 1); + if (escaped != null) { + int charsSkipped = index - unescapedChunkStart; + + // This is the size needed to add the replacement, not the full + // size needed by the string. We only regrow when we absolutely must. + int sizeNeeded = destIndex + charsSkipped + escaped.length; + if (dest.length < sizeNeeded) { + int destLength = sizeNeeded + (end - index) + DEST_PAD; + dest = growBuffer(dest, destIndex, destLength); + } + // If we have skipped any characters, we need to copy them now. + if (charsSkipped > 0) { + s.getChars(unescapedChunkStart, index, dest, destIndex); + destIndex += charsSkipped; + } + if (escaped.length > 0) { + System.arraycopy(escaped, 0, dest, destIndex, escaped.length); + destIndex += escaped.length; + } + // If we dealt with an escaped character, reset the unescaped range. + unescapedChunkStart = nextIndex; + } + index = nextEscapeIndex(s, nextIndex, end, unsafeOctets); + } + + // Process trailing unescaped characters - no need to account for escaped + // length or padding the allocation. + int charsSkipped = end - unescapedChunkStart; + if (charsSkipped > 0) { + int endIndex = destIndex + charsSkipped; + if (dest.length < endIndex) { + dest = growBuffer(dest, destIndex, endIndex); + } + s.getChars(unescapedChunkStart, end, dest, destIndex); + destIndex = endIndex; + } + return new String(dest, 0, destIndex); + } + + private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean[] unsafeOctets) { + for (; index < end; index++) { + char c = csq.charAt(index); + if (c <= unsafeOctets.length && unsafeOctets[c]) { + break; + } + } + return index; + } + + /** Escapes the given Unicode code point in UTF-8. */ + @CheckForNull + @SuppressWarnings("UngroupedOverloads") + private static char[] escape(int cp, AtomicInteger size, 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]) { + return null; + } else if (cp <= 0x7F) { + // Single byte UTF-8 characters + // Start with "%--" and fill in the blanks + char[] dest = new char[3]; + dest[0] = '%'; + dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; + dest[1] = UPPER_HEX_DIGITS[cp >>> 4]; + size.incrementAndGet(); + return dest; + } else if (cp <= 0x7ff) { + // Two byte UTF-8 characters [cp >= 0x80 && cp <= 0x7ff] + // Start with "%--%--" and fill in the blanks + char[] dest = new char[6]; + dest[0] = '%'; + dest[3] = '%'; + dest[5] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[1] = UPPER_HEX_DIGITS[0xC | cp]; + size.addAndGet(2); + return dest; + } else if (cp <= 0xffff) { + // Three byte UTF-8 characters [cp >= 0x800 && cp <= 0xffff] + // Start with "%E-%--%--" and fill in the blanks + char[] dest = new char[9]; + dest[0] = '%'; + dest[1] = 'E'; + dest[3] = '%'; + dest[6] = '%'; + dest[8] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[7] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[5] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[2] = UPPER_HEX_DIGITS[cp]; + size.addAndGet(3); + return dest; + } else if (cp <= 0x10ffff) { + char[] dest = new char[12]; + // Four byte UTF-8 characters [cp >= 0xffff && cp <= 0x10ffff] + // Start with "%F-%--%--%--" and fill in the blanks + dest[0] = '%'; + dest[1] = 'F'; + dest[3] = '%'; + dest[6] = '%'; + dest[9] = '%'; + dest[11] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[10] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[8] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[7] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[5] = UPPER_HEX_DIGITS[cp & 0xF]; + cp >>>= 4; + dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; + cp >>>= 2; + dest[2] = UPPER_HEX_DIGITS[cp & 0x7]; + size.addAndGet(4); + return dest; + } else { + // If this ever happens it is due to bug in UnicodeEscaper, not bad input. + throw new IllegalArgumentException("Invalid unicode character value " + cp); + } + } + + /** + * Returns the Unicode code point of the character at the given index. + * + *

Unlike {@link Character#codePointAt(CharSequence, int)} or {@link String#codePointAt(int)} + * this method will never fail silently when encountering an invalid surrogate pair. + * + *

The behaviour of this method is as follows: + * + *

    + *
  1. If {@code index >= end}, {@link IndexOutOfBoundsException} is thrown. + *
  2. If the character at the specified index is not a surrogate, it is returned. + *
  3. If the first character was a high surrogate value, then an attempt is made to read the + * next character. + *
      + *
    1. If the end of the sequence was reached, the negated value of the trailing high + * surrogate is returned. + *
    2. If the next character was a valid low surrogate, the code point value of the + * high/low surrogate pair is returned. + *
    3. If the next character was not a low surrogate value, then {@link + * IllegalArgumentException} is thrown. + *
    + *
  4. If the first character was a low surrogate value, {@link IllegalArgumentException} is + * thrown. + *
+ * + * @param seq the sequence of characters from which to decode the code point + * @param index the index of the first character to decode + * @param end the index beyond the last valid character to decode + * @return the Unicode code point for the given index or the negated value of the trailing high + * surrogate character at the end of the sequence + */ + private static int codePointAt(CharSequence seq, int index, int end) { + if (index < end) { + char c1 = seq.charAt(index++); + if (c1 < Character.MIN_HIGH_SURROGATE || c1 > Character.MAX_LOW_SURROGATE) { + // Fast path (first test is probably all we need to do) + return c1; + } else if (c1 <= Character.MAX_HIGH_SURROGATE) { + // If the high surrogate was the last character, return its inverse + if (index == end) { + return -c1; + } + // Otherwise look for the low surrogate following it + char c2 = seq.charAt(index); + if (Character.isLowSurrogate(c2)) { + return Character.toCodePoint(c1, c2); + } + throw new IllegalArgumentException( + "Expected low surrogate but got char '" + + c2 + + "' with value " + + (int) c2 + + " at index " + + index + + " in '" + + seq + + "'"); + } else { + throw new IllegalArgumentException( + "Unexpected low surrogate character '" + + c1 + + "' with value " + + (int) c1 + + " at index " + + (index - 1) + + " in '" + + seq + + "'"); + } + } + throw new IndexOutOfBoundsException("Index exceeds specified range"); + } + + /** + * Helper method to grow the character buffer as needed, this only happens once in a while so it's + * ok if it's in a method call. If the index passed in is 0 then no copying will be done. + */ + private static char[] growBuffer(char[] dest, int index, int size) { + if (size < 0) { // overflow - should be OutOfMemoryError but GWT/j2cl don't support it + throw new AssertionError("Cannot increase internal buffer any further"); + } + char[] copy = new char[size]; + if (index > 0) { + System.arraycopy(dest, 0, copy, 0, index); + } + return copy; + } +} From 831b564475a9628714420f855bbc3f9c3c9549e0 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 27 Feb 2025 15:25:59 -0500 Subject: [PATCH 24/31] nit fixes --- .../datadog/trace/core/baggage/BaggagePropagator.java | 2 +- .../datadog/trace/core/propagation/HttpCodec.java | 4 ++++ .../java/datadog/trace/core/util/PercentEscaper.java | 11 +---------- 3 files changed, 6 insertions(+), 11 deletions(-) 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 c772c44130f..4d72731f1c8 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 @@ -23,7 +23,7 @@ public class BaggagePropagator implements Propagator { private static final Logger log = LoggerFactory.getLogger(BaggagePropagator.class); private static final PercentEscaper UTF_ESCAPER = PercentEscaper.create(); static final String BAGGAGE_KEY = "baggage"; - private Config config; + private final Config config; private final boolean injectBaggage; private final boolean extractBaggage; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 120d14f7ee5..d98c4d86d13 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -125,6 +125,8 @@ private static Map createInjectors( case TRACECONTEXT: result.put(style, W3CHttpCodec.newInjector(reverseBaggageMapping)); break; + case BAGGAGE: + break; default: log.debug("No implementation found to inject propagation style: {}", style); break; @@ -159,6 +161,8 @@ public static Extractor createExtractor( case TRACECONTEXT: extractors.add(W3CHttpCodec.newExtractor(config, traceConfigSupplier)); break; + case BAGGAGE: + break; default: log.debug("No implementation found to extract propagation style: {}", style); break; 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 0745ca72da1..1b81ca60aec 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 @@ -70,14 +70,6 @@ public final class PercentEscaper { /** The amount of padding (chars) to use when growing the escape buffer. */ private static final int DEST_PAD = 32; - private static final String SAFE_CHARS = - "-._~" // Unreserved characters. - + "!$'()*" - // + ",;=" // baggage delimiters, so let's escape them - + "&" // The subdelim characters. - + "@:" // The gendelim characters permitted in paths. - + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"; - private static final String UNSAFE_CHARACTERS_KEY = "\",;\\()/:<=>?@[]{} "; private static final String UNSAFE_CHARACTERS_VALUE = "\",;\\ "; @@ -85,11 +77,10 @@ public final class PercentEscaper { private static final char[] UPPER_HEX_DIGITS = "0123456789ABCDEF".toCharArray(); /** - * An array of flags where for any {@code char c} if {@code safeOctets[c]} is true then {@code c} + * Arrays of flags where for any {@code char c} if {@code safeOctets[c]} is true then {@code c} * should remain unmodified in the output. If {@code c >= safeOctets.length} then it should be * escaped. */ - // private static final boolean[] safeOctets = createSafeOctets(SAFE_CHARS); private static final boolean[] unsafeKeyOctets = createUnsafeOctets(UNSAFE_CHARACTERS_KEY); private static final boolean[] unsafeValOctets = createUnsafeOctets(UNSAFE_CHARACTERS_VALUE); From e0a97095eed001b466405ebddbe5018eb9b24eff Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 27 Feb 2025 15:37:05 -0500 Subject: [PATCH 25/31] hiding unsafeKey and unsafeValue --- .../datadog/trace/core/baggage/BaggagePropagator.java | 6 ++---- .../main/java/datadog/trace/core/util/PercentEscaper.java | 8 ++++---- 2 files changed, 6 insertions(+), 8 deletions(-) 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 4d72731f1c8..4ac3ec71d2d 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 @@ -77,11 +77,9 @@ public void inject(Context context, C carrier, CarrierSetter setter) { pairSize.incrementAndGet(); } - baggageText.append( - UTF_ESCAPER.escape(entry.getKey(), pairSize, UTF_ESCAPER.getUnsafeKeyOctets())); + baggageText.append(UTF_ESCAPER.escapeKey(entry.getKey(), pairSize)); baggageText.append('='); - baggageText.append( - UTF_ESCAPER.escape(entry.getValue(), pairSize, UTF_ESCAPER.getUnsafeValOctets())); + baggageText.append(UTF_ESCAPER.escapeValue(entry.getValue(), pairSize)); processedBaggage++; // reached the max number of baggage items allowed 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 1b81ca60aec..cfa71a60fb3 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 @@ -108,12 +108,12 @@ private static boolean[] createUnsafeOctets(String safeChars) { return octets; } - public boolean[] getUnsafeKeyOctets() { - return unsafeKeyOctets; + public String escapeKey(String s, AtomicInteger size) { + return escape(s, size, unsafeKeyOctets); } - public boolean[] getUnsafeValOctets() { - return unsafeValOctets; + public String escapeValue(String s, AtomicInteger size) { + return escape(s, size, unsafeKeyOctets); } /** Escape the provided String, using percent-style URL Encoding. */ From b2731e8c13e0211a5bb4a3cedde7ce13d3a89f65 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Thu, 27 Feb 2025 16:03:16 -0500 Subject: [PATCH 26/31] change from atomicinteger to int[] --- .../trace/core/baggage/BaggagePropagator.java | 9 ++++---- .../trace/core/util/PercentEscaper.java | 22 +++++++++---------- 2 files changed, 14 insertions(+), 17 deletions(-) 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 4ac3ec71d2d..79931336d03 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 @@ -12,7 +12,6 @@ import java.util.Collections; import java.util.HashMap; import java.util.Map; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BiConsumer; import javax.annotation.ParametersAreNonnullByDefault; import org.slf4j.Logger; @@ -70,11 +69,11 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); for (final Map.Entry entry : baggageContext.asMap().entrySet()) { - AtomicInteger pairSize = new AtomicInteger(1); + int[] pairSize = {1}; // if there are already baggage items processed, add and allocate bytes for a comma if (processedBaggage != 0) { baggageText.append(','); - pairSize.incrementAndGet(); + pairSize[0]++; } baggageText.append(UTF_ESCAPER.escapeKey(entry.getKey(), pairSize)); @@ -87,11 +86,11 @@ public void inject(Context context, C carrier, CarrierSetter setter) { break; } // Drop newest k/v pair if adding it leads to exceeding the limit - if (currentBytes + pairSize.get() > maxBytes) { + if (currentBytes + pairSize[0] > maxBytes) { baggageText.setLength(currentBytes); break; } - currentBytes += pairSize.get(); + currentBytes += pairSize[0]; } String baggageString = baggageText.toString(); 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 cfa71a60fb3..129d5255dc8 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 @@ -20,7 +20,6 @@ package datadog.trace.core.util; -import java.util.concurrent.atomic.AtomicInteger; import javax.annotation.CheckForNull; /** @@ -108,23 +107,23 @@ private static boolean[] createUnsafeOctets(String safeChars) { return octets; } - public String escapeKey(String s, AtomicInteger size) { + public String escapeKey(String s, int[] size) { return escape(s, size, unsafeKeyOctets); } - public String escapeValue(String s, AtomicInteger size) { + public String escapeValue(String s, int[] size) { return escape(s, size, unsafeKeyOctets); } /** Escape the provided String, using percent-style URL Encoding. */ - public String escape(String s, AtomicInteger size, boolean[] unsafeOctets) { + public String escape(String s, int[] size, 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, size, unsafeOctets); } - size.incrementAndGet(); + size[0]++; } return s; } @@ -148,8 +147,7 @@ public String escape(String s, AtomicInteger size, boolean[] unsafeOctets) { * @throws NullPointerException if {@code string} is null * @throws IllegalArgumentException if invalid surrogate characters are encountered */ - private static String escapeSlow( - String s, int index, AtomicInteger size, boolean[] unsafeOctets) { + private static String escapeSlow(String s, int index, int[] size, boolean[] unsafeOctets) { int end = s.length(); // Get a destination buffer and setup some loop variables. @@ -219,7 +217,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, AtomicInteger size, boolean[] unsafeOctets) { + private static char[] escape(int cp, int[] size, 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]) { @@ -231,7 +229,7 @@ private static char[] escape(int cp, AtomicInteger size, boolean[] unsafeOctets) dest[0] = '%'; dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; dest[1] = UPPER_HEX_DIGITS[cp >>> 4]; - size.incrementAndGet(); + size[0]++; return dest; } else if (cp <= 0x7ff) { // Two byte UTF-8 characters [cp >= 0x80 && cp <= 0x7ff] @@ -246,7 +244,7 @@ private static char[] escape(int cp, AtomicInteger size, boolean[] unsafeOctets) dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; cp >>>= 4; dest[1] = UPPER_HEX_DIGITS[0xC | cp]; - size.addAndGet(2); + size[0] += 2; return dest; } else if (cp <= 0xffff) { // Three byte UTF-8 characters [cp >= 0x800 && cp <= 0xffff] @@ -265,7 +263,7 @@ private static char[] escape(int cp, AtomicInteger size, boolean[] unsafeOctets) dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; cp >>>= 2; dest[2] = UPPER_HEX_DIGITS[cp]; - size.addAndGet(3); + size[0] += 3; return dest; } else if (cp <= 0x10ffff) { char[] dest = new char[12]; @@ -289,7 +287,7 @@ private static char[] escape(int cp, AtomicInteger size, boolean[] unsafeOctets) dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; cp >>>= 2; dest[2] = UPPER_HEX_DIGITS[cp & 0x7]; - size.addAndGet(4); + size[0] += 4; return dest; } else { // If this ever happens it is due to bug in UnicodeEscaper, not bad input. From e7f533d2c4bfb391aff774f449e16727c8b4fba6 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Fri, 28 Feb 2025 10:53:47 -0500 Subject: [PATCH 27/31] creating EscapedData class to store escaped data --- .../trace/core/baggage/BaggagePropagator.java | 16 +++++---- .../datadog/trace/core/util/EscapedData.java | 36 +++++++++++++++++++ .../trace/core/util/PercentEscaper.java | 36 ++++++++++--------- 3 files changed, 66 insertions(+), 22 deletions(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.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 79931336d03..67f8702291e 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 @@ -6,6 +6,7 @@ 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.core.util.PercentEscaper; import java.io.UnsupportedEncodingException; import java.net.URLDecoder; @@ -69,16 +70,19 @@ public void inject(Context context, C carrier, CarrierSetter setter) { int currentBytes = 0; StringBuilder baggageText = new StringBuilder(); for (final Map.Entry entry : baggageContext.asMap().entrySet()) { - int[] pairSize = {1}; // if there are already baggage items processed, add and allocate bytes for a comma + int extraBytes = 1; if (processedBaggage != 0) { baggageText.append(','); - pairSize[0]++; + extraBytes++; } - baggageText.append(UTF_ESCAPER.escapeKey(entry.getKey(), pairSize)); + EscapedData escapedKey = UTF_ESCAPER.escapeKey(entry.getKey()); + EscapedData escapedVal = UTF_ESCAPER.escapeValue(entry.getValue()); + + baggageText.append(escapedKey.getData()); baggageText.append('='); - baggageText.append(UTF_ESCAPER.escapeValue(entry.getValue(), pairSize)); + baggageText.append(escapedVal.getData()); processedBaggage++; // reached the max number of baggage items allowed @@ -86,11 +90,11 @@ public void inject(Context context, C carrier, CarrierSetter setter) { break; } // Drop newest k/v pair if adding it leads to exceeding the limit - if (currentBytes + pairSize[0] > maxBytes) { + if (currentBytes + escapedKey.getSize() + escapedVal.getSize() + extraBytes > maxBytes) { baggageText.setLength(currentBytes); break; } - currentBytes += pairSize[0]; + currentBytes += escapedKey.getSize() + escapedVal.getSize() + extraBytes; } String baggageString = baggageText.toString(); 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 new file mode 100644 index 00000000000..fef31fbbaac --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/util/EscapedData.java @@ -0,0 +1,36 @@ +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 129d5255dc8..0a9411b3f1b 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,25 +107,25 @@ private static boolean[] createUnsafeOctets(String safeChars) { return octets; } - public String escapeKey(String s, int[] size) { - return escape(s, size, unsafeKeyOctets); + public EscapedData escapeKey(String s) { + return escape(s, unsafeKeyOctets); } - public String escapeValue(String s, int[] size) { - return escape(s, size, unsafeKeyOctets); + public EscapedData escapeValue(String s) { + return escape(s, unsafeValOctets); } /** Escape the provided String, using percent-style URL Encoding. */ - public String escape(String s, int[] size, boolean[] unsafeOctets) { + public EscapedData escape(String s, boolean[] unsafeOctets) { + int size = 0; 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, size, unsafeOctets); + return escapeSlow(s, index, unsafeOctets); } - size[0]++; } - return s; + return new EscapedData(s, slen); } /* @@ -147,13 +147,14 @@ public String escape(String s, int[] size, boolean[] unsafeOctets) { * @throws NullPointerException if {@code string} is null * @throws IllegalArgumentException if invalid surrogate characters are encountered */ - private static String escapeSlow(String s, int index, int[] size, boolean[] unsafeOctets) { + private static EscapedData 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); while (index < end) { int cp = codePointAt(s, index, end); @@ -163,7 +164,7 @@ private static String escapeSlow(String s, int index, int[] size, boolean[] unsa // 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, size, unsafeOctets); + char[] escaped = escape(cp, data, unsafeOctets); int nextIndex = index + (Character.isSupplementaryCodePoint(cp) ? 2 : 1); if (escaped != null) { int charsSkipped = index - unescapedChunkStart; @@ -201,7 +202,10 @@ private static String escapeSlow(String s, int index, int[] size, boolean[] unsa s.getChars(unescapedChunkStart, end, dest, destIndex); destIndex = endIndex; } - return new String(dest, 0, destIndex); + data.addSize(charsSkipped); //Adding characters in-between characters that want to be encoded + + data.setData(new String(dest, 0, destIndex)); + return data; } private static int nextEscapeIndex(CharSequence csq, int index, int end, boolean[] unsafeOctets) { @@ -217,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, int[] size, boolean[] unsafeOctets) { + private static char[] escape(int cp, EscapedData data, 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]) { @@ -229,7 +233,7 @@ private static char[] escape(int cp, int[] size, boolean[] unsafeOctets) { dest[0] = '%'; dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; dest[1] = UPPER_HEX_DIGITS[cp >>> 4]; - size[0]++; + data.incrementSize(); return dest; } else if (cp <= 0x7ff) { // Two byte UTF-8 characters [cp >= 0x80 && cp <= 0x7ff] @@ -244,7 +248,7 @@ private static char[] escape(int cp, int[] size, boolean[] unsafeOctets) { dest[2] = UPPER_HEX_DIGITS[cp & 0xF]; cp >>>= 4; dest[1] = UPPER_HEX_DIGITS[0xC | cp]; - size[0] += 2; + data.addSize(2); return dest; } else if (cp <= 0xffff) { // Three byte UTF-8 characters [cp >= 0x800 && cp <= 0xffff] @@ -263,7 +267,7 @@ private static char[] escape(int cp, int[] size, boolean[] unsafeOctets) { dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; cp >>>= 2; dest[2] = UPPER_HEX_DIGITS[cp]; - size[0] += 3; + data.addSize(3); return dest; } else if (cp <= 0x10ffff) { char[] dest = new char[12]; @@ -287,7 +291,7 @@ private static char[] escape(int cp, int[] size, boolean[] unsafeOctets) { dest[4] = UPPER_HEX_DIGITS[0x8 | (cp & 0x3)]; cp >>>= 2; dest[2] = UPPER_HEX_DIGITS[cp & 0x7]; - size[0] += 4; + data.addSize(4); return dest; } else { // If this ever happens it is due to bug in UnicodeEscaper, not bad input. From 99165900493c2516fbfc08f35a2288264f4a4a45 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Fri, 28 Feb 2025 11:21:36 -0500 Subject: [PATCH 28/31] nit comments and merge conflicts --- .../datadog/trace/core/util/EscapedData.java | 16 ++++++++-------- .../datadog/trace/core/util/PercentEscaper.java | 2 +- internal-api/build.gradle | 4 +++- 3 files changed, 12 insertions(+), 10 deletions(-) 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 index fef31fbbaac..0ff428d2f60 100644 --- 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 @@ -4,33 +4,33 @@ public class EscapedData { private String data; private int size; - public EscapedData(String data, int size){ + public EscapedData(String data, int size) { this.data = data; this.size = size; } - public EscapedData(){ + public EscapedData() { this.data = ""; this.size = 0; } - public String getData(){ + public String getData() { return data; } - public int getSize(){ + public int getSize() { return size; } - public void setData(String data){ + public void setData(String data) { this.data = data; } - public void incrementSize(){ + public void incrementSize() { size++; } - public void addSize(int delta){ - size+=delta; + 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 0a9411b3f1b..8bd6a7b21d5 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 @@ -202,7 +202,7 @@ 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 + data.addSize(charsSkipped); // Adding characters in-between characters that want to be encoded data.setData(new String(dest, 0, destIndex)); return data; diff --git a/internal-api/build.gradle b/internal-api/build.gradle index 60b3581db97..0df48ba71ae 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -205,7 +205,9 @@ excludedClassesCoverage += [ 'datadog.trace.util.stacktrace.StackTraceFrame', 'datadog.trace.api.iast.VulnerabilityMarks', 'datadog.trace.api.iast.securitycontrol.SecurityControlHelper', - 'datadog.trace.api.iast.securitycontrol.SecurityControl' + 'datadog.trace.api.iast.securitycontrol.SecurityControl', + // Otel encoder class tested in BaggagePropagatorTest + 'datadog.trace.core.util.PercentEscaper' ] excludedClassesBranchCoverage = [ 'datadog.trace.api.ProductActivationConfig', From 586468c93bacaa8e0711998077b108fa81176e06 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Fri, 28 Feb 2025 11:28:22 -0500 Subject: [PATCH 29/31] spotless --- .../trace/bootstrap/instrumentation/api/AgentPropagation.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java index f707ccb51b8..a25c0abfee5 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentPropagation.java @@ -14,7 +14,7 @@ public final class AgentPropagation { public static final Concern TRACING_CONCERN = named("tracing"); public static final Concern BAGGAGE_CONCERN = named("baggage"); public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray"); - + // TODO DSM propagator should run after the other propagators as it stores the pathway context // TODO into the span context for now. Remove priority after the migration is complete. public static final Concern DSM_CONCERN = withPriority("data-stream-monitoring", 110); From bebdf7cb1fa035363259f60f4c6e415678d48414 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Mar 2025 11:23:38 -0500 Subject: [PATCH 30/31] updating build file --- internal-api/build.gradle | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal-api/build.gradle b/internal-api/build.gradle index 0df48ba71ae..c3ca56589c6 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -206,8 +206,8 @@ excludedClassesCoverage += [ 'datadog.trace.api.iast.VulnerabilityMarks', 'datadog.trace.api.iast.securitycontrol.SecurityControlHelper', 'datadog.trace.api.iast.securitycontrol.SecurityControl', - // Otel encoder class tested in BaggagePropagatorTest - 'datadog.trace.core.util.PercentEscaper' + // BaggageContext class tested in BaggagePropagatorTest in dd-trace-core + 'datadog.trace.core.util.BaggageContext' ] excludedClassesBranchCoverage = [ 'datadog.trace.api.ProductActivationConfig', From c4f6d7d60336830cf3254f4ddae413b208db9b67 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Mar 2025 11:56:13 -0500 Subject: [PATCH 31/31] nit update --- internal-api/build.gradle | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/internal-api/build.gradle b/internal-api/build.gradle index c3ca56589c6..5dbe7d16de6 100644 --- a/internal-api/build.gradle +++ b/internal-api/build.gradle @@ -181,6 +181,8 @@ 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", @@ -205,9 +207,7 @@ excludedClassesCoverage += [ 'datadog.trace.util.stacktrace.StackTraceFrame', 'datadog.trace.api.iast.VulnerabilityMarks', 'datadog.trace.api.iast.securitycontrol.SecurityControlHelper', - 'datadog.trace.api.iast.securitycontrol.SecurityControl', - // BaggageContext class tested in BaggagePropagatorTest in dd-trace-core - 'datadog.trace.core.util.BaggageContext' + 'datadog.trace.api.iast.securitycontrol.SecurityControl' ] excludedClassesBranchCoverage = [ 'datadog.trace.api.ProductActivationConfig',