From ae71f519f1d20adfc44407e8c2e804c6821ef84f Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Feb 2025 16:34:29 -0500 Subject: [PATCH 1/2] 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 4eb414f3394..99a3573f3a7 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 e7b2efbf020..18134ac41fc 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -188,6 +188,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; @@ -986,6 +988,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 = @@ -2215,6 +2223,14 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public int getTraceBaggageMaxItems() { + return traceBaggageMaxItems; + } + + public int getTraceBaggageMaxBytes() { + return traceBaggageMaxBytes; + } + public int getClockSyncPeriod() { return clockSyncPeriod; } From 22bc0098deb79fa9678fa9540c002bbef64fd51d Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Mon, 3 Feb 2025 16:55:30 -0500 Subject: [PATCH 2/2] 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;