From 796b57e503894a1c247ed253148049c1fdfdcc3e Mon Sep 17 00:00:00 2001 From: Jordan Wong <61242306+jordan-wong@users.noreply.github.com> Date: Wed, 7 May 2025 13:37:58 -0400 Subject: [PATCH 01/27] Implement APIGW Inferred Proxy Spans (#8336) * Creates inferred proxy spans for API Gateway calls via presence of http headers --------- Co-authored-by: Zarir Hamza --- .../datadog/context/InferredProxyContext.java | 50 ++ .../propagation/InferredProxyPropagator.java | 74 +++ .../context/InferredProxyHandlingTest.java | 465 ++++++++++++++++++ .../decorator/HttpServerDecorator.java | 451 +++++++++++++++++ .../trace/api/config/TracerConfig.java | 3 + .../java/datadog/trace/core/CoreTracer.java | 5 + .../main/java/datadog/trace/api/Config.java | 7 + .../instrumentation/api/AgentPropagation.java | 2 +- 8 files changed, 1056 insertions(+), 1 deletion(-) create mode 100644 components/context/src/main/java/datadog/context/InferredProxyContext.java create mode 100644 components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java create mode 100644 components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java diff --git a/components/context/src/main/java/datadog/context/InferredProxyContext.java b/components/context/src/main/java/datadog/context/InferredProxyContext.java new file mode 100644 index 00000000000..51eecc4cc02 --- /dev/null +++ b/components/context/src/main/java/datadog/context/InferredProxyContext.java @@ -0,0 +1,50 @@ +package datadog.context; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class InferredProxyContext implements ImplicitContextKeyed { + public static final ContextKey CONTEXT_KEY = + ContextKey.named("inferred-proxy-key"); + private final Map inferredProxy; + + public static InferredProxyContext fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + public InferredProxyContext(Map contextInfo) { + this.inferredProxy = + (contextInfo == null || contextInfo.isEmpty()) + ? new HashMap<>() + : new HashMap<>(contextInfo); + } + + public InferredProxyContext() { + this.inferredProxy = new HashMap<>(); + } + + public Map getInferredProxyContext() { + return Collections.unmodifiableMap(inferredProxy); + } + + public void putInferredProxyInfo(String key, String value) { + inferredProxy.put(key, value); + } + + public void removeInferredProxyInfo(String key) { + inferredProxy.remove(key); + } + + /** + * Creates a new context with this value under its chosen key. + * + * @param context the context to copy the original values from. + * @return the new context with the implicitly keyed value. + * @see Context#with(ImplicitContextKeyed) + */ + @Override + public Context storeInto(Context context) { + return context.with(CONTEXT_KEY, this); + } +} diff --git a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java new file mode 100644 index 00000000000..69e5a0e896e --- /dev/null +++ b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java @@ -0,0 +1,74 @@ +package datadog.context.propagation; + +import datadog.context.Context; +import datadog.context.InferredProxyContext; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; + +public class InferredProxyPropagator implements Propagator { + public static final String INFERRED_PROXY_KEY = "x-dd-proxy"; + /** + * METHOD STUB: InferredProxy is currently not meant to be injected to downstream services Injects + * a context into a downstream service using the given carrier. + * + * @param context the context containing the values to be injected. + * @param carrier the instance that will receive the key/value pairs to propagate. + * @param setter the callback to set key/value pairs into the carrier. + */ + @Override + public void inject(Context context, C carrier, CarrierSetter setter) {} + + /** + * Extracts a context from un upstream service. + * + * @param context the base context to store the extracted values on top, use {@link + * Context#root()} for a default base context. + * @param carrier the instance to fetch the propagated key/value pairs from. + * @param visitor the callback to walk over the carrier and extract its key/value pais. + * @return A context with the extracted values on top of the given base context. + */ + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + if (context == null || carrier == null || visitor == null) { + return context; + } + InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); + visitor.forEachKeyValue(carrier, extractor); + + InferredProxyContext extractedContext = extractor.extractedContext; + if (extractedContext == null) { + return context; + } + return extractedContext.storeInto(context); + } + + public static class InferredProxyContextExtractor implements BiConsumer { + private InferredProxyContext extractedContext; + + InferredProxyContextExtractor() {} + + private Map parseInferredProxyHeaders(String input) { + Map parsedHeaders = new HashMap<>(); + return parsedHeaders; + } + + /** + * Performs this operation on the given arguments. + * + * @param key the first input argument from an http header + * @param value the second input argument from an http header + */ + @Override + public void accept(String key, String value) { + if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY)) { + return; + } + Map inferredProxyMap = parseInferredProxyHeaders(value); + if (extractedContext == null) { + extractedContext = new InferredProxyContext(); + } + extractedContext.putInferredProxyInfo(key, value); + } + } +} diff --git a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java new file mode 100644 index 00000000000..53ddf5cb12a --- /dev/null +++ b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java @@ -0,0 +1,465 @@ +package datadog.context; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.InferredProxyPropagator; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; // For @Test on nested class methods +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +class InferredProxyHandlingTest { + + // Define header key constants locally for the test + static final String PROXY_SYSTEM_KEY = "x-dd-proxy-system"; + static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; + static final String PROXY_PATH_KEY = "x-dd-proxy-path"; + static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; + static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; + + private InferredProxyPropagator propagator; + + @BeforeEach + void setUp() { + propagator = new InferredProxyPropagator(); + } + + // Moved @MethodSource providers to the outer class and made them static + static Stream validHeadersProviderForPropagator() { + Map allStandard = new HashMap<>(); + allStandard.put(PROXY_SYSTEM_KEY, "aws-apigw"); // The only currently supported system + allStandard.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + allStandard.put(PROXY_PATH_KEY, "/foo"); + allStandard.put(PROXY_HTTP_METHOD_KEY, "GET"); + allStandard.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); + + return Stream.of( + Arguments.of( + "all standard headers (aws-apigw)", + allStandard, + "aws-apigw", + "12345", + "/foo", + "GET", + "api.example.com", + null, + null)); + } + + static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + missingSystem.put(PROXY_PATH_KEY, "/foo"); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); + missingTime.put(PROXY_PATH_KEY, "/foo"); + + return Stream.of( + Arguments.of("PROXY_SYSTEM_KEY missing", missingSystem), + Arguments.of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); + } + + // Simple Map visitor for tests (can remain static or non-static in outer class) + static class MapVisitor implements CarrierVisitor> { + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + if (carrier == null) { + return; + } + carrier.forEach(visitor); + } + } + + // Custom visitor to test null key path in the extractor - MOVED HERE and made static + static class NullKeyTestVisitor implements CarrierVisitor> { + private final BiConsumer actualExtractorAccept; + + NullKeyTestVisitor(BiConsumer actualExtractorAccept) { + this.actualExtractorAccept = actualExtractorAccept; + } + + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + if (actualExtractorAccept != null) { + actualExtractorAccept.accept(null, "valueForNullKey"); + } + } + } + + @Nested + @DisplayName("InferredProxyPropagator Tests") + class PropagatorTests { // Kept non-static + + @ParameterizedTest(name = "{0}") + @MethodSource( + "datadog.context.InferredProxyHandlingTest#validHeadersProviderForPropagator") // Fully + // qualified + // name + @DisplayName("Should extract InferredProxyContext when valid headers are present") + void testSuccessfulExtraction( + String description, + Map headers, + String expectedSystem, + String expectedTimeMs, + String expectedPath, + String expectedMethod, + String expectedDomain, + String expectedExtraKey, + String expectedExtraValue) { + + Context rootContext = Context.root(); + // Now accesses the outer class's propagator instance field + Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull( + inferredProxyContext, "InferredProxyContext should not be null for: " + description); + Map actualProxyData = inferredProxyContext.getInferredProxyContext(); + assertEquals(expectedSystem, actualProxyData.get(PROXY_SYSTEM_KEY)); + assertEquals(expectedTimeMs, actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + assertEquals(expectedPath, actualProxyData.get(PROXY_PATH_KEY)); + assertEquals(expectedMethod, actualProxyData.get(PROXY_HTTP_METHOD_KEY)); + assertEquals(expectedDomain, actualProxyData.get(PROXY_DOMAIN_NAME_KEY)); + if (expectedExtraKey != null) { + assertEquals(expectedExtraValue, actualProxyData.get(expectedExtraKey)); + } + } + + @ParameterizedTest(name = "{0}") + @MethodSource( + "datadog.context.InferredProxyHandlingTest#invalidOrMissingHeadersProviderForPropagator") // Fully qualified name + @DisplayName("Should create InferredProxyContext even if some critical headers are missing") + void testExtractionWithMissingCriticalHeaders(String description, Map headers) { + Context rootContext = Context.root(); + Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull( + inferredProxyContext, + "InferredProxyContext should still be created if any x-dd-proxy-* header is present for: " + + description); + Map actualProxyData = inferredProxyContext.getInferredProxyContext(); + + if (headers.containsKey(PROXY_SYSTEM_KEY)) { + assertEquals(headers.get(PROXY_SYSTEM_KEY), actualProxyData.get(PROXY_SYSTEM_KEY)); + } else { + assertNull(actualProxyData.get(PROXY_SYSTEM_KEY)); + } + if (headers.containsKey(PROXY_REQUEST_TIME_MS_KEY)) { + assertEquals( + headers.get(PROXY_REQUEST_TIME_MS_KEY), actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + } else { + assertNull(actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); + } + } + + @Test + @DisplayName("Should not extract InferredProxyContext if no relevant headers are present") + void testNoRelevantHeaders() { + Map carrier = new HashMap<>(); + carrier.put("x-unrelated-header", "value"); + carrier.put("another-header", "othervalue"); + Context rootContext = Context.root(); + + Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNull( + inferredProxyContext, + "InferredProxyContext should be null if no x-dd-proxy-* headers are found"); + } + + @Test + @DisplayName("Should return original context if carrier is null") + void testNullCarrier() { + InferredProxyContext initialData = + new InferredProxyContext(Collections.singletonMap("test", "value")); + Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); + + Context extractedOuterContext = propagator.extract(rootContext, null, new MapVisitor()); + + assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); + assertEquals( + "value", + InferredProxyContext.fromContext(extractedOuterContext) + .getInferredProxyContext() + .get("test")); + } + + @Test + @DisplayName("Should return original context if visitor is null") + void testNullVisitor() { + Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); + InferredProxyContext initialData = + new InferredProxyContext(Collections.singletonMap("test", "value")); + Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); + + Context extractedOuterContext = propagator.extract(rootContext, carrier, null); + + assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); + assertEquals( + "value", + InferredProxyContext.fromContext(extractedOuterContext) + .getInferredProxyContext() + .get("test")); + } + + @Test + @DisplayName("Should return original context if context is null") + void testNullContext() { + Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); + Context extractedOuterContext = propagator.extract(null, carrier, new MapVisitor()); + assertNull(extractedOuterContext, "Context should remain null if passed as null"); + } + + @Test + @DisplayName("Extractor should handle multiple proxy headers") + void testMultipleProxyHeaders() { + Map carrier = new HashMap<>(); + carrier.put(PROXY_SYSTEM_KEY, "aws-apigw"); + carrier.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + carrier.put("x-dd-proxy-custom", "value1"); // First proxy header + carrier.put("x-dd-proxy-another", "value2"); // Second proxy header + + Context rootContext = Context.root(); + Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); + InferredProxyContext inferredProxyContext = + InferredProxyContext.fromContext(extractedOuterContext); + + assertNotNull(inferredProxyContext); + // Check if both headers were stored (covers extractedContext == null being false) + assertEquals( + "value1", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-custom")); + assertEquals( + "value2", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-another")); + assertEquals( + "aws-apigw", inferredProxyContext.getInferredProxyContext().get(PROXY_SYSTEM_KEY)); + } + + @Test + @DisplayName("Extractor accept method should handle null/empty keys") + void testExtractorAcceptNullEmptyKeys() { + Context rootContext = Context.root(); + + // Test null key - HashMap doesn't allow null keys. Standard HTTP visitors + // also typically don't yield null keys. Testing this branch is difficult + // without a custom visitor or modifying the source. Relying on coverage report + // or assuming standard carriers won't provide null keys. + + // Test empty key + Map carrierWithEmptyKey = new HashMap<>(); + carrierWithEmptyKey.put("", "emptyKeyValue"); // Add empty key + carrierWithEmptyKey.put(PROXY_SYSTEM_KEY, "aws-apigw"); // Add a valid key too + + Context contextAfterEmpty = + propagator.extract(rootContext, carrierWithEmptyKey, new MapVisitor()); + InferredProxyContext ipcEmpty = InferredProxyContext.fromContext(contextAfterEmpty); + + // The propagator should ignore the empty key entry entirely. + assertNotNull(ipcEmpty, "Context should be created due to valid key"); + assertNull(ipcEmpty.getInferredProxyContext().get(""), "Empty key should not be stored"); + assertEquals( + "aws-apigw", + ipcEmpty.getInferredProxyContext().get(PROXY_SYSTEM_KEY), + "Valid key should still be stored"); + assertEquals(1, ipcEmpty.getInferredProxyContext().size(), "Only valid key should be stored"); + } + + @Test + @DisplayName( + "Extractor accept method should handle explicitly passed null key via custom visitor") + void testExtractorAcceptExplicitNullKey() { + Context rootContext = Context.root(); + Map carrier = new HashMap<>(); // Carrier can be empty for this test + + // We need to get a handle to the internal BiConsumer (the InferredProxyContextExtractor + // instance). + // The extract method will create one. We can pass a visitor that captures it. + + final BiConsumer[] extractorHolder = new BiConsumer[1]; + + CarrierVisitor> capturingVisitor = + (cr, bic) -> { + extractorHolder[0] = bic; // Capture the BiConsumer + // Optionally, call the original MapVisitor if we still want normal processing after + // capture + // new MapVisitor().forEachKeyValue(cr, bic); + }; + + // This first call is primarily to get a reference to the internal extractor + propagator.extract(rootContext, carrier, capturingVisitor); + + assertNotNull(extractorHolder[0], "Failed to capture the internal extractor instance"); + + // Now use a new custom visitor to specifically test the null key path + // on the captured extractor instance (though this isn't how extract is typically used). + // A more direct way to test the BiConsumer if it were accessible or if the design allowed it. + // For now, we directly call accept on the captured one. + extractorHolder[0].accept(null, "valueForNullKey"); + + // The goal is JaCoCo coverage. Asserting internal state of the extractor is hard without + // reflection. + // We can verify that the context remains unchanged or as expected if no valid headers + // processed. + InferredProxyContext ipc = + InferredProxyContext.fromContext( + rootContext); // or context returned by a second extract call + assertNull(ipc, "Context should not have InferredProxyContext from only a null key call"); + } + } + + @Nested + @DisplayName("InferredProxyContext Tests") + class ContextUnitTests { + + @Test + @DisplayName("Default constructor should create an empty context map") + void testDefaultConstructor() { + InferredProxyContext ipc = new InferredProxyContext(); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + + @Test + @DisplayName("Constructor with map should initialize context map") + void testMapConstructor() { + Map initialData = new HashMap<>(); + initialData.put("key1", "value1"); + initialData.put("key2", "value2"); + + InferredProxyContext ipc = new InferredProxyContext(initialData); + assertNotNull(ipc.getInferredProxyContext()); + assertEquals(2, ipc.getInferredProxyContext().size()); + assertEquals("value1", ipc.getInferredProxyContext().get("key1")); + assertEquals("value2", ipc.getInferredProxyContext().get("key2")); + + initialData.put("key3", "value3"); // Modify original map + assertNull(ipc.getInferredProxyContext().get("key3"), "Internal map should be a copy"); + } + + @Test + @DisplayName("putInferredProxyInfo should add to the context map") + void testPutInfo() { + InferredProxyContext ipc = new InferredProxyContext(); + ipc.putInferredProxyInfo("system", "aws-apigw"); + ipc.putInferredProxyInfo("time", "12345"); + + Map contextMap = ipc.getInferredProxyContext(); + assertEquals(2, contextMap.size()); + assertEquals("aws-apigw", contextMap.get("system")); + assertEquals("12345", contextMap.get("time")); + + ipc.putInferredProxyInfo("system", "azure-func"); // Overwrite + assertEquals("azure-func", contextMap.get("system")); + assertEquals(2, contextMap.size()); + } + + @Test + @DisplayName("removeInferredProxyInfo should remove from the context map") + void testRemoveInfo() { + Map initialData = new HashMap<>(); + initialData.put("key1", "value1"); + initialData.put("key2", "value2"); + InferredProxyContext ipc = new InferredProxyContext(initialData); + + ipc.removeInferredProxyInfo("key1"); + Map contextMap = ipc.getInferredProxyContext(); + assertEquals(1, contextMap.size()); + assertNull(contextMap.get("key1")); + assertEquals("value2", contextMap.get("key2")); + + ipc.removeInferredProxyInfo("nonexistent"); // Remove non-existent + assertEquals(1, contextMap.size()); + } + + @Test + @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") + void testStoreAndFromContext() { + InferredProxyContext ipcToStore = new InferredProxyContext(); + ipcToStore.putInferredProxyInfo("customKey", "customValue"); + + Context rootContext = Context.root(); + Context contextWithValue = ipcToStore.storeInto(rootContext); + assertNotNull(contextWithValue); + + InferredProxyContext retrievedIpc = InferredProxyContext.fromContext(contextWithValue); + assertNotNull(retrievedIpc); + assertEquals("customValue", retrievedIpc.getInferredProxyContext().get("customKey")); + + assertNull( + InferredProxyContext.fromContext(rootContext), + "Original root context should not be affected"); + + Context cleanContext = Context.root(); + assertNull( + InferredProxyContext.fromContext(cleanContext), + "fromContext on clean context should be null"); + } + + @Test + @DisplayName("getInferredProxyContext should return an unmodifiable map or a copy") + void testGetInferredProxyContextImmutability() { + InferredProxyContext ipc = new InferredProxyContext(); + ipc.putInferredProxyInfo("key1", "value1"); + + Map retrievedMap = ipc.getInferredProxyContext(); + assertNotNull(retrievedMap); + assertEquals("value1", retrievedMap.get("key1")); + + boolean threwUnsupported = false; + try { + retrievedMap.put("newKey", "newValue"); + } catch (UnsupportedOperationException e) { + threwUnsupported = true; + } + // Depending on whether InferredProxyContext.getInferredProxyContext() returns a direct + // reference or a copy, + // this assertion might change. If it returns a direct mutable reference, threwUnsupported + // will be false. + // If it returns an unmodifiable view or a copy, attempts to modify might throw or simply not + // affect the original. + // For now, we check that the original context was not changed. + assertEquals( + 1, ipc.getInferredProxyContext().size(), "Internal map size should remain unchanged"); + assertEquals( + "value1", + ipc.getInferredProxyContext().get("key1"), + "Internal map content should remain unchanged"); + // If it MUST be unmodifiable, add: assertTrue(threwUnsupported, "Retrieved map should be + // unmodifiable"); + } + + @Test + @DisplayName("Constructor with null map should create an empty context map") + void testNullMapConstructor() { + InferredProxyContext ipc = new InferredProxyContext(null); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + + @Test + @DisplayName("Constructor with empty map should create an empty context map") + void testEmptyMapConstructor() { + Map emptyMap = Collections.emptyMap(); + InferredProxyContext ipc = new InferredProxyContext(emptyMap); + assertNotNull(ipc.getInferredProxyContext()); + assertTrue(ipc.getInferredProxyContext().isEmpty()); + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 25684aeb42e..5d8a7ee5974 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -11,8 +11,12 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.context.Context; import datadog.context.propagation.Propagators; +import datadog.context.InferredProxyContext; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.DDTags; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.TraceConfig; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; import datadog.trace.api.gateway.BlockResponseFunction; @@ -22,10 +26,12 @@ import datadog.trace.api.gateway.IGSpanInfo; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.interceptor.MutableSpan; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -38,6 +44,8 @@ import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver; import java.net.InetAddress; import java.util.BitSet; +import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.Locale; import java.util.Map; import java.util.concurrent.ExecutionException; @@ -52,9 +60,401 @@ public abstract class HttpServerDecorator extends ServerDecorator { + class InferredProxySpanGroup implements AgentSpan { + private final AgentSpan inferredProxySpan; + private final AgentSpan serverSpan; + + InferredProxySpanGroup(AgentSpan inferredProxySpan, AgentSpan serverSpan) { + this.inferredProxySpan = inferredProxySpan; + this.serverSpan = serverSpan; + } + + @Override + public DDTraceId getTraceId() { + return serverSpan.getTraceId(); + } + + @Override + public long getSpanId() { + return serverSpan.getSpanId(); + } + + @Override + public AgentSpan setTag(String key, boolean value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, int value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, long value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, double value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, String value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, CharSequence value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setTag(String key, Object value) { + return serverSpan.setTag(key, value); + } + + /** + * @param map + * @return + */ + @Override + public AgentSpan setAllTags(Map map) { + return null; + } + + @Override + public AgentSpan setTag(String key, Number value) { + return serverSpan.setTag(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, int value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, long value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setMetric(CharSequence key, double value) { + return serverSpan.setMetric(key, value); + } + + @Override + public AgentSpan setSpanType(CharSequence type) { + return serverSpan.setSpanType(type); + } + + @Override + public Object getTag(String key) { + return serverSpan.getTag(key); + } + + @Override + public AgentSpan setError(boolean error) { + serverSpan.setError(error); + if (inferredProxySpan != null) { + inferredProxySpan.setError(error); + } + return this; + } + + @Override + public AgentSpan setError(boolean error, byte priority) { + serverSpan.setError(error, priority); + if (inferredProxySpan != null) { + inferredProxySpan.setError(error, priority); + } + return this; + } + + @Override + public AgentSpan setMeasured(boolean measured) { + return serverSpan.setMeasured(measured); + } + + @Override + public AgentSpan setErrorMessage(String errorMessage) { + return serverSpan.setErrorMessage(errorMessage); + } + + @Override + public AgentSpan addThrowable(Throwable throwable) { + serverSpan.addThrowable(throwable); + if (inferredProxySpan != null) { + inferredProxySpan.addThrowable(throwable); + } + return this; + } + + @Override + public AgentSpan addThrowable(Throwable throwable, byte errorPriority) { + serverSpan.addThrowable(throwable, errorPriority); + if (inferredProxySpan != null) { + inferredProxySpan.addThrowable(throwable, errorPriority); + } + return this; + } + + @Override + public AgentSpan getLocalRootSpan() { + return serverSpan.getLocalRootSpan(); + } + + @Override + public boolean isSameTrace(AgentSpan otherSpan) { + return serverSpan.isSameTrace(otherSpan); + } + + @Override + public AgentSpanContext context() { + return serverSpan.context(); + } + + @Override + public String getBaggageItem(String key) { + return serverSpan.getBaggageItem(key); + } + + @Override + public AgentSpan setBaggageItem(String key, String value) { + return serverSpan.setBaggageItem(key, value); + } + + @Override + public AgentSpan setHttpStatusCode(int statusCode) { + serverSpan.setHttpStatusCode(statusCode); + if (inferredProxySpan != null) { + inferredProxySpan.setHttpStatusCode(statusCode); + } + return this; + } + + @Override + public short getHttpStatusCode() { + return serverSpan.getHttpStatusCode(); + } + + @Override + public void finish() { + serverSpan.finish(); + if (inferredProxySpan != null) { + inferredProxySpan.finish(); + } + } + + @Override + public void finish(long finishMicros) { + serverSpan.finish(finishMicros); + if (inferredProxySpan != null) { + inferredProxySpan.finish(finishMicros); + } + } + + @Override + public void finishWithDuration(long durationNanos) { + serverSpan.finishWithDuration(durationNanos); + if (inferredProxySpan != null) { + inferredProxySpan.finishWithDuration(durationNanos); + } + } + + @Override + public void beginEndToEnd() { + serverSpan.beginEndToEnd(); + } + + @Override + public void finishWithEndToEnd() { + serverSpan.finishWithEndToEnd(); + if (inferredProxySpan != null) { + inferredProxySpan.finishWithEndToEnd(); + } + } + + @Override + public boolean phasedFinish() { + final boolean ret = serverSpan.phasedFinish(); + if (inferredProxySpan != null) { + inferredProxySpan.phasedFinish(); + } + return ret; + } + + @Override + public void publish() { + serverSpan.publish(); + } + + @Override + public CharSequence getSpanName() { + return serverSpan.getSpanName(); + } + + @Override + public void setSpanName(CharSequence spanName) { + serverSpan.setSpanName(spanName); + } + + @Deprecated + @Override + public boolean hasResourceName() { + return serverSpan.hasResourceName(); + } + + @Override + public byte getResourceNamePriority() { + return serverSpan.getResourceNamePriority(); + } + + @Override + public AgentSpan setResourceName(CharSequence resourceName) { + return serverSpan.setResourceName(resourceName); + } + + @Override + public AgentSpan setResourceName(CharSequence resourceName, byte priority) { + return serverSpan.setResourceName(resourceName, priority); + } + + @Override + public RequestContext getRequestContext() { + return serverSpan.getRequestContext(); + } + + @Override + public Integer forceSamplingDecision() { + return serverSpan.forceSamplingDecision(); + } + + @Override + public AgentSpan setSamplingPriority(int newPriority, int samplingMechanism) { + return serverSpan.setSamplingPriority(newPriority, samplingMechanism); + } + + @Override + public TraceConfig traceConfig() { + return serverSpan.traceConfig(); + } + + @Override + public void addLink(AgentSpanLink link) { + serverSpan.addLink(link); + } + + @Override + public AgentSpan setMetaStruct(String field, Object value) { + return serverSpan.setMetaStruct(field, value); + } + + @Override + public boolean isOutbound() { + return serverSpan.isOutbound(); + } + + @Override + public AgentSpan asAgentSpan() { + return serverSpan.asAgentSpan(); + } + + @Override + public long getStartTime() { + return serverSpan.getStartTime(); + } + + @Override + public long getDurationNano() { + return serverSpan.getDurationNano(); + } + + @Override + public CharSequence getOperationName() { + return serverSpan.getOperationName(); + } + + @Override + public MutableSpan setOperationName(CharSequence serviceName) { + return serverSpan.setOperationName(serviceName); + } + + @Override + public String getServiceName() { + return serverSpan.getServiceName(); + } + + @Override + public MutableSpan setServiceName(String serviceName) { + return serverSpan.setServiceName(serviceName); + } + + @Override + public CharSequence getResourceName() { + return serverSpan.getResourceName(); + } + + @Override + public Integer getSamplingPriority() { + return serverSpan.getSamplingPriority(); + } + + @Deprecated + @Override + public MutableSpan setSamplingPriority(int newPriority) { + return serverSpan.setSamplingPriority(newPriority); + } + + @Override + public String getSpanType() { + return serverSpan.getSpanType(); + } + + @Override + public Map getTags() { + return serverSpan.getTags(); + } + + @Override + public boolean isError() { + return serverSpan.isError(); + } + + @Deprecated + @Override + public MutableSpan getRootSpan() { + return serverSpan.getRootSpan(); + } + + @Override + public void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba) { + serverSpan.setRequestBlockingAction(rba); + } + + @Override + public Flow.Action.RequestBlockingAction getRequestBlockingAction() { + return serverSpan.getRequestBlockingAction(); + } + } + private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class); private static final int UNSET_PORT = 0; + public static final String PROXY_SYSTEM = "x-dd-proxy"; + public static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; + public static final String PROXY_PATH = "x-dd-proxy-path"; + public static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; + public static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; + public static final String STAGE = "x-dd-proxy-stage"; + + public static final Map SUPPORTED_PROXIES; + + static { + SUPPORTED_PROXIES = new HashMap<>(); + SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); + } + public static final String DD_SPAN_ATTRIBUTE = "datadog.span"; public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch"; public static final String DD_RUM_INJECTED = "datadog.rum.injected"; @@ -161,6 +561,57 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { return context.with(span); } + private AgentSpan startSpanWithInferredProxy( + String instrumentationName, + datadog.context.Context fullContextForInferredProxy, + AgentSpanContext.Extracted standardExtractedContext) { + + InferredProxyContext inferredProxy = + InferredProxyContext.fromContext(fullContextForInferredProxy); + + if (inferredProxy == null) { + return null; + } + + Map headers = inferredProxy.getInferredProxyContext(); + + // Check if timestamp and proxy system are present + String startTimeStr = headers.get(PROXY_START_TIME_MS); + String proxySystem = headers.get(PROXY_SYSTEM); + + if (startTimeStr == null + || proxySystem == null + || !SUPPORTED_PROXIES.containsKey(proxySystem)) { + return null; + } + + long startTime; + try { + startTime = Long.parseLong(startTimeStr) * 1000; // Convert to microseconds + } catch (NumberFormatException e) { + return null; // Invalid timestamp + } + + AgentSpan apiGtwSpan = + tracer() + .startSpan( + "inferred_proxy", + SUPPORTED_PROXIES.get(proxySystem), + callIGCallbackStart(standardExtractedContext), + startTime); + + apiGtwSpan.setTag(Tags.COMPONENT, proxySystem); + apiGtwSpan.setTag( + DDTags.RESOURCE_NAME, headers.get(PROXY_HTTP_METHOD) + " " + headers.get(PROXY_PATH)); + apiGtwSpan.setTag(DDTags.SERVICE_NAME, headers.get(PROXY_DOMAIN_NAME)); + apiGtwSpan.setTag(DDTags.SPAN_TYPE, "web"); + apiGtwSpan.setTag(Tags.HTTP_METHOD, headers.get(PROXY_HTTP_METHOD)); + apiGtwSpan.setTag(Tags.HTTP_URL, headers.get(PROXY_DOMAIN_NAME) + headers.get(PROXY_PATH)); + apiGtwSpan.setTag("stage", headers.get(STAGE)); + apiGtwSpan.setTag("_dd.inferred_span", 1); + return apiGtwSpan; + } + public AgentSpan onRequest( final AgentSpan span, final CONNECTION connection, 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 0300cef2070..ae5d515cb62 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 @@ -100,6 +100,9 @@ public final class TracerConfig { public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes"; public static final String TRACE_BAGGAGE_TAG_KEYS = "trace.baggage.tag.keys"; + public static final String TRACE_INFERRED_PROXY_SERVICES_ENABLED = + "trace.inferred.proxy.services.enabled"; + public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled"; public static final String CLIENT_IP_ENABLED = "trace.client-ip.enabled"; 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 4c81a579f83..56aeedc6e71 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.TracePropagationBehaviorExtract.IGNORE; 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.INFERRED_PROXY_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; @@ -22,6 +23,7 @@ import datadog.communication.ddagent.SharedCommunicationObjects; import datadog.communication.monitor.Monitoring; import datadog.communication.monitor.Recording; +import datadog.context.propagation.InferredProxyPropagator; import datadog.context.propagation.Propagators; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; @@ -820,6 +822,9 @@ private CoreTracer( && config.getTracePropagationBehaviorExtract() != IGNORE) { Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator(config)); } + if (config.isInferredProxyPropagationEnabled()) { + Propagators.register(INFERRED_PROXY_CONCERN, new InferredProxyPropagator()); + } if (config.isCiVisibilityEnabled()) { if (config.isCiVisibilityTraceSanitationEnabled()) { 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 aae53ec7e07..a2e58598ea5 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -827,6 +827,7 @@ public static String getHostName() { private final int traceBaggageMaxItems; private final int traceBaggageMaxBytes; private final List traceBaggageTagKeys; + private final boolean traceInferredProxyEnabled; private final int clockSyncPeriod; private final boolean logsInjectionEnabled; @@ -1730,6 +1731,8 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins tracePropagationExtractFirst = configProvider.getBoolean( TRACE_PROPAGATION_EXTRACT_FIRST, DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST); + traceInferredProxyEnabled = + configProvider.getBoolean(TRACE_INFERRED_PROXY_SERVICES_ENABLED, false); clockSyncPeriod = configProvider.getInteger(CLOCK_SYNC_PERIOD, DEFAULT_CLOCK_SYNC_PERIOD); @@ -3118,6 +3121,10 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } + public boolean isInferredProxyPropagationEnabled() { + return traceInferredProxyEnabled; + } + public boolean isBaggageExtract() { return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); } 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 fba659bea74..0fbb4484353 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 @@ -18,7 +18,7 @@ public final class AgentPropagation { // TODO: remove this priority once we have a story for replacing TagContext with the Context API public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105); public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray"); - + public static final Concern INFERRED_PROXY_CONCERN = named("inferred-proxy"); // 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 e8c2bafaff3aa416da5b3aa456101ccb0f18092b Mon Sep 17 00:00:00 2001 From: Bruce Bujon Date: Mon, 12 May 2025 11:26:35 +0200 Subject: [PATCH 02/27] feat(serverless): Fix gateway inferred span design Avoid duplicate expensive context extraction Avoid subclassing tracing span for serverless but used serverless context element instead to store / track inferred span while keep tracing feature untouched Improved propagator to not create / capture inferred span context element on invalid data Rework context element to hold the inferred spans and its captured data Release captured data as soon as they span start (never read after this point so reclaiming memory) Refactor context element and propagator into the right package, not context component (product / feature agnostic) Refactor unit tests --- .github/CODEOWNERS | 6 +- .../datadog/context/InferredProxyContext.java | 50 -- .../propagation/InferredProxyPropagator.java | 74 --- .../context/InferredProxyHandlingTest.java | 465 ----------------- .../decorator/HttpServerDecorator.java | 481 +----------------- .../java/datadog/trace/core/CoreTracer.java | 2 +- .../propagation/InferredProxyPropagator.java | 60 +++ .../InferredProxyPropagatorTests.java | 96 ++++ .../main/java/datadog/trace/api/Config.java | 3 + .../trace/api/gateway/InferredProxySpan.java | 108 ++++ .../api/gateway/InferredProxySpanTests.java | 95 ++++ 11 files changed, 393 insertions(+), 1047 deletions(-) delete mode 100644 components/context/src/main/java/datadog/context/InferredProxyContext.java delete mode 100644 components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java delete mode 100644 components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java create mode 100644 dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java create mode 100644 internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java create mode 100644 internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ef437eebae2..96ab0c6452e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -28,8 +28,10 @@ /internal-api/src/test/groovy/datadog/trace/api/sampling @DataDog/apm-sdk-api-java # @DataDog/apm-serverless -/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless -/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless +**/InferredProxy*.java @DataDog/apm-serverless +**/InferredProxy*.groovy @DataDog/apm-serverless # @DataDog/apm-lang-platform-java /.circleci/ @DataDog/apm-lang-platform-java diff --git a/components/context/src/main/java/datadog/context/InferredProxyContext.java b/components/context/src/main/java/datadog/context/InferredProxyContext.java deleted file mode 100644 index 51eecc4cc02..00000000000 --- a/components/context/src/main/java/datadog/context/InferredProxyContext.java +++ /dev/null @@ -1,50 +0,0 @@ -package datadog.context; - -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; - -public class InferredProxyContext implements ImplicitContextKeyed { - public static final ContextKey CONTEXT_KEY = - ContextKey.named("inferred-proxy-key"); - private final Map inferredProxy; - - public static InferredProxyContext fromContext(Context context) { - return context.get(CONTEXT_KEY); - } - - public InferredProxyContext(Map contextInfo) { - this.inferredProxy = - (contextInfo == null || contextInfo.isEmpty()) - ? new HashMap<>() - : new HashMap<>(contextInfo); - } - - public InferredProxyContext() { - this.inferredProxy = new HashMap<>(); - } - - public Map getInferredProxyContext() { - return Collections.unmodifiableMap(inferredProxy); - } - - public void putInferredProxyInfo(String key, String value) { - inferredProxy.put(key, value); - } - - public void removeInferredProxyInfo(String key) { - inferredProxy.remove(key); - } - - /** - * Creates a new context with this value under its chosen key. - * - * @param context the context to copy the original values from. - * @return the new context with the implicitly keyed value. - * @see Context#with(ImplicitContextKeyed) - */ - @Override - public Context storeInto(Context context) { - return context.with(CONTEXT_KEY, this); - } -} diff --git a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java b/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java deleted file mode 100644 index 69e5a0e896e..00000000000 --- a/components/context/src/main/java/datadog/context/propagation/InferredProxyPropagator.java +++ /dev/null @@ -1,74 +0,0 @@ -package datadog.context.propagation; - -import datadog.context.Context; -import datadog.context.InferredProxyContext; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; - -public class InferredProxyPropagator implements Propagator { - public static final String INFERRED_PROXY_KEY = "x-dd-proxy"; - /** - * METHOD STUB: InferredProxy is currently not meant to be injected to downstream services Injects - * a context into a downstream service using the given carrier. - * - * @param context the context containing the values to be injected. - * @param carrier the instance that will receive the key/value pairs to propagate. - * @param setter the callback to set key/value pairs into the carrier. - */ - @Override - public void inject(Context context, C carrier, CarrierSetter setter) {} - - /** - * Extracts a context from un upstream service. - * - * @param context the base context to store the extracted values on top, use {@link - * Context#root()} for a default base context. - * @param carrier the instance to fetch the propagated key/value pairs from. - * @param visitor the callback to walk over the carrier and extract its key/value pais. - * @return A context with the extracted values on top of the given base context. - */ - @Override - public Context extract(Context context, C carrier, CarrierVisitor visitor) { - if (context == null || carrier == null || visitor == null) { - return context; - } - InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); - visitor.forEachKeyValue(carrier, extractor); - - InferredProxyContext extractedContext = extractor.extractedContext; - if (extractedContext == null) { - return context; - } - return extractedContext.storeInto(context); - } - - public static class InferredProxyContextExtractor implements BiConsumer { - private InferredProxyContext extractedContext; - - InferredProxyContextExtractor() {} - - private Map parseInferredProxyHeaders(String input) { - Map parsedHeaders = new HashMap<>(); - return parsedHeaders; - } - - /** - * Performs this operation on the given arguments. - * - * @param key the first input argument from an http header - * @param value the second input argument from an http header - */ - @Override - public void accept(String key, String value) { - if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY)) { - return; - } - Map inferredProxyMap = parseInferredProxyHeaders(value); - if (extractedContext == null) { - extractedContext = new InferredProxyContext(); - } - extractedContext.putInferredProxyInfo(key, value); - } - } -} diff --git a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java b/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java deleted file mode 100644 index 53ddf5cb12a..00000000000 --- a/components/context/src/test/java/datadog/context/InferredProxyHandlingTest.java +++ /dev/null @@ -1,465 +0,0 @@ -package datadog.context; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import datadog.context.propagation.CarrierVisitor; -import datadog.context.propagation.InferredProxyPropagator; -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; -import java.util.stream.Stream; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Nested; -import org.junit.jupiter.api.Test; // For @Test on nested class methods -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; - -class InferredProxyHandlingTest { - - // Define header key constants locally for the test - static final String PROXY_SYSTEM_KEY = "x-dd-proxy-system"; - static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; - static final String PROXY_PATH_KEY = "x-dd-proxy-path"; - static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; - static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; - - private InferredProxyPropagator propagator; - - @BeforeEach - void setUp() { - propagator = new InferredProxyPropagator(); - } - - // Moved @MethodSource providers to the outer class and made them static - static Stream validHeadersProviderForPropagator() { - Map allStandard = new HashMap<>(); - allStandard.put(PROXY_SYSTEM_KEY, "aws-apigw"); // The only currently supported system - allStandard.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - allStandard.put(PROXY_PATH_KEY, "/foo"); - allStandard.put(PROXY_HTTP_METHOD_KEY, "GET"); - allStandard.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); - - return Stream.of( - Arguments.of( - "all standard headers (aws-apigw)", - allStandard, - "aws-apigw", - "12345", - "/foo", - "GET", - "api.example.com", - null, - null)); - } - - static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed - Map missingSystem = new HashMap<>(); - missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - missingSystem.put(PROXY_PATH_KEY, "/foo"); - - Map missingTime = new HashMap<>(); - missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); - missingTime.put(PROXY_PATH_KEY, "/foo"); - - return Stream.of( - Arguments.of("PROXY_SYSTEM_KEY missing", missingSystem), - Arguments.of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); - } - - // Simple Map visitor for tests (can remain static or non-static in outer class) - static class MapVisitor implements CarrierVisitor> { - @Override - public void forEachKeyValue(Map carrier, BiConsumer visitor) { - if (carrier == null) { - return; - } - carrier.forEach(visitor); - } - } - - // Custom visitor to test null key path in the extractor - MOVED HERE and made static - static class NullKeyTestVisitor implements CarrierVisitor> { - private final BiConsumer actualExtractorAccept; - - NullKeyTestVisitor(BiConsumer actualExtractorAccept) { - this.actualExtractorAccept = actualExtractorAccept; - } - - @Override - public void forEachKeyValue(Map carrier, BiConsumer visitor) { - if (actualExtractorAccept != null) { - actualExtractorAccept.accept(null, "valueForNullKey"); - } - } - } - - @Nested - @DisplayName("InferredProxyPropagator Tests") - class PropagatorTests { // Kept non-static - - @ParameterizedTest(name = "{0}") - @MethodSource( - "datadog.context.InferredProxyHandlingTest#validHeadersProviderForPropagator") // Fully - // qualified - // name - @DisplayName("Should extract InferredProxyContext when valid headers are present") - void testSuccessfulExtraction( - String description, - Map headers, - String expectedSystem, - String expectedTimeMs, - String expectedPath, - String expectedMethod, - String expectedDomain, - String expectedExtraKey, - String expectedExtraValue) { - - Context rootContext = Context.root(); - // Now accesses the outer class's propagator instance field - Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull( - inferredProxyContext, "InferredProxyContext should not be null for: " + description); - Map actualProxyData = inferredProxyContext.getInferredProxyContext(); - assertEquals(expectedSystem, actualProxyData.get(PROXY_SYSTEM_KEY)); - assertEquals(expectedTimeMs, actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - assertEquals(expectedPath, actualProxyData.get(PROXY_PATH_KEY)); - assertEquals(expectedMethod, actualProxyData.get(PROXY_HTTP_METHOD_KEY)); - assertEquals(expectedDomain, actualProxyData.get(PROXY_DOMAIN_NAME_KEY)); - if (expectedExtraKey != null) { - assertEquals(expectedExtraValue, actualProxyData.get(expectedExtraKey)); - } - } - - @ParameterizedTest(name = "{0}") - @MethodSource( - "datadog.context.InferredProxyHandlingTest#invalidOrMissingHeadersProviderForPropagator") // Fully qualified name - @DisplayName("Should create InferredProxyContext even if some critical headers are missing") - void testExtractionWithMissingCriticalHeaders(String description, Map headers) { - Context rootContext = Context.root(); - Context extractedOuterContext = propagator.extract(rootContext, headers, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull( - inferredProxyContext, - "InferredProxyContext should still be created if any x-dd-proxy-* header is present for: " - + description); - Map actualProxyData = inferredProxyContext.getInferredProxyContext(); - - if (headers.containsKey(PROXY_SYSTEM_KEY)) { - assertEquals(headers.get(PROXY_SYSTEM_KEY), actualProxyData.get(PROXY_SYSTEM_KEY)); - } else { - assertNull(actualProxyData.get(PROXY_SYSTEM_KEY)); - } - if (headers.containsKey(PROXY_REQUEST_TIME_MS_KEY)) { - assertEquals( - headers.get(PROXY_REQUEST_TIME_MS_KEY), actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - } else { - assertNull(actualProxyData.get(PROXY_REQUEST_TIME_MS_KEY)); - } - } - - @Test - @DisplayName("Should not extract InferredProxyContext if no relevant headers are present") - void testNoRelevantHeaders() { - Map carrier = new HashMap<>(); - carrier.put("x-unrelated-header", "value"); - carrier.put("another-header", "othervalue"); - Context rootContext = Context.root(); - - Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNull( - inferredProxyContext, - "InferredProxyContext should be null if no x-dd-proxy-* headers are found"); - } - - @Test - @DisplayName("Should return original context if carrier is null") - void testNullCarrier() { - InferredProxyContext initialData = - new InferredProxyContext(Collections.singletonMap("test", "value")); - Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); - - Context extractedOuterContext = propagator.extract(rootContext, null, new MapVisitor()); - - assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); - assertEquals( - "value", - InferredProxyContext.fromContext(extractedOuterContext) - .getInferredProxyContext() - .get("test")); - } - - @Test - @DisplayName("Should return original context if visitor is null") - void testNullVisitor() { - Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); - InferredProxyContext initialData = - new InferredProxyContext(Collections.singletonMap("test", "value")); - Context rootContext = Context.root().with(InferredProxyContext.CONTEXT_KEY, initialData); - - Context extractedOuterContext = propagator.extract(rootContext, carrier, null); - - assertEquals(rootContext, extractedOuterContext, "Context should be unchanged"); - assertEquals( - "value", - InferredProxyContext.fromContext(extractedOuterContext) - .getInferredProxyContext() - .get("test")); - } - - @Test - @DisplayName("Should return original context if context is null") - void testNullContext() { - Map carrier = Collections.singletonMap(PROXY_SYSTEM_KEY, "aws-apigw"); - Context extractedOuterContext = propagator.extract(null, carrier, new MapVisitor()); - assertNull(extractedOuterContext, "Context should remain null if passed as null"); - } - - @Test - @DisplayName("Extractor should handle multiple proxy headers") - void testMultipleProxyHeaders() { - Map carrier = new HashMap<>(); - carrier.put(PROXY_SYSTEM_KEY, "aws-apigw"); - carrier.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - carrier.put("x-dd-proxy-custom", "value1"); // First proxy header - carrier.put("x-dd-proxy-another", "value2"); // Second proxy header - - Context rootContext = Context.root(); - Context extractedOuterContext = propagator.extract(rootContext, carrier, new MapVisitor()); - InferredProxyContext inferredProxyContext = - InferredProxyContext.fromContext(extractedOuterContext); - - assertNotNull(inferredProxyContext); - // Check if both headers were stored (covers extractedContext == null being false) - assertEquals( - "value1", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-custom")); - assertEquals( - "value2", inferredProxyContext.getInferredProxyContext().get("x-dd-proxy-another")); - assertEquals( - "aws-apigw", inferredProxyContext.getInferredProxyContext().get(PROXY_SYSTEM_KEY)); - } - - @Test - @DisplayName("Extractor accept method should handle null/empty keys") - void testExtractorAcceptNullEmptyKeys() { - Context rootContext = Context.root(); - - // Test null key - HashMap doesn't allow null keys. Standard HTTP visitors - // also typically don't yield null keys. Testing this branch is difficult - // without a custom visitor or modifying the source. Relying on coverage report - // or assuming standard carriers won't provide null keys. - - // Test empty key - Map carrierWithEmptyKey = new HashMap<>(); - carrierWithEmptyKey.put("", "emptyKeyValue"); // Add empty key - carrierWithEmptyKey.put(PROXY_SYSTEM_KEY, "aws-apigw"); // Add a valid key too - - Context contextAfterEmpty = - propagator.extract(rootContext, carrierWithEmptyKey, new MapVisitor()); - InferredProxyContext ipcEmpty = InferredProxyContext.fromContext(contextAfterEmpty); - - // The propagator should ignore the empty key entry entirely. - assertNotNull(ipcEmpty, "Context should be created due to valid key"); - assertNull(ipcEmpty.getInferredProxyContext().get(""), "Empty key should not be stored"); - assertEquals( - "aws-apigw", - ipcEmpty.getInferredProxyContext().get(PROXY_SYSTEM_KEY), - "Valid key should still be stored"); - assertEquals(1, ipcEmpty.getInferredProxyContext().size(), "Only valid key should be stored"); - } - - @Test - @DisplayName( - "Extractor accept method should handle explicitly passed null key via custom visitor") - void testExtractorAcceptExplicitNullKey() { - Context rootContext = Context.root(); - Map carrier = new HashMap<>(); // Carrier can be empty for this test - - // We need to get a handle to the internal BiConsumer (the InferredProxyContextExtractor - // instance). - // The extract method will create one. We can pass a visitor that captures it. - - final BiConsumer[] extractorHolder = new BiConsumer[1]; - - CarrierVisitor> capturingVisitor = - (cr, bic) -> { - extractorHolder[0] = bic; // Capture the BiConsumer - // Optionally, call the original MapVisitor if we still want normal processing after - // capture - // new MapVisitor().forEachKeyValue(cr, bic); - }; - - // This first call is primarily to get a reference to the internal extractor - propagator.extract(rootContext, carrier, capturingVisitor); - - assertNotNull(extractorHolder[0], "Failed to capture the internal extractor instance"); - - // Now use a new custom visitor to specifically test the null key path - // on the captured extractor instance (though this isn't how extract is typically used). - // A more direct way to test the BiConsumer if it were accessible or if the design allowed it. - // For now, we directly call accept on the captured one. - extractorHolder[0].accept(null, "valueForNullKey"); - - // The goal is JaCoCo coverage. Asserting internal state of the extractor is hard without - // reflection. - // We can verify that the context remains unchanged or as expected if no valid headers - // processed. - InferredProxyContext ipc = - InferredProxyContext.fromContext( - rootContext); // or context returned by a second extract call - assertNull(ipc, "Context should not have InferredProxyContext from only a null key call"); - } - } - - @Nested - @DisplayName("InferredProxyContext Tests") - class ContextUnitTests { - - @Test - @DisplayName("Default constructor should create an empty context map") - void testDefaultConstructor() { - InferredProxyContext ipc = new InferredProxyContext(); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - - @Test - @DisplayName("Constructor with map should initialize context map") - void testMapConstructor() { - Map initialData = new HashMap<>(); - initialData.put("key1", "value1"); - initialData.put("key2", "value2"); - - InferredProxyContext ipc = new InferredProxyContext(initialData); - assertNotNull(ipc.getInferredProxyContext()); - assertEquals(2, ipc.getInferredProxyContext().size()); - assertEquals("value1", ipc.getInferredProxyContext().get("key1")); - assertEquals("value2", ipc.getInferredProxyContext().get("key2")); - - initialData.put("key3", "value3"); // Modify original map - assertNull(ipc.getInferredProxyContext().get("key3"), "Internal map should be a copy"); - } - - @Test - @DisplayName("putInferredProxyInfo should add to the context map") - void testPutInfo() { - InferredProxyContext ipc = new InferredProxyContext(); - ipc.putInferredProxyInfo("system", "aws-apigw"); - ipc.putInferredProxyInfo("time", "12345"); - - Map contextMap = ipc.getInferredProxyContext(); - assertEquals(2, contextMap.size()); - assertEquals("aws-apigw", contextMap.get("system")); - assertEquals("12345", contextMap.get("time")); - - ipc.putInferredProxyInfo("system", "azure-func"); // Overwrite - assertEquals("azure-func", contextMap.get("system")); - assertEquals(2, contextMap.size()); - } - - @Test - @DisplayName("removeInferredProxyInfo should remove from the context map") - void testRemoveInfo() { - Map initialData = new HashMap<>(); - initialData.put("key1", "value1"); - initialData.put("key2", "value2"); - InferredProxyContext ipc = new InferredProxyContext(initialData); - - ipc.removeInferredProxyInfo("key1"); - Map contextMap = ipc.getInferredProxyContext(); - assertEquals(1, contextMap.size()); - assertNull(contextMap.get("key1")); - assertEquals("value2", contextMap.get("key2")); - - ipc.removeInferredProxyInfo("nonexistent"); // Remove non-existent - assertEquals(1, contextMap.size()); - } - - @Test - @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") - void testStoreAndFromContext() { - InferredProxyContext ipcToStore = new InferredProxyContext(); - ipcToStore.putInferredProxyInfo("customKey", "customValue"); - - Context rootContext = Context.root(); - Context contextWithValue = ipcToStore.storeInto(rootContext); - assertNotNull(contextWithValue); - - InferredProxyContext retrievedIpc = InferredProxyContext.fromContext(contextWithValue); - assertNotNull(retrievedIpc); - assertEquals("customValue", retrievedIpc.getInferredProxyContext().get("customKey")); - - assertNull( - InferredProxyContext.fromContext(rootContext), - "Original root context should not be affected"); - - Context cleanContext = Context.root(); - assertNull( - InferredProxyContext.fromContext(cleanContext), - "fromContext on clean context should be null"); - } - - @Test - @DisplayName("getInferredProxyContext should return an unmodifiable map or a copy") - void testGetInferredProxyContextImmutability() { - InferredProxyContext ipc = new InferredProxyContext(); - ipc.putInferredProxyInfo("key1", "value1"); - - Map retrievedMap = ipc.getInferredProxyContext(); - assertNotNull(retrievedMap); - assertEquals("value1", retrievedMap.get("key1")); - - boolean threwUnsupported = false; - try { - retrievedMap.put("newKey", "newValue"); - } catch (UnsupportedOperationException e) { - threwUnsupported = true; - } - // Depending on whether InferredProxyContext.getInferredProxyContext() returns a direct - // reference or a copy, - // this assertion might change. If it returns a direct mutable reference, threwUnsupported - // will be false. - // If it returns an unmodifiable view or a copy, attempts to modify might throw or simply not - // affect the original. - // For now, we check that the original context was not changed. - assertEquals( - 1, ipc.getInferredProxyContext().size(), "Internal map size should remain unchanged"); - assertEquals( - "value1", - ipc.getInferredProxyContext().get("key1"), - "Internal map content should remain unchanged"); - // If it MUST be unmodifiable, add: assertTrue(threwUnsupported, "Retrieved map should be - // unmodifiable"); - } - - @Test - @DisplayName("Constructor with null map should create an empty context map") - void testNullMapConstructor() { - InferredProxyContext ipc = new InferredProxyContext(null); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - - @Test - @DisplayName("Constructor with empty map should create an empty context map") - void testEmptyMapConstructor() { - Map emptyMap = Collections.emptyMap(); - InferredProxyContext ipc = new InferredProxyContext(emptyMap); - assertNotNull(ipc.getInferredProxyContext()); - assertTrue(ipc.getInferredProxyContext().isEmpty()); - } - } -} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 5d8a7ee5974..91251be57f9 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -11,12 +11,8 @@ import datadog.appsec.api.blocking.BlockingException; import datadog.context.Context; import datadog.context.propagation.Propagators; -import datadog.context.InferredProxyContext; -import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.DDTags; -import datadog.trace.api.DDTraceId; -import datadog.trace.api.TraceConfig; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; import datadog.trace.api.gateway.BlockResponseFunction; @@ -24,14 +20,13 @@ import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction; import datadog.trace.api.gateway.IGSpanInfo; +import datadog.trace.api.gateway.InferredProxySpan; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; -import datadog.trace.api.interceptor.MutableSpan; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentPropagation; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; -import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -44,8 +39,6 @@ import datadog.trace.bootstrap.instrumentation.decorator.http.ClientIpAddressResolver; import java.net.InetAddress; import java.util.BitSet; -import java.util.HashMap; -import java.util.LinkedHashMap; import java.util.Locale; import java.util.Map; import java.util.concurrent.ExecutionException; @@ -60,401 +53,9 @@ public abstract class HttpServerDecorator extends ServerDecorator { - class InferredProxySpanGroup implements AgentSpan { - private final AgentSpan inferredProxySpan; - private final AgentSpan serverSpan; - - InferredProxySpanGroup(AgentSpan inferredProxySpan, AgentSpan serverSpan) { - this.inferredProxySpan = inferredProxySpan; - this.serverSpan = serverSpan; - } - - @Override - public DDTraceId getTraceId() { - return serverSpan.getTraceId(); - } - - @Override - public long getSpanId() { - return serverSpan.getSpanId(); - } - - @Override - public AgentSpan setTag(String key, boolean value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, int value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, long value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, double value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, String value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, CharSequence value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setTag(String key, Object value) { - return serverSpan.setTag(key, value); - } - - /** - * @param map - * @return - */ - @Override - public AgentSpan setAllTags(Map map) { - return null; - } - - @Override - public AgentSpan setTag(String key, Number value) { - return serverSpan.setTag(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, int value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, long value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setMetric(CharSequence key, double value) { - return serverSpan.setMetric(key, value); - } - - @Override - public AgentSpan setSpanType(CharSequence type) { - return serverSpan.setSpanType(type); - } - - @Override - public Object getTag(String key) { - return serverSpan.getTag(key); - } - - @Override - public AgentSpan setError(boolean error) { - serverSpan.setError(error); - if (inferredProxySpan != null) { - inferredProxySpan.setError(error); - } - return this; - } - - @Override - public AgentSpan setError(boolean error, byte priority) { - serverSpan.setError(error, priority); - if (inferredProxySpan != null) { - inferredProxySpan.setError(error, priority); - } - return this; - } - - @Override - public AgentSpan setMeasured(boolean measured) { - return serverSpan.setMeasured(measured); - } - - @Override - public AgentSpan setErrorMessage(String errorMessage) { - return serverSpan.setErrorMessage(errorMessage); - } - - @Override - public AgentSpan addThrowable(Throwable throwable) { - serverSpan.addThrowable(throwable); - if (inferredProxySpan != null) { - inferredProxySpan.addThrowable(throwable); - } - return this; - } - - @Override - public AgentSpan addThrowable(Throwable throwable, byte errorPriority) { - serverSpan.addThrowable(throwable, errorPriority); - if (inferredProxySpan != null) { - inferredProxySpan.addThrowable(throwable, errorPriority); - } - return this; - } - - @Override - public AgentSpan getLocalRootSpan() { - return serverSpan.getLocalRootSpan(); - } - - @Override - public boolean isSameTrace(AgentSpan otherSpan) { - return serverSpan.isSameTrace(otherSpan); - } - - @Override - public AgentSpanContext context() { - return serverSpan.context(); - } - - @Override - public String getBaggageItem(String key) { - return serverSpan.getBaggageItem(key); - } - - @Override - public AgentSpan setBaggageItem(String key, String value) { - return serverSpan.setBaggageItem(key, value); - } - - @Override - public AgentSpan setHttpStatusCode(int statusCode) { - serverSpan.setHttpStatusCode(statusCode); - if (inferredProxySpan != null) { - inferredProxySpan.setHttpStatusCode(statusCode); - } - return this; - } - - @Override - public short getHttpStatusCode() { - return serverSpan.getHttpStatusCode(); - } - - @Override - public void finish() { - serverSpan.finish(); - if (inferredProxySpan != null) { - inferredProxySpan.finish(); - } - } - - @Override - public void finish(long finishMicros) { - serverSpan.finish(finishMicros); - if (inferredProxySpan != null) { - inferredProxySpan.finish(finishMicros); - } - } - - @Override - public void finishWithDuration(long durationNanos) { - serverSpan.finishWithDuration(durationNanos); - if (inferredProxySpan != null) { - inferredProxySpan.finishWithDuration(durationNanos); - } - } - - @Override - public void beginEndToEnd() { - serverSpan.beginEndToEnd(); - } - - @Override - public void finishWithEndToEnd() { - serverSpan.finishWithEndToEnd(); - if (inferredProxySpan != null) { - inferredProxySpan.finishWithEndToEnd(); - } - } - - @Override - public boolean phasedFinish() { - final boolean ret = serverSpan.phasedFinish(); - if (inferredProxySpan != null) { - inferredProxySpan.phasedFinish(); - } - return ret; - } - - @Override - public void publish() { - serverSpan.publish(); - } - - @Override - public CharSequence getSpanName() { - return serverSpan.getSpanName(); - } - - @Override - public void setSpanName(CharSequence spanName) { - serverSpan.setSpanName(spanName); - } - - @Deprecated - @Override - public boolean hasResourceName() { - return serverSpan.hasResourceName(); - } - - @Override - public byte getResourceNamePriority() { - return serverSpan.getResourceNamePriority(); - } - - @Override - public AgentSpan setResourceName(CharSequence resourceName) { - return serverSpan.setResourceName(resourceName); - } - - @Override - public AgentSpan setResourceName(CharSequence resourceName, byte priority) { - return serverSpan.setResourceName(resourceName, priority); - } - - @Override - public RequestContext getRequestContext() { - return serverSpan.getRequestContext(); - } - - @Override - public Integer forceSamplingDecision() { - return serverSpan.forceSamplingDecision(); - } - - @Override - public AgentSpan setSamplingPriority(int newPriority, int samplingMechanism) { - return serverSpan.setSamplingPriority(newPriority, samplingMechanism); - } - - @Override - public TraceConfig traceConfig() { - return serverSpan.traceConfig(); - } - - @Override - public void addLink(AgentSpanLink link) { - serverSpan.addLink(link); - } - - @Override - public AgentSpan setMetaStruct(String field, Object value) { - return serverSpan.setMetaStruct(field, value); - } - - @Override - public boolean isOutbound() { - return serverSpan.isOutbound(); - } - - @Override - public AgentSpan asAgentSpan() { - return serverSpan.asAgentSpan(); - } - - @Override - public long getStartTime() { - return serverSpan.getStartTime(); - } - - @Override - public long getDurationNano() { - return serverSpan.getDurationNano(); - } - - @Override - public CharSequence getOperationName() { - return serverSpan.getOperationName(); - } - - @Override - public MutableSpan setOperationName(CharSequence serviceName) { - return serverSpan.setOperationName(serviceName); - } - - @Override - public String getServiceName() { - return serverSpan.getServiceName(); - } - - @Override - public MutableSpan setServiceName(String serviceName) { - return serverSpan.setServiceName(serviceName); - } - - @Override - public CharSequence getResourceName() { - return serverSpan.getResourceName(); - } - - @Override - public Integer getSamplingPriority() { - return serverSpan.getSamplingPriority(); - } - - @Deprecated - @Override - public MutableSpan setSamplingPriority(int newPriority) { - return serverSpan.setSamplingPriority(newPriority); - } - - @Override - public String getSpanType() { - return serverSpan.getSpanType(); - } - - @Override - public Map getTags() { - return serverSpan.getTags(); - } - - @Override - public boolean isError() { - return serverSpan.isError(); - } - - @Deprecated - @Override - public MutableSpan getRootSpan() { - return serverSpan.getRootSpan(); - } - - @Override - public void setRequestBlockingAction(Flow.Action.RequestBlockingAction rba) { - serverSpan.setRequestBlockingAction(rba); - } - - @Override - public Flow.Action.RequestBlockingAction getRequestBlockingAction() { - return serverSpan.getRequestBlockingAction(); - } - } - private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class); private static final int UNSET_PORT = 0; - public static final String PROXY_SYSTEM = "x-dd-proxy"; - public static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; - public static final String PROXY_PATH = "x-dd-proxy-path"; - public static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; - public static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; - public static final String STAGE = "x-dd-proxy-stage"; - - public static final Map SUPPORTED_PROXIES; - - static { - SUPPORTED_PROXIES = new HashMap<>(); - SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); - } - public static final String DD_SPAN_ATTRIBUTE = "datadog.span"; public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch"; public static final String DD_RUM_INJECTED = "datadog.rum.injected"; @@ -547,69 +148,30 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { instrumentationNames != null && instrumentationNames.length > 0 ? instrumentationNames[0] : DEFAULT_INSTRUMENTATION_NAME; - AgentSpanContext.Extracted extracted = callIGCallbackStart(getExtractedSpanContext(context)); + AgentSpanContext extracted = getExtractedSpanContext(context); + // Call IG callbacks + extracted = callIGCallbackStart(extracted); + // Create gateway inferred span if needed + extracted = startInferredProxySpan(context, extracted); AgentSpan span = tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true); + // Apply RequestBlockingAction if any Flow flow = callIGCallbackRequestHeaders(span, carrier); if (flow.getAction() instanceof RequestBlockingAction) { span.setRequestBlockingAction((RequestBlockingAction) flow.getAction()); } - AgentPropagation.ContextVisitor getter = getter(); - if (null != carrier && null != getter) { - tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); - } + // DSM Checkpoint + tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer()); return context.with(span); } - private AgentSpan startSpanWithInferredProxy( - String instrumentationName, - datadog.context.Context fullContextForInferredProxy, - AgentSpanContext.Extracted standardExtractedContext) { - - InferredProxyContext inferredProxy = - InferredProxyContext.fromContext(fullContextForInferredProxy); - - if (inferredProxy == null) { - return null; - } - - Map headers = inferredProxy.getInferredProxyContext(); - - // Check if timestamp and proxy system are present - String startTimeStr = headers.get(PROXY_START_TIME_MS); - String proxySystem = headers.get(PROXY_SYSTEM); - - if (startTimeStr == null - || proxySystem == null - || !SUPPORTED_PROXIES.containsKey(proxySystem)) { - return null; - } - - long startTime; - try { - startTime = Long.parseLong(startTimeStr) * 1000; // Convert to microseconds - } catch (NumberFormatException e) { - return null; // Invalid timestamp + protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) { + InferredProxySpan span; + if (!Config.get().isInferredProxyPropagationEnabled() + || (span = InferredProxySpan.fromContext(context)) == null) { + return extracted; } - - AgentSpan apiGtwSpan = - tracer() - .startSpan( - "inferred_proxy", - SUPPORTED_PROXIES.get(proxySystem), - callIGCallbackStart(standardExtractedContext), - startTime); - - apiGtwSpan.setTag(Tags.COMPONENT, proxySystem); - apiGtwSpan.setTag( - DDTags.RESOURCE_NAME, headers.get(PROXY_HTTP_METHOD) + " " + headers.get(PROXY_PATH)); - apiGtwSpan.setTag(DDTags.SERVICE_NAME, headers.get(PROXY_DOMAIN_NAME)); - apiGtwSpan.setTag(DDTags.SPAN_TYPE, "web"); - apiGtwSpan.setTag(Tags.HTTP_METHOD, headers.get(PROXY_HTTP_METHOD)); - apiGtwSpan.setTag(Tags.HTTP_URL, headers.get(PROXY_DOMAIN_NAME) + headers.get(PROXY_PATH)); - apiGtwSpan.setTag("stage", headers.get(STAGE)); - apiGtwSpan.setTag("_dd.inferred_span", 1); - return apiGtwSpan; + return span.start(extracted); } public AgentSpan onRequest( @@ -832,8 +394,7 @@ public AgentSpan onResponse(final AgentSpan span, final RESPONSE response) { return span; } - private AgentSpanContext.Extracted callIGCallbackStart( - @Nullable final AgentSpanContext.Extracted extracted) { + private AgentSpanContext callIGCallbackStart(@Nullable final AgentSpanContext extracted) { AgentTracer.TracerAPI tracer = tracer(); Supplier> startedCbAppSec = tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted()); @@ -969,10 +530,20 @@ private Flow callIGCallbackURI( @Override public AgentSpan beforeFinish(AgentSpan span) { + // TODO Migrate beforeFinish to Context API onRequestEndForInstrumentationGateway(span); + // Close Serverless Gateway Inferred Span if any + // finishInferredProxySpan(context); return super.beforeFinish(span); } + protected void finishInferredProxySpan(Context context) { + InferredProxySpan span; + if ((span = InferredProxySpan.fromContext(context)) != null) { + span.finish(); + } + } + private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) { if (span.getLocalRootSpan() != span) { return; 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 56aeedc6e71..71cf4fe76c4 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 @@ -23,7 +23,6 @@ import datadog.communication.ddagent.SharedCommunicationObjects; import datadog.communication.monitor.Monitoring; import datadog.communication.monitor.Recording; -import datadog.context.propagation.InferredProxyPropagator; import datadog.context.propagation.Propagators; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; @@ -93,6 +92,7 @@ import datadog.trace.core.monitor.TracerHealthMetrics; import datadog.trace.core.propagation.ExtractedContext; import datadog.trace.core.propagation.HttpCodec; +import datadog.trace.core.propagation.InferredProxyPropagator; import datadog.trace.core.propagation.PropagationTags; import datadog.trace.core.propagation.TracingPropagator; import datadog.trace.core.propagation.XRayPropagator; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java new file mode 100644 index 00000000000..da51bf4643b --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java @@ -0,0 +1,60 @@ +package datadog.trace.core.propagation; + +import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; + +import datadog.context.Context; +import datadog.context.propagation.CarrierSetter; +import datadog.context.propagation.CarrierVisitor; +import datadog.context.propagation.Propagator; +import datadog.trace.api.gateway.InferredProxySpan; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import javax.annotation.ParametersAreNonnullByDefault; + +/** Inferred proxy propagator. Only extract, not meant for injection. */ +@ParametersAreNonnullByDefault +public class InferredProxyPropagator implements Propagator { + private static final String INFERRED_PROXY_KEY_PREFIX = "x-dd-proxy"; + + @Override + public void inject(Context context, C carrier, CarrierSetter setter) {} + + @Override + public Context extract(Context context, C carrier, CarrierVisitor visitor) { + if (context == null || carrier == null || visitor == null) { + return context; + } + InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); + visitor.forEachKeyValue(carrier, extractor); + InferredProxySpan inferredProxySpan = extractor.inferredProxySpan(); + if (inferredProxySpan != null) { + context = context.with(inferredProxySpan); + } + return context; + } + + /** Extract inferred proxy related headers into a map. */ + private static class InferredProxyContextExtractor implements BiConsumer { + private Map values; + + @Override + public void accept(String key, String value) { + if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY_PREFIX)) { + return; + } + if (values == null) { + this.values = new HashMap<>(); + } + this.values.put(key, value); + } + + public InferredProxySpan inferredProxySpan() { + if (this.values == null) { + return null; + } + InferredProxySpan inferredProxySpan = fromHeaders(this.values); + return inferredProxySpan.isValid() ? inferredProxySpan : null; + } + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java new file mode 100644 index 00000000000..5830a160cc7 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java @@ -0,0 +1,96 @@ +package datadog.trace.core.propagation; + +import static datadog.context.Context.root; +import static datadog.trace.api.gateway.InferredProxySpan.fromContext; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.params.provider.Arguments.of; + +import datadog.context.Context; +import datadog.context.propagation.CarrierVisitor; +import datadog.trace.api.gateway.InferredProxySpan; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.stream.Stream; +import javax.annotation.ParametersAreNonnullByDefault; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +@DisplayName("InferredProxyPropagator Tests") +class InferredProxyPropagatorTests { + private static final String PROXY_SYSTEM_KEY = "x-dd-proxy"; + private static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; + private static final String PROXY_PATH_KEY = "x-dd-proxy-path"; + private static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; + private static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; + private static final MapVisitor MAP_VISITOR = new MapVisitor(); + + private InferredProxyPropagator propagator; + + @BeforeEach + void setUp() { + this.propagator = new InferredProxyPropagator(); + } + + @Test + @DisplayName("Should extract InferredProxySpan when valid headers are present") + void testSuccessfulExtraction() { + Map headers = new HashMap<>(); + headers.put(PROXY_SYSTEM_KEY, "aws-apigateway"); + headers.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + headers.put(PROXY_PATH_KEY, "/foo"); + headers.put(PROXY_HTTP_METHOD_KEY, "GET"); + headers.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); + + Context context = this.propagator.extract(root(), headers, MAP_VISITOR); + InferredProxySpan inferredProxySpan = fromContext(context); + assertNotNull(inferredProxySpan); + assertTrue(inferredProxySpan.isValid()); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("invalidOrMissingHeadersProviderForPropagator") + @DisplayName("Should not create InferredProxySpan if some critical headers are missing") + void testExtractionWithMissingCriticalHeaders(String description, Map headers) { + Context rootContext = root(); + Context extractedOuterContext = this.propagator.extract(rootContext, headers, MAP_VISITOR); + InferredProxySpan inferredProxySpan = fromContext(extractedOuterContext); + assertNull(inferredProxySpan, "Invalid inferred proxy span should not be extracted"); + } + + static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); + missingSystem.put(PROXY_PATH_KEY, "/foo"); + + Map emptyValue = new HashMap<>(); + emptyValue.put(PROXY_SYSTEM_KEY, ""); + + Map nullValue = new HashMap<>(); + nullValue.put(PROXY_SYSTEM_KEY, null); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); + missingTime.put(PROXY_PATH_KEY, "/foo"); + + return Stream.of( + of("PROXY_SYSTEM_KEY missing", missingSystem), + of("PROXY_SYSTEM_KEY empty", emptyValue), + of("PROXY_SYSTEM_KEY null", nullValue), + of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); + } + + @ParametersAreNonnullByDefault + private static class MapVisitor implements CarrierVisitor> { + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + carrier.forEach(visitor); + } + } +} 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 a2e58598ea5..de78bd9dbd8 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -598,6 +598,7 @@ import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_ERROR_STATUSES; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_PATH_RESOURCE_NAME_MAPPING; +import static datadog.trace.api.config.TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_KEEP_LATENCY_THRESHOLD_MS; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_FLUSH_INTERVAL; @@ -5409,6 +5410,8 @@ public String toString() { + tracePropagationBehaviorExtract + ", tracePropagationExtractFirst=" + tracePropagationExtractFirst + + ", traceInferredProxyEnabled=" + + traceInferredProxyEnabled + ", clockSyncPeriod=" + clockSyncPeriod + ", jmxFetchEnabled=" diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java new file mode 100644 index 00000000000..a8256c1b38a --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java @@ -0,0 +1,108 @@ +package datadog.trace.api.gateway; + +import static datadog.context.ContextKey.named; +import static datadog.trace.api.DDTags.RESOURCE_NAME; +import static datadog.trace.api.DDTags.SERVICE_NAME; +import static datadog.trace.api.DDTags.SPAN_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_METHOD; +import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ImplicitContextKeyed; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class InferredProxySpan implements ImplicitContextKeyed { + private static final ContextKey CONTEXT_KEY = named("inferred-proxy-key"); + static final String PROXY_SYSTEM = "x-dd-proxy"; + static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; + static final String PROXY_PATH = "x-dd-proxy-path"; + static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; + static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; + static final String STAGE = "x-dd-proxy-stage"; + static final Map SUPPORTED_PROXIES; + static final String INSTRUMENTATION_NAME = "inferred_proxy"; + + static { + SUPPORTED_PROXIES = new HashMap<>(); + SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); + } + + private final Map headers; + private AgentSpan span; + + public static InferredProxySpan fromHeaders(Map values) { + return new InferredProxySpan(values); + } + + public static InferredProxySpan fromContext(Context context) { + return context.get(CONTEXT_KEY); + } + + private InferredProxySpan(Map headers) { + this.headers = headers == null ? Collections.emptyMap() : headers; + } + + public boolean isValid() { + String startTimeStr = header(PROXY_START_TIME_MS); + String proxySystem = header(PROXY_SYSTEM); + return startTimeStr != null + && proxySystem != null + && SUPPORTED_PROXIES.containsKey(proxySystem); + } + + public AgentSpanContext start(AgentSpanContext extracted) { + if (this.span != null || !isValid()) { + return extracted; + } + + long startTime; + try { + startTime = Long.parseLong(header(PROXY_START_TIME_MS)) * 1000; // Convert to microseconds + } catch (NumberFormatException e) { + return extracted; // Invalid timestamp + } + + String proxySystem = header(PROXY_SYSTEM); + String proxy = SUPPORTED_PROXIES.get(proxySystem); + AgentSpan span = AgentTracer.get().startSpan(INSTRUMENTATION_NAME, proxy, extracted, startTime); + + span.setTag(COMPONENT, proxySystem); + span.setTag(RESOURCE_NAME, header(PROXY_HTTP_METHOD) + " " + header(PROXY_PATH)); + span.setTag(SERVICE_NAME, header(PROXY_DOMAIN_NAME)); + span.setTag(SPAN_TYPE, "web"); + span.setTag(HTTP_METHOD, header(PROXY_HTTP_METHOD)); + span.setTag(HTTP_URL, header(PROXY_DOMAIN_NAME) + header(PROXY_PATH)); + span.setTag("stage", header(STAGE)); + span.setTag("_dd.inferred_span", 1); + + // Free collected headers + this.headers.clear(); + // Store inferred span + this.span = span; + // Return inferred span as new parent context + return this.span.context(); + } + + private String header(String name) { + return this.headers.get(name); + } + + public void finish() { + if (this.span != null) { + this.span.finish(); + this.span = null; + } + } + + @Override + public Context storeInto(Context context) { + return context.with(CONTEXT_KEY, this); + } +} diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java new file mode 100644 index 00000000000..dc61d4a23aa --- /dev/null +++ b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java @@ -0,0 +1,95 @@ +package datadog.trace.api.gateway; + +import static datadog.context.Context.root; +import static datadog.trace.api.gateway.InferredProxySpan.PROXY_START_TIME_MS; +import static datadog.trace.api.gateway.InferredProxySpan.PROXY_SYSTEM; +import static datadog.trace.api.gateway.InferredProxySpan.fromContext; +import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; +import static java.util.Collections.emptyMap; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.params.provider.Arguments.of; + +import datadog.context.Context; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Stream; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +@DisplayName("InferredProxyContext Tests") +class InferredProxySpanTests { + @Test + @DisplayName("Valid headers should make valid span") + void testMapConstructor() { + Map headers = new HashMap<>(); + headers.put(PROXY_START_TIME_MS, "12345"); + headers.put(PROXY_SYSTEM, "aws-apigateway"); + + InferredProxySpan inferredProxySpan = InferredProxySpan.fromHeaders(headers); + assertTrue(inferredProxySpan.isValid()); + assertNotNull( + inferredProxySpan.start(null), "inferred proxy span start and return new parent context"); + assertNull(inferredProxySpan.start(null), "inferred proxy span not should start twice"); + inferredProxySpan.finish(); + } + + @ParameterizedTest(name = "{0}") + @DisplayName("Invalid headers should make invalid span") + @MethodSource("invalidHeaders") + void testInvalidHeaders(String useCase, Map headers) { + InferredProxySpan inferredProxySpan = fromHeaders(headers); + assertFalse(inferredProxySpan.isValid(), useCase + " should not be valid"); + assertNull(inferredProxySpan.start(null), "Invalid inferred proxy span should not start"); + } + + static Stream invalidHeaders() { // Renamed + Map missingSystem = new HashMap<>(); + missingSystem.put(PROXY_START_TIME_MS, "12345"); + + Map missingTime = new HashMap<>(); + missingTime.put(PROXY_SYSTEM, "aws-apigateway"); + Map invalidSystem = new HashMap<>(); + invalidSystem.put(PROXY_START_TIME_MS, "12345"); + invalidSystem.put(PROXY_SYSTEM, "invalidSystem"); + + return Stream.of( + of("Missing system headers", missingSystem), + of("Missing start time headers", missingTime), + of("Invalid system headers", invalidSystem)); + } + + @Test + @DisplayName("Constructor with null should not crash") + void testNullMapConstructor() { + InferredProxySpan inferredProxySpan = fromHeaders(null); + assertNotNull(inferredProxySpan); + assertFalse(inferredProxySpan.isValid()); + } + + @Test + @DisplayName("Constructor with empty map should be invalid") + void testEmptyMapConstructor() { + InferredProxySpan inferredProxySpan = fromHeaders(emptyMap()); + assertNotNull(inferredProxySpan); + assertFalse(inferredProxySpan.isValid()); + } + + @Test + @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") + void testStoreAndFromContext() { + InferredProxySpan inferredProxySpan = fromHeaders(null); + Context context = inferredProxySpan.storeInto(root()); + assertNotNull(context); + + InferredProxySpan retrieved = fromContext(context); + assertNotNull(retrieved); + + assertNull(fromContext(root()), "fromContext on empty context should be null"); + } +} From 737e783d5606b497060d74451b6128c8912d9274 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Mon, 25 Aug 2025 22:37:52 -0400 Subject: [PATCH 03/27] init commit of instrumentations --- .../decorator/BaseDecorator.java | 7 ++++++- .../decorator/HttpServerDecorator.java | 14 ++++++++++++- .../akkahttp/DatadogAsyncHandlerWrapper.java | 8 ++++---- ...tadogServerRequestResponseFlowWrapper.java | 6 +++--- .../akkahttp/DatadogWrapperHelper.java | 10 ++++++---- .../axway/HTTPPluginAdvice.java | 2 +- .../instrumentation/axway/StateAdvice.java | 2 +- .../AzureFunctionsInstrumentation.java | 2 +- .../finatra/FinatraInstrumentation.java | 2 +- .../WriteFinalNettyResponseAdvice.java | 2 +- .../pekkohttp/DatadogAsyncHandlerWrapper.java | 6 +++--- ...tadogServerRequestResponseFlowWrapper.java | 7 +++---- .../pekkohttp/DatadogWrapperHelper.java | 10 ++++++---- .../instrumentation/play23/PlayAdvice.java | 4 ++-- .../play23/RequestCompleteCallback.java | 10 +++++++--- .../instrumentation/play24/PlayAdvice.java | 4 ++-- .../play24/RequestCompleteCallback.java | 10 +++++++--- .../instrumentation/play26/PlayAdvice.java | 4 ++-- .../play26/RequestCompleteCallback.java | 10 +++++++--- .../restlet/RestletInstrumentation.java | 2 +- .../servlet2/Servlet2Advice.java | 3 +-- .../servlet3/FinishAsyncDispatchListener.java | 20 +++++++++++-------- .../servlet3/Servlet3Advice.java | 6 +++--- .../instrumentation/spray/SprayHelper.scala | 2 +- .../DispatcherServletInstrumentation.java | 2 +- .../HandlerAdapterInstrumentation.java | 2 +- .../springweb6/ControllerAdvice.java | 2 +- .../springweb6/RenderAdvice.java | 2 +- 28 files changed, 98 insertions(+), 63 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java index bb575d92080..7472b660f54 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java @@ -4,6 +4,7 @@ import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT; import static datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver.hostName; +import datadog.context.Context; import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.api.DDTags; @@ -77,7 +78,7 @@ public AgentSpan afterStart(final AgentSpan span) { } public AgentScope beforeFinish(final AgentScope scope) { - beforeFinish(scope.span()); + beforeFinish(scope.context()); return scope; } @@ -85,6 +86,10 @@ public AgentSpan beforeFinish(final AgentSpan span) { return span; } + public Context beforeFinish(final Context context) { + return context; + } + public AgentScope onError(final AgentScope scope, final Throwable throwable) { onError(scope.span(), throwable); return scope; diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 91251be57f9..92b5ede5592 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -528,12 +528,24 @@ private Flow callIGCallbackURI( return Flow.ResultFlow.empty(); } + @Override + public Context beforeFinish(Context context) { + AgentSpan span = AgentSpan.fromContext(context); + onRequestEndForInstrumentationGateway(span); + + // Close Serverless Gateway Inferred Span if any + finishInferredProxySpan(context); + + return super.beforeFinish(context); + } + @Override public AgentSpan beforeFinish(AgentSpan span) { - // TODO Migrate beforeFinish to Context API onRequestEndForInstrumentationGateway(span); + // Close Serverless Gateway Inferred Span if any // finishInferredProxySpan(context); + return super.beforeFinish(span); } diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java index 195af7e84c9..324993dd75c 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java @@ -38,7 +38,7 @@ public Future apply(final HttpRequest request) { request.discardEntityBytes(materializer); HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request); span.getRequestContext().getTraceSegment().effectivelyBlocked(); - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); return FastFuture$.MODULE$.successful().apply(response); } @@ -46,7 +46,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope, t); throw t; } @@ -67,14 +67,14 @@ public HttpResponse apply(HttpResponse response) { response = newResponse; } - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope, t); return t; } }, diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java index 4becaf8fe2c..2a23d586ac0 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java @@ -78,7 +78,7 @@ public void onPush() throws Exception { skipNextPull[0] = true; requestContext.getTraceSegment().effectivelyBlocked(); emit(responseOutlet, response); - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); pull(requestInlet); scope.close(); return; @@ -142,7 +142,7 @@ public void onPush() throws Exception { response.discardEntityBytes(materializer()); response = newResponse; } - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); // Check if the active span matches the scope from when the request came in, // and close it. If it's not, then it will be cleaned up actor message // processing instrumentation that drives this state machine @@ -172,7 +172,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception { if (scope != null) { // Mark the span as failed AgentSpan span = fromContext(scope.context()); - DatadogWrapperHelper.finishSpan(span, ex); + DatadogWrapperHelper.finishSpan(scope, ex); } // We will not receive any more responses from the user code, so clean up any // remaining spans diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index 4c45b352c83..571936a6966 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java @@ -20,17 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) { return context.attach(); } - public static void finishSpan(final AgentSpan span, final HttpResponse response) { + public static void finishSpan(final ContextScope scope, final HttpResponse response) { + final AgentSpan span = fromContext(scope.context()); DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } - public static void finishSpan(final AgentSpan span, final Throwable t) { + public static void finishSpan(final ContextScope scope, final Throwable t) { + final AgentSpan span = fromContext(scope.context()); DECORATE.onError(span, t); span.setHttpStatusCode(500); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java index b698fbdedb9..d03f8e16dba 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java @@ -45,7 +45,7 @@ public static void onExit( if (throwable != null) { DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } finally { scope.close(); span.finish(); diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java index c9cc01ba869..fabff3d1346 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java @@ -38,7 +38,7 @@ public static void onExit( if (throwable != null) { DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } finally { scope.close(); span.finish(); diff --git a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java index 472f45bbc77..5edb0f4030e 100644 --- a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java +++ b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java @@ -85,7 +85,7 @@ public static void methodExit( final AgentSpan span = fromContext(scope.context()); DECORATE.onError(span, throwable); DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java index a534b6d3c6d..da9f9ce6426 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java @@ -92,7 +92,7 @@ public static void setupCallback( final AgentSpan span = scope.span(); if (throwable != null) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); scope.close(); return; diff --git a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java index 778b35ddd43..33dbfef94ac 100644 --- a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java +++ b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java @@ -22,7 +22,7 @@ public static void beginRequest( try (final AgentScope scope = activateSpan(span)) { DECORATE.onResponse(span, message); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java index 8842e6e9be0..90cc371e59c 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java @@ -32,7 +32,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope, t); throw t; } final Future wrapped = @@ -40,14 +40,14 @@ public Future apply(final HttpRequest request) { new AbstractFunction1() { @Override public HttpResponse apply(final HttpResponse response) { - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope, t); return t; } }, diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java index 533a77e1ec4..81166d91ab0 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java @@ -114,12 +114,12 @@ public void onPush() throws Exception { final HttpResponse response = grab(responseInlet); final ContextScope scope = scopes.poll(); if (scope != null) { - AgentSpan span = fromContext(scope.context()); - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope, response); // Check if the active span matches the scope from when the request came in, // and close it. If it's not, then it will be cleaned up actor message // processing instrumentation that drives this state machine AgentSpan activeSpan = activeSpan(); + AgentSpan span = fromContext(scope.context()); if (activeSpan == span) { scope.close(); } @@ -144,8 +144,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception { ContextScope scope = scopes.poll(); if (scope != null) { // Mark the span as failed - AgentSpan span = fromContext(scope.context()); - DatadogWrapperHelper.finishSpan(span, ex); + DatadogWrapperHelper.finishSpan(scope, ex); } // We will not receive any more responses from the user code, so clean up any // remaining spans diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index 5c15a75fa46..a29ab47768d 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -20,17 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) { return context.attach(); } - public static void finishSpan(final AgentSpan span, final HttpResponse response) { + public static void finishSpan(final ContextScope scope, final HttpResponse response) { + final AgentSpan span = fromContext(scope.context()); DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } - public static void finishSpan(final AgentSpan span, final Throwable t) { + public static void finishSpan(final ContextScope scope, final Throwable t) { + final AgentSpan span = fromContext(scope.context()); DECORATE.onError(span, t); span.setHttpStatusCode(500); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } diff --git a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java index b3014547ce7..fdbc677b0c5 100644 --- a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java @@ -55,14 +55,14 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); if (REPORT_HTTP_STATUS) { playControllerSpan.setHttpStatusCode(500); } - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java index 68b56de4ddb..6a51c2cf294 100644 --- a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play23; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -12,10 +14,12 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1, Object> { private static final Logger log = LoggerFactory.getLogger(RequestCompleteCallback.class); + private final ContextScope scope; private final AgentSpan span; - public RequestCompleteCallback(final AgentSpan span) { - this.span = span; + public RequestCompleteCallback(final ContextScope scope) { + this.scope = scope; + this.span = spanFromContext(scope.context()); } @Override @@ -28,7 +32,7 @@ public Object apply(final Try result) { DECORATE.onResponse(span, result.get()); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java index ff617daa07a..ce5713e5afa 100644 --- a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java @@ -70,14 +70,14 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); if (REPORT_HTTP_STATUS) { playControllerSpan.setHttpStatusCode(500); } - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java index 7841c6c4977..07fd7efa2d6 100644 --- a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play24; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -12,9 +14,11 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1, Object> { private static final Logger log = LoggerFactory.getLogger(RequestCompleteCallback.class); private final AgentSpan span; + private final ContextScope scope; - public RequestCompleteCallback(final AgentSpan span) { - this.span = span; + public RequestCompleteCallback(final ContextScope scope) { + this.span = spanFromContext(scope.context()); + this.scope = scope; } @Override @@ -27,7 +31,7 @@ public Object apply(final Try result) { DECORATE.onResponse(span, result.get()); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java index 83d6a74cf76..89103f8e2b6 100644 --- a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java @@ -72,11 +72,11 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java index 14c96f8a3bd..97e84394109 100644 --- a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play26; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,9 +16,11 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1 result) { DECORATE.updateOn404Only(span, response); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java index 6c6b16a938c..5a498a65829 100644 --- a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java +++ b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java @@ -81,7 +81,7 @@ public static void finishRequest( DECORATE.onError(span, error); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index 120fbb039a0..a87230f1135 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java @@ -116,9 +116,8 @@ public static void stopSpan( } DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); - span.finish(); } } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java index 3c70e9be54c..8d7080bba95 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.servlet3; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.TIMEOUT; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE; +import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; @@ -26,14 +28,16 @@ public class FinishAsyncDispatchListener implements AsyncListener, Runnable { private final AtomicBoolean activated; private final AgentSpan span; private final boolean doOnResponse; + private final ContextScope scope; - public FinishAsyncDispatchListener(final AgentSpan span, boolean doOnResponse) { - this(span, new AtomicBoolean(), doOnResponse); + public FinishAsyncDispatchListener(final ContextScope scope, boolean doOnResponse) { + this(scope, new AtomicBoolean(), doOnResponse); } public FinishAsyncDispatchListener( - final AgentSpan span, AtomicBoolean activated, boolean doOnResponse) { - this.span = span; + final ContextScope scope, AtomicBoolean activated, boolean doOnResponse) { + this.scope = scope; + this.span = spanFromContext(scope.context()); this.activated = activated; this.doOnResponse = doOnResponse; } @@ -54,7 +58,7 @@ public void onComplete(final AsyncEvent event) throws IOException { DECORATE.onError(span, (Throwable) error); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -62,7 +66,7 @@ public void onComplete(final AsyncEvent event) throws IOException { @Override public void run() { if (activated.compareAndSet(false, true)) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -74,7 +78,7 @@ public void onTimeout(final AsyncEvent event) throws IOException { span.setError(true); } span.setTag(TIMEOUT, event.getAsyncContext().getTimeout()); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -83,7 +87,7 @@ public void onTimeout(final AsyncEvent event) throws IOException { public void onError(final AsyncEvent event) throws IOException { if (event.getThrowable() != null && activated.compareAndSet(false, true)) { DECORATE.onError(span, event.getThrowable()); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java index 9a8a55f159d..fc5563840e3 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java @@ -144,7 +144,7 @@ public static void stopSpan( if (request.isAsyncStarted()) { AtomicBoolean activated = new AtomicBoolean(); FinishAsyncDispatchListener asyncListener = - new FinishAsyncDispatchListener(span, activated, !isDispatch); + new FinishAsyncDispatchListener(scope, activated, !isDispatch); // Jetty doesn't always call the listener, if the request ends before request.setAttribute(DD_FIN_DISP_LIST_SPAN_ATTRIBUTE, asyncListener); try { @@ -158,7 +158,7 @@ public static void stopSpan( DECORATE.onResponse(span, resp); } if (finishSpan) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } @@ -181,7 +181,7 @@ public static void stopSpan( DECORATE.onError(span, throwable); } if (finishSpan) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala index 29436ce4bd4..802d22f10ea 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala @@ -23,7 +23,7 @@ object SprayHelper { case throwable: Throwable => DECORATE.onError(span, throwable) case x => } - DECORATE.beforeFinish(span) + DECORATE.beforeFinish(parentContext) span.finish() message }) diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java index 82a310fe133..b0945179d19 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java @@ -106,7 +106,7 @@ public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { public static void stopSpan( @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { DECORATE_RENDER.onError(scope, throwable); - DECORATE_RENDER.beforeFinish(scope); + DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); scope.span().finish(); } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index 9304644ea12..123fbc0a638 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -122,7 +122,7 @@ public static void stopSpan( finish = true; } if (finish) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index 670b4f2ac45..8d56521fce0 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -73,7 +73,7 @@ public static void stopSpan( finish = true; } if (finish) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java index 44e88d34688..a829cfdbeba 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java @@ -22,7 +22,7 @@ public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { public static void stopSpan( @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { SpringWebHttpServerDecorator.DECORATE_RENDER.onError(scope, throwable); - SpringWebHttpServerDecorator.DECORATE_RENDER.beforeFinish(scope); + SpringWebHttpServerDecorator.DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); scope.span().finish(); } From c9a2e21bb5b7484d7a6bf0342debe5885087fd11 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 12:10:31 -0400 Subject: [PATCH 04/27] fix tests --- .../trace/instrumentation/servlet2/Servlet2Advice.java | 1 + .../datadog/trace/instrumentation/spray/SprayHelper.scala | 7 +++++-- .../spray/SprayHttpServerRunSealedRouteAdvice.java | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index a87230f1135..1299b6d0bb6 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java @@ -119,5 +119,6 @@ public static void stopSpan( DECORATE.beforeFinish(scope.context()); scope.close(); + span.finish(); } } diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala index 802d22f10ea..f1c994b33d2 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.spray import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; @@ -14,7 +15,8 @@ object SprayHelper { def wrapRequestContext( ctx: RequestContext, span: AgentSpan, - parentContext: Context + parentContext: Context, + scope: ContextScope ): RequestContext = { ctx.withRouteResponseMapped(message => { DECORATE.onRequest(span, ctx, ctx.request, parentContext) @@ -23,7 +25,8 @@ object SprayHelper { case throwable: Throwable => DECORATE.onError(span, throwable) case x => } - DECORATE.beforeFinish(parentContext) + DECORATE.beforeFinish(scope.context()) + scope.close() span.finish() message }) diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java index 9a44db3476d..7bb8e735eb8 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java @@ -36,7 +36,7 @@ public static ContextScope enter( ContextScope scope = context.attach(); DECORATE.afterStart(span); - ctx = SprayHelper.wrapRequestContext(ctx, span, parentContext); + ctx = SprayHelper.wrapRequestContext(ctx, span, parentContext, scope); return scope; } From a0ecac7ea818154fca820380f9ff9c7bf6f8f40b Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 13:10:43 -0400 Subject: [PATCH 05/27] round 2 of context migration --- .../SynapseServerInstrumentation.java | 36 ++++++++++--------- .../SynapseServerWorkerInstrumentation.java | 2 +- .../tomcat/TomcatDecorator.java | 1 + .../tomcat/RequestInstrumentation.java | 17 +++++---- .../tomcat/ResponseInstrumentation.java | 17 +++++---- .../tomcat/TomcatServerInstrumentation.java | 3 ++ .../undertow/ExchangeEndSpanListener.java | 19 +++++----- .../HttpRequestParserInstrumentation.java | 2 +- 8 files changed, 58 insertions(+), 39 deletions(-) diff --git a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java index 6a01290b458..0982e1b9cb3 100644 --- a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java @@ -74,8 +74,9 @@ public static ContextScope beginRequest( DECORATE.afterStart(span); DECORATE.onRequest(span, connection, request, context); - // capture span to be finished by one of the various server response advices - connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span); + // capture context (which contains span) to be finished by one of the various server response + // advices + connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, context); return scope; } @@ -90,10 +91,10 @@ public static final class ServerResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope beginResponse( @Advice.Argument(0) final NHttpServerConnection connection) { - // check and remove span from context so it won't be finished twice - AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); - if (null != span) { - return span.attach(); + // check and remove context so it won't be finished twice + Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); + if (null != context) { + return context.attach(); } return null; } @@ -111,7 +112,7 @@ public static void responseReady( if (null != error) { DECORATE.onError(span, error); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } @@ -122,16 +123,19 @@ public static final class ServerErrorResponseAdvice { public static void errorResponse( @Advice.Argument(0) final NHttpServerConnection connection, @Advice.Argument(value = 1, optional = true) final Object error) { - // check and remove span from context so it won't be finished twice - AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); - if (null != span) { - if (error instanceof Throwable) { - DECORATE.onError(span, (Throwable) error); - } else { - span.setError(true); + // check and remove context so it won't be finished twice + Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); + if (null != context) { + AgentSpan span = spanFromContext(context); + if (null != span) { + if (error instanceof Throwable) { + DECORATE.onError(span, (Throwable) error); + } else { + span.setError(true); + } + DECORATE.beforeFinish(context); + span.finish(); } - DECORATE.beforeFinish(span); - span.finish(); } } } diff --git a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java index 6bfa5a5db2e..12fceba2509 100644 --- a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java @@ -97,7 +97,7 @@ public static void responseReady( // (if there's an ACK response or error we might not get a separate response event) if ((null != httpResponse || null != error) && null != request.getConnection().getContext().removeAttribute(SYNAPSE_SPAN_KEY)) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } else { diff --git a/dd-java-agent/instrumentation/tomcat-5.5-common/src/main/java/datadog/trace/instrumentation/tomcat/TomcatDecorator.java b/dd-java-agent/instrumentation/tomcat-5.5-common/src/main/java/datadog/trace/instrumentation/tomcat/TomcatDecorator.java index 5a4de186d88..ca856ae97d2 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5-common/src/main/java/datadog/trace/instrumentation/tomcat/TomcatDecorator.java +++ b/dd-java-agent/instrumentation/tomcat-5.5-common/src/main/java/datadog/trace/instrumentation/tomcat/TomcatDecorator.java @@ -21,6 +21,7 @@ public class TomcatDecorator public static final TomcatDecorator DECORATE = new TomcatDecorator(); public static final String DD_PARENT_CONTEXT_ATTRIBUTE = "datadog.parent-context"; + public static final String DD_CONTEXT_ATTRIBUTE = "datadog.context"; public static final String DD_CONTEXT_PATH_ATTRIBUTE = "datadog.context.path"; public static final String DD_SERVLET_PATH_ATTRIBUTE = "datadog.servlet.path"; public static final String DD_REAL_STATUS_CODE = "datadog.servlet.real_status_code"; diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java index 9c292ae4e1d..43f63890860 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java @@ -2,11 +2,13 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,19 +82,22 @@ public static void stopSpan(@Advice.This final Request req) { Response resp = req.getResponse(); Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { + if (spanObj instanceof AgentSpan && contextObj instanceof Context) { /** - * This advice will be called for both Request and Response. The span is removed from the - * request so the advice only applies the first invocation. (So it doesn't matter which is - * recycled first.) + * This advice will be called for both Request and Response. The span and context are + * removed from the request so the advice only applies the first invocation. (So it doesn't + * matter which is recycled first.) */ - // value set on the coyote request, so we must remove directly from there. + // values set on the coyote request, so we must remove directly from there. req.getCoyoteRequest().setAttribute(DD_SPAN_ATTRIBUTE, null); + req.getCoyoteRequest().setAttribute(DD_CONTEXT_ATTRIBUTE, null); final AgentSpan span = (AgentSpan) spanObj; + final Context context = (Context) contextObj; DECORATE.onResponse(span, resp); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(context); span.finish(); } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java index e17ad93f3c0..869943182bf 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java @@ -2,10 +2,12 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -57,19 +59,22 @@ public static void stopSpan(@Advice.This final Response resp) { Request req = resp.getRequest(); Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { + if (spanObj instanceof AgentSpan && contextObj instanceof Context) { /** - * This advice will be called for both Request and Response. The span is removed from the - * request so the advice only applies the first invocation. (So it doesn't matter which is - * recycled first.) + * This advice will be called for both Request and Response. The span and context are + * removed from the request so the advice only applies the first invocation. (So it doesn't + * matter which is recycled first.) */ - // value set on the coyote request, so we must remove directly from there. + // values set on the coyote request, so we must remove directly from there. req.getCoyoteRequest().setAttribute(DD_SPAN_ATTRIBUTE, null); + req.getCoyoteRequest().setAttribute(DD_CONTEXT_ATTRIBUTE, null); final AgentSpan span = (AgentSpan) spanObj; + final Context context = (Context) contextObj; DECORATE.onResponse(span, resp); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(context); span.finish(); } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index 4534a35b8e5..baf0ca56f04 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -8,6 +8,7 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; +import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; import static java.util.Collections.singletonMap; @@ -132,7 +133,9 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque final AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); + // Store both span and context req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index 76cb6392f4c..c3c555e4990 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -21,17 +21,18 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener return; } - AgentSpan span = continuation.span(); + try (AgentScope scope = continuation.activate()) { + AgentSpan span = scope.span(); - Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION); - if (throwable != null) { - DECORATE.onError(span, throwable); - } + Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION); + if (throwable != null) { + DECORATE.onError(span, throwable); + } - DECORATE.onResponse(span, exchange); - DECORATE.beforeFinish(span); - continuation.cancel(); - span.finish(); + DECORATE.onResponse(span, exchange); + DECORATE.beforeFinish(scope.context()); + span.finish(); + } nextListener.proceed(); } } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java index 5977c2af5b4..13f126deca2 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java @@ -85,7 +85,7 @@ public static void afterRequestParse( DECORATE.onError(span, throwable); // because we know that a http 400 will be thrown DECORATE.onResponseStatus(span, 400); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } finally { if (span != null) { span.finish(); From 2ef3dc78c1db28353a60883d307ddf7bca724a15 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 13:31:38 -0400 Subject: [PATCH 06/27] netty changes --- .../netty38/ChannelFutureListenerInstrumentation.java | 2 +- .../netty38/server/HttpServerRequestTracingHandler.java | 2 +- .../netty38/server/HttpServerResponseTracingHandler.java | 2 +- .../netty40/ChannelFutureListenerInstrumentation.java | 2 +- .../netty40/server/HttpServerRequestTracingHandler.java | 2 +- .../netty40/server/HttpServerResponseTracingHandler.java | 2 +- .../netty41/ChannelFutureListenerInstrumentation.java | 2 +- .../netty41/server/HttpServerRequestTracingHandler.java | 2 +- .../netty41/server/HttpServerResponseTracingHandler.java | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 4a211023143..6836a3bb72e 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -111,7 +111,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f errorSpan.context().setIntegrationName(NETTY); try (final AgentScope scope = activateSpan(errorSpan)) { DECORATE.onError(errorSpan, cause); - DECORATE.beforeFinish(errorSpan); + DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java index 862dc01cb2d..ae7c9a463dc 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java @@ -70,7 +70,7 @@ public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent ctx.sendUpstream(msg); } catch (final Throwable throwable) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false throw throwable; } diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index a32891cb820..2b77f94c660 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -58,7 +58,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m && (response.getStatus() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java index e0b426f7697..255fb792bab 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java @@ -96,7 +96,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f errorSpan.context().setIntegrationName(NETTY); try (final AgentScope scope = activateSpan(errorSpan)) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); - NettyHttpServerDecorator.DECORATE.beforeFinish(errorSpan); + NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java index 58bb49cb647..6908510ea2d 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java @@ -67,7 +67,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { ctx.fireChannelRead(msg); } catch (final Throwable throwable) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); throw throwable; diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java index 47399950b44..82ac4439eed 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java @@ -58,7 +58,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann && (response.getStatus() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); span.finish(); // Finish the span manually since finishSpanOnClose was false } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java index 23211173f4a..a552db7704b 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java @@ -96,7 +96,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f errorSpan.context().setIntegrationName(NETTY); try (final AgentScope scope = activateSpan(errorSpan)) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); - NettyHttpServerDecorator.DECORATE.beforeFinish(errorSpan); + NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java index 30197668254..b5a9fa2c95e 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java @@ -72,7 +72,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { The handler chain failed with exception - need to finish the span here */ DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); throw throwable; diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java index 07235193fc8..ba1b389b69c 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java @@ -52,7 +52,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann if (response.status() != HttpResponseStatus.CONTINUE && (response.status() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); } From a14848fc738e9b6d628b0dbf1ae94329d16ebe8d Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 14:13:03 -0400 Subject: [PATCH 07/27] tomcat changes --- .../decorator/HttpServerDecorator.java | 1 + .../instrumentation/tomcat/TomcatDecorator.java | 1 - .../tomcat/RequestInstrumentation.java | 14 ++++++-------- .../tomcat/ResponseInstrumentation.java | 14 ++++++-------- .../tomcat/TomcatServerInstrumentation.java | 14 ++++++-------- 5 files changed, 19 insertions(+), 25 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 92b5ede5592..e3aed68b4cf 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -57,6 +57,7 @@ public abstract class HttpServerDecorator Date: Tue, 26 Aug 2025 14:39:38 -0400 Subject: [PATCH 08/27] liberty changes --- .../LibertyServerInstrumentation.java | 10 +++--- .../RequestFinishInstrumentation.java | 23 ++++++++----- .../ResponseFinishInstrumentation.java | 34 ++++++++++++------- .../LibertyServerInstrumentation.java | 10 +++--- .../RequestFinishInstrumentation.java | 23 ++++++++----- .../ResponseFinishInstrumentation.java | 34 ++++++++++++------- 6 files changed, 80 insertions(+), 54 deletions(-) diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java index 1cb460dfbcb..cc984c4d071 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -96,9 +96,9 @@ public static class HandleRequestAdvice { // if we try to get an attribute that doesn't exist open liberty might complain with an // exception try { - Object existingSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - scope = ((AgentSpan) existingSpan).attach(); + Object existingContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + scope = ((Context) existingContext).attach(); return false; } } catch (NullPointerException e) { @@ -123,7 +123,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java index fa50b98f53b..7daf4045919 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty20; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -9,6 +10,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletRequest; import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.servlet.IExtendedResponse; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -57,19 +59,22 @@ public static void stopSpan(@Advice.This SRTServletRequest req) { // this should be a servlet response if (resp instanceof SRTServletResponse) { SRTServletResponse httpResp = (SRTServletResponse) resp; - Object spanObj = null; + Object contextObj = null; try { - spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - if (spanObj instanceof AgentSpan) { - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, httpResp); - DECORATE.beforeFinish(span); - span.finish(); + if (contextObj instanceof Context) { + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, httpResp); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java index 3bd85a7f2d1..5041e9ee6a6 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty20; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -11,6 +12,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.WebContainerRequestState; import com.ibm.wsspi.webcontainer.servlet.IExtendedRequest; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,33 +82,39 @@ static void after(@Advice.Argument(0) boolean releaseChannel) { @SuppressFBWarnings("DCN_NULLPOINTER_EXCEPTION") public static class ResponseFinishAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentSpan onEnter(@Advice.This SRTServletResponse resp) { + static Context onEnter(@Advice.This SRTServletResponse resp) { // this is the last opportunity to have any meaningful // interaction with the response - AgentSpan span = null; + Context context = null; IExtendedRequest req = resp.getRequest(); try { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - span = (AgentSpan) spanObj; - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - DECORATE.onResponse(span, resp); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + context = (Context) contextObj; + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, resp); + } } } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - return span; + return context; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.This SRTServletResponse resp, @Advice.Enter AgentSpan span) { - if (span == null) { + @Advice.This SRTServletResponse resp, @Advice.Enter Context context) { + if (context == null) { return; } - DECORATE.beforeFinish(span); - span.finish(); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java index 2fb526942d7..644644bee89 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -98,9 +98,9 @@ public static class HandleRequestAdvice { // if we try to get an attribute that doesn't exist open liberty might complain with an // exception try { - Object existingSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - scope = ((AgentSpan) existingSpan).attach(); + Object existingContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + scope = ((Context) existingContext).attach(); return false; } } catch (NullPointerException e) { @@ -125,7 +125,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java index a2b2172dde4..a6e633bcfba 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty23; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -9,6 +10,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletRequest; import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.servlet.IExtendedResponse; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -57,19 +59,22 @@ public static void stopSpan(@Advice.This SRTServletRequest req) { // this should be a servlet response if (resp instanceof SRTServletResponse) { SRTServletResponse httpResp = (SRTServletResponse) resp; - Object spanObj = null; + Object contextObj = null; try { - spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - if (spanObj instanceof AgentSpan) { - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, httpResp); - DECORATE.beforeFinish(span); - span.finish(); + if (contextObj instanceof Context) { + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, httpResp); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java index 96ebcc7252e..d8b4c2a013e 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty23; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -11,6 +12,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.WebContainerRequestState; import com.ibm.wsspi.webcontainer.servlet.IExtendedRequest; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,33 +82,39 @@ static void after(@Advice.Argument(0) boolean releaseChannel) { @SuppressFBWarnings("DCN_NULLPOINTER_EXCEPTION") public static class ResponseFinishAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentSpan onEnter(@Advice.This SRTServletResponse resp) { + static Context onEnter(@Advice.This SRTServletResponse resp) { // this is the last opportunity to have any meaningful // interaction with the response - AgentSpan span = null; + Context context = null; IExtendedRequest req = resp.getRequest(); try { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - span = (AgentSpan) spanObj; - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - DECORATE.onResponse(span, resp); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + context = (Context) contextObj; + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, resp); + } } } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - return span; + return context; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.This SRTServletResponse resp, @Advice.Enter AgentSpan span) { - if (span == null) { + @Advice.This SRTServletResponse resp, @Advice.Enter Context context) { + if (context == null) { return; } - DECORATE.beforeFinish(span); - span.finish(); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.beforeFinish(context); + span.finish(); + } } } } From 909d174fbcf5ff53e2329d6c34fefb72edf2b663 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 15:42:29 -0400 Subject: [PATCH 09/27] store span --- .../liberty20/LibertyServerInstrumentation.java | 2 ++ .../liberty23/LibertyServerInstrumentation.java | 2 ++ .../instrumentation/tomcat/TomcatServerInstrumentation.java | 3 ++- 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java index cc984c4d071..1417e9a33b0 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; @@ -123,6 +124,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); + request.setAttribute(DD_SPAN_ATTRIBUTE, span); request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java index 644644bee89..9630e0b9c1d 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; @@ -125,6 +126,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); + request.setAttribute(DD_SPAN_ATTRIBUTE, span); request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index 293e18d83fe..d643195a9f5 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -132,7 +132,8 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque final AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); - // Store context + // Store both span and context + req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); From c883e8daa2d0ea5bfee20615a596882a795e91df Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 16:50:34 -0400 Subject: [PATCH 10/27] jetty --- .../jetty11/JettyServerAdvice.java | 24 ++++++++------ .../jetty12/JettyServerAdvice.java | 23 +++++++------ .../jetty70/JettyServerInstrumentation.java | 25 +++++++++------ .../jetty70/ServerHandleInstrumentation.java | 32 +++++++++++++++---- .../jetty76/JettyServerInstrumentation.java | 25 +++++++++------ .../jetty76/ServerHandleInstrumentation.java | 32 +++++++++++++++---- .../jetty9/JettyServerInstrumentation.java | 23 +++++++------ .../jetty9/ServerHandleInstrumentation.java | 32 +++++++++++++++---- .../instrumentation/jetty10/HandleAdvice.java | 8 +++-- .../instrumentation/jetty10/ResetAdvice.java | 19 +++++++---- .../jetty10/ServerHandleAdvice.java | 32 +++++++++++++++---- .../tomcat/TomcatServerInstrumentation.java | 1 - 12 files changed, 189 insertions(+), 87 deletions(-) diff --git a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java index 65dbaffbe4e..a800d14d492 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty11; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE; @@ -21,9 +21,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -35,6 +35,7 @@ public static ContextScope onEnter( DECORATE.onRequest(span, req, req, parentContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -58,12 +59,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java index ed8c193b334..a376b0026d5 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.jetty12; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty12.JettyDecorator.DECORATE; @@ -20,9 +21,9 @@ public static void onExit( @Advice.This final HttpChannelState channel, @Advice.Return(readOnly = false) Runnable ret) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - try (final ContextScope scope = ((AgentSpan) existingSpan).attach()) { + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + try (final ContextScope scope = ((Context) existingContext).attach()) { ret = JettyRunnableWrapper.wrapIfNeeded(ret); return; } @@ -37,6 +38,7 @@ public static void onExit( DECORATE.onRequest(span, req, req, parentContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); ret = JettyRunnableWrapper.wrapIfNeeded(ret); @@ -52,12 +54,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannelState channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java index af7f8a4c7d6..c3b4045d65f 100644 --- a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java @@ -2,6 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; import static java.util.Collections.singletonMap; @@ -149,10 +150,10 @@ public static ContextScope onEnter( @Advice.This final HttpConnection connection, @Advice.Local("agentSpan") AgentSpan span) { Request req = connection.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // Request already gone through initial processing, so just activate the span. - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + // Request already gone through initial processing, so just activate the context. + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -163,6 +164,7 @@ public static ContextScope onEnter( DECORATE.onRequest(span, req, req, parentContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -183,12 +185,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpConnection channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java index 58ddd318ab4..426c856a589 100644 --- a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java @@ -1,16 +1,18 @@ package datadog.trace.instrumentation.jetty70; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpConnection; @@ -53,12 +55,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpConnection connection, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = connection.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // see comments in HandleRequestAdvice for jetty-9 Object dispatchSpan; synchronized (req) { @@ -66,8 +71,15 @@ static AgentScope onEnter( } if (dispatchSpan instanceof AgentSpan) { span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -75,7 +87,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -83,12 +95,18 @@ public static void onExit( return; } + // Extract the span from the context if we didn't get it from @Local + if (span == null) { + span = spanFromContext(scope.context()); + } + if (t != null) { DECORATE.onError(span, t); } if (!req.getAsyncContinuation().isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java index ec9ef262fb1..60385bf1510 100644 --- a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java @@ -2,6 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; @@ -150,10 +151,10 @@ public static ContextScope onEnter( @Advice.Local("newSpan") AgentSpan span) { Request req = connection.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // Request already gone through initial processing, so just activate the span. - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + // Request already gone through initial processing, so just activate the context. + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -164,6 +165,7 @@ public static ContextScope onEnter( DECORATE.onRequest(span, req, req, parentContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -184,12 +186,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final AbstractHttpConnection connection) { Request req = connection.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, connection); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, connection); + DECORATE.beforeFinish(context); + span.finish(); + } } // Jetty doesn't always call async listeners diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java index 86f239b9e69..5faff27ff31 100644 --- a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java @@ -1,16 +1,18 @@ package datadog.trace.instrumentation.jetty76; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.AbstractHttpConnection; @@ -53,12 +55,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) AbstractHttpConnection connection, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = connection.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // see comments in HandleRequestAdvice for jetty-9 Object dispatchSpan; synchronized (req) { @@ -66,8 +71,15 @@ static AgentScope onEnter( } if (dispatchSpan instanceof AgentSpan) { span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -75,7 +87,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("agentSpan") AgentSpan span, @Advice.Local("request") Request req, @Advice.Thrown Throwable t) { @@ -83,12 +95,18 @@ public static void onExit( return; } + // Extract the span from the context if we didn't get it from @Local + if (span == null) { + span = spanFromContext(scope.context()); + } + if (t != null) { DECORATE.onError(span, t); } if (!req.getAsyncContinuation().isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java index b85a639fb2b..a31c158d9f4 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java @@ -4,6 +4,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; @@ -165,9 +166,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -178,6 +179,7 @@ public static ContextScope onEnter( DECORATE.onRequest(span, req, req, parentContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -197,12 +199,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java index 477e60d80b5..2e7c9ab2ac5 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.jetty9; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; @@ -9,9 +10,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -69,12 +71,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span, @Advice.Local("request") Request req) { req = channel.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // same logic as in Servlet3Advice. We need to activate/finish the dispatch span here // because we don't know if a servlet is going to be called and therefore whether // Servlet3Advice will have an opportunity to run. @@ -95,8 +100,15 @@ static AgentScope onEnter( // the root span, stored in DD_SPAN_ATTRIBUTE. // req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -104,7 +116,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -112,6 +124,11 @@ public static void onExit( return; } + // Extract the span from the context if we didn't get it from @Local + if (span == null) { + span = spanFromContext(scope.context()); + } + boolean registeredFinishListener; synchronized (req) { req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); @@ -131,7 +148,8 @@ public static void onExit( } if (!(req.isAsyncStarted() && registeredFinishListener)) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java index ece1c7720d2..9ecb84c5496 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.jetty10; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; @@ -19,9 +20,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -32,6 +33,7 @@ public static ContextScope onEnter( final ContextScope scope = context.attach(); req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ResetAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ResetAdvice.java index 8e7e7c68cd2..f398135f14f 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ResetAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ResetAdvice.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -16,12 +18,15 @@ public class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - JettyDecorator.OnResponse.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + JettyDecorator.OnResponse.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java index 4c759e94ea0..6fad9b03ded 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java @@ -1,10 +1,12 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -12,12 +14,15 @@ class ServerHandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpChannel channel, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = channel.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // same logic as in Servlet3Advice. We need to activate/finish the dispatch span here // because we don't know if a servlet is going to be called and therefore whether // Servlet3Advice will have an opportunity to run. @@ -33,8 +38,15 @@ static AgentScope onEnter( // the root span, stored in DD_SPAN_ATTRIBUTE. // req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -42,7 +54,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -50,12 +62,18 @@ public static void onExit( return; } + // Extract the span from the context if we didn't get it from @Local + if (span == null) { + span = spanFromContext(scope.context()); + } + if (t != null) { DECORATE.onError(span, t); } if (!req.isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index d643195a9f5..f4a849c9371 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -132,7 +132,6 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque final AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); - // Store both span and context req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); From b8474013dd73b402796245557d866fa2c82aebbd Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 17:45:16 -0400 Subject: [PATCH 11/27] grizzly --- .../decorator/HttpServerDecorator.java | 10 --------- .../grizzly/GrizzlyBlockingHelper.java | 21 ++++++++++++------- .../GrizzlyHttpHandlerInstrumentation.java | 4 +++- .../grizzly/SpanClosingListener.java | 20 ++++++++++++------ .../grizzlyhttp232/GrizzlyDecorator.java | 6 +++++- 5 files changed, 36 insertions(+), 25 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index e3aed68b4cf..e3b2b5d6da7 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -540,16 +540,6 @@ public Context beforeFinish(Context context) { return super.beforeFinish(context); } - @Override - public AgentSpan beforeFinish(AgentSpan span) { - onRequestEndForInstrumentationGateway(span); - - // Close Serverless Gateway Inferred Span if any - // finishInferredProxySpan(context); - - return super.beforeFinish(span); - } - protected void finishInferredProxySpan(Context context) { InferredProxySpan span; if ((span = InferredProxySpan.fromContext(context)) != null) { diff --git a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java index 0be7ae6ee01..076d3ab66ae 100644 --- a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java +++ b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.grizzly; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.blocking.BlockingActionHelper; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -35,14 +37,14 @@ public class GrizzlyBlockingHelper { private GrizzlyBlockingHelper() {} public static boolean block( - Request request, Response response, Flow.Action.RequestBlockingAction rba, AgentSpan span) { + Request request, Response response, Flow.Action.RequestBlockingAction rba, Context context) { return block( request, response, rba.getStatusCode(), rba.getBlockingContentType(), rba.getExtraHeaders(), - span); + context); } public static boolean block( @@ -51,11 +53,12 @@ public static boolean block( int statusCode, BlockingContentType bct, Map extraHeaders, - AgentSpan span) { + Context context) { if (GET_OUTPUT_STREAM == null) { return false; } + AgentSpan span = spanFromContext(context); try { OutputStream os = (OutputStream) GET_OUTPUT_STREAM.invoke(response); response.setStatus(BlockingActionHelper.getHttpCode(statusCode)); @@ -76,13 +79,17 @@ public static boolean block( os.close(); response.finish(); - span.getRequestContext().getTraceSegment().effectivelyBlocked(); + if (span != null) { + span.getRequestContext().getTraceSegment().effectivelyBlocked(); + } SpanClosingListener.LISTENER.onAfterService(request); } catch (Throwable e) { log.info("Error committing blocking response", e); - DECORATE.onError(span, e); - DECORATE.beforeFinish(span); - span.finish(); + if (span != null) { + DECORATE.onError(span, e); + DECORATE.beforeFinish(context); + span.finish(); + } } return true; diff --git a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java index cc12314e046..f7dd6eb5548 100644 --- a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java @@ -2,6 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -81,13 +82,14 @@ public static class HandleAdvice { scope = context.attach(); request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); if (rba != null) { - boolean success = GrizzlyBlockingHelper.block(request, response, rba, span); + boolean success = GrizzlyBlockingHelper.block(request, response, rba, context); if (success) { return true; /* skip body */ } diff --git a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java index c5c243e3369..5faa06411b4 100644 --- a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java +++ b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java @@ -1,8 +1,11 @@ package datadog.trace.instrumentation.grizzly; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.glassfish.grizzly.http.server.AfterServiceListener; import org.glassfish.grizzly.http.server.Request; @@ -12,13 +15,18 @@ public class SpanClosingListener implements AfterServiceListener { @Override public void onAfterService(final Request request) { - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanAttr instanceof AgentSpan) { + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { request.removeAttribute(DD_SPAN_ATTRIBUTE); - final AgentSpan span = (AgentSpan) spanAttr; - DECORATE.onResponse(span, request.getResponse()); - DECORATE.beforeFinish(span); - span.finish(); + request.removeAttribute(DD_CONTEXT_ATTRIBUTE); + + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, request.getResponse()); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index 978a0d65824..e9e3db0010e 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.grizzlyhttp232; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import datadog.appsec.api.blocking.BlockingContentType; import datadog.context.Context; @@ -121,6 +122,7 @@ public static NextAction onHttpCodecFilterExit( DECORATE.afterStart(span); ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span); ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse); + ctx.getAttributes().setAttribute(DD_CONTEXT_ATTRIBUTE, context); DECORATE.onRequest(span, httpRequest, httpRequest, parentContext); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -146,13 +148,15 @@ public static NextAction onHttpCodecFilterExit( public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable) { AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); + Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); if (null != span) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span).finish(); + DECORATE.beforeFinish(context); span.finish(); } ctx.getAttributes().removeAttribute(DD_SPAN_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); + ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); } @Override From f76415fbc97345a6f2f7103e2647ea8a06d26bd0 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 26 Aug 2025 18:45:59 -0400 Subject: [PATCH 12/27] fix grizzly instrumentation --- .../grizzlyhttp232/GrizzlyDecorator.java | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index e9e3db0010e..f97f8ada9af 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -97,11 +97,13 @@ public static void onHttpServerFilterPrepareResponseEnter( public static void onHttpServerFilterPrepareResponseExit( FilterChainContext ctx, HttpResponsePacket responsePacket) { AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (null != span) { - DECORATE.beforeFinish(span); + Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + if (null != span && null != context) { + DECORATE.beforeFinish(context); span.finish(); } ctx.getAttributes().removeAttribute(DD_SPAN_ATTRIBUTE); + ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); } @@ -149,14 +151,14 @@ public static NextAction onHttpCodecFilterExit( public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable) { AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); - if (null != span) { + if (null != span && null != context) { DECORATE.onError(span, throwable); DECORATE.beforeFinish(context); span.finish(); } ctx.getAttributes().removeAttribute(DD_SPAN_ATTRIBUTE); - ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); + ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); } @Override From 762c23381630426651d90fe491d3412666b6d5a3 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 11:33:26 -0400 Subject: [PATCH 13/27] remove serverless --- .github/CODEOWNERS | 2 - .../decorator/HttpServerDecorator.java | 22 ---- .../trace/api/config/TracerConfig.java | 3 - .../java/datadog/trace/core/CoreTracer.java | 5 - .../propagation/InferredProxyPropagator.java | 60 ---------- .../InferredProxyPropagatorTests.java | 96 ---------------- .../main/java/datadog/trace/api/Config.java | 10 -- .../trace/api/gateway/InferredProxySpan.java | 108 ------------------ .../instrumentation/api/AgentPropagation.java | 1 - .../api/gateway/InferredProxySpanTests.java | 95 --------------- 10 files changed, 402 deletions(-) delete mode 100644 dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java delete mode 100644 dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java delete mode 100644 internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java delete mode 100644 internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 96ab0c6452e..e0cd6d19f9d 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -30,8 +30,6 @@ # @DataDog/apm-serverless /dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless /dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless -**/InferredProxy*.java @DataDog/apm-serverless -**/InferredProxy*.groovy @DataDog/apm-serverless # @DataDog/apm-lang-platform-java /.circleci/ @DataDog/apm-lang-platform-java diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index e3b2b5d6da7..43e534de194 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -20,7 +20,6 @@ import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction; import datadog.trace.api.gateway.IGSpanInfo; -import datadog.trace.api.gateway.InferredProxySpan; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.naming.SpanNaming; @@ -152,8 +151,6 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { AgentSpanContext extracted = getExtractedSpanContext(context); // Call IG callbacks extracted = callIGCallbackStart(extracted); - // Create gateway inferred span if needed - extracted = startInferredProxySpan(context, extracted); AgentSpan span = tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true); // Apply RequestBlockingAction if any @@ -166,15 +163,6 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) { return context.with(span); } - protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanContext extracted) { - InferredProxySpan span; - if (!Config.get().isInferredProxyPropagationEnabled() - || (span = InferredProxySpan.fromContext(context)) == null) { - return extracted; - } - return span.start(extracted); - } - public AgentSpan onRequest( final AgentSpan span, final CONNECTION connection, @@ -534,19 +522,9 @@ public Context beforeFinish(Context context) { AgentSpan span = AgentSpan.fromContext(context); onRequestEndForInstrumentationGateway(span); - // Close Serverless Gateway Inferred Span if any - finishInferredProxySpan(context); - return super.beforeFinish(context); } - protected void finishInferredProxySpan(Context context) { - InferredProxySpan span; - if ((span = InferredProxySpan.fromContext(context)) != null) { - span.finish(); - } - } - private void onRequestEndForInstrumentationGateway(@Nonnull final AgentSpan span) { if (span.getLocalRootSpan() != span) { return; 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 ae5d515cb62..0300cef2070 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 @@ -100,9 +100,6 @@ public final class TracerConfig { public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes"; public static final String TRACE_BAGGAGE_TAG_KEYS = "trace.baggage.tag.keys"; - public static final String TRACE_INFERRED_PROXY_SERVICES_ENABLED = - "trace.inferred.proxy.services.enabled"; - public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled"; public static final String CLIENT_IP_ENABLED = "trace.client-ip.enabled"; 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 71cf4fe76c4..4c81a579f83 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,7 +7,6 @@ import static datadog.trace.api.TracePropagationBehaviorExtract.IGNORE; 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.INFERRED_PROXY_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; @@ -92,7 +91,6 @@ import datadog.trace.core.monitor.TracerHealthMetrics; import datadog.trace.core.propagation.ExtractedContext; import datadog.trace.core.propagation.HttpCodec; -import datadog.trace.core.propagation.InferredProxyPropagator; import datadog.trace.core.propagation.PropagationTags; import datadog.trace.core.propagation.TracingPropagator; import datadog.trace.core.propagation.XRayPropagator; @@ -822,9 +820,6 @@ private CoreTracer( && config.getTracePropagationBehaviorExtract() != IGNORE) { Propagators.register(BAGGAGE_CONCERN, new BaggagePropagator(config)); } - if (config.isInferredProxyPropagationEnabled()) { - Propagators.register(INFERRED_PROXY_CONCERN, new InferredProxyPropagator()); - } if (config.isCiVisibilityEnabled()) { if (config.isCiVisibilityTraceSanitationEnabled()) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java deleted file mode 100644 index da51bf4643b..00000000000 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/InferredProxyPropagator.java +++ /dev/null @@ -1,60 +0,0 @@ -package datadog.trace.core.propagation; - -import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; - -import datadog.context.Context; -import datadog.context.propagation.CarrierSetter; -import datadog.context.propagation.CarrierVisitor; -import datadog.context.propagation.Propagator; -import datadog.trace.api.gateway.InferredProxySpan; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; -import javax.annotation.ParametersAreNonnullByDefault; - -/** Inferred proxy propagator. Only extract, not meant for injection. */ -@ParametersAreNonnullByDefault -public class InferredProxyPropagator implements Propagator { - private static final String INFERRED_PROXY_KEY_PREFIX = "x-dd-proxy"; - - @Override - public void inject(Context context, C carrier, CarrierSetter setter) {} - - @Override - public Context extract(Context context, C carrier, CarrierVisitor visitor) { - if (context == null || carrier == null || visitor == null) { - return context; - } - InferredProxyContextExtractor extractor = new InferredProxyContextExtractor(); - visitor.forEachKeyValue(carrier, extractor); - InferredProxySpan inferredProxySpan = extractor.inferredProxySpan(); - if (inferredProxySpan != null) { - context = context.with(inferredProxySpan); - } - return context; - } - - /** Extract inferred proxy related headers into a map. */ - private static class InferredProxyContextExtractor implements BiConsumer { - private Map values; - - @Override - public void accept(String key, String value) { - if (key == null || key.isEmpty() || !key.startsWith(INFERRED_PROXY_KEY_PREFIX)) { - return; - } - if (values == null) { - this.values = new HashMap<>(); - } - this.values.put(key, value); - } - - public InferredProxySpan inferredProxySpan() { - if (this.values == null) { - return null; - } - InferredProxySpan inferredProxySpan = fromHeaders(this.values); - return inferredProxySpan.isValid() ? inferredProxySpan : null; - } - } -} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java deleted file mode 100644 index 5830a160cc7..00000000000 --- a/dd-trace-core/src/test/java/datadog/trace/core/propagation/InferredProxyPropagatorTests.java +++ /dev/null @@ -1,96 +0,0 @@ -package datadog.trace.core.propagation; - -import static datadog.context.Context.root; -import static datadog.trace.api.gateway.InferredProxySpan.fromContext; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.params.provider.Arguments.of; - -import datadog.context.Context; -import datadog.context.propagation.CarrierVisitor; -import datadog.trace.api.gateway.InferredProxySpan; -import java.util.HashMap; -import java.util.Map; -import java.util.function.BiConsumer; -import java.util.stream.Stream; -import javax.annotation.ParametersAreNonnullByDefault; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; - -@DisplayName("InferredProxyPropagator Tests") -class InferredProxyPropagatorTests { - private static final String PROXY_SYSTEM_KEY = "x-dd-proxy"; - private static final String PROXY_REQUEST_TIME_MS_KEY = "x-dd-proxy-request-time-ms"; - private static final String PROXY_PATH_KEY = "x-dd-proxy-path"; - private static final String PROXY_HTTP_METHOD_KEY = "x-dd-proxy-httpmethod"; - private static final String PROXY_DOMAIN_NAME_KEY = "x-dd-proxy-domain-name"; - private static final MapVisitor MAP_VISITOR = new MapVisitor(); - - private InferredProxyPropagator propagator; - - @BeforeEach - void setUp() { - this.propagator = new InferredProxyPropagator(); - } - - @Test - @DisplayName("Should extract InferredProxySpan when valid headers are present") - void testSuccessfulExtraction() { - Map headers = new HashMap<>(); - headers.put(PROXY_SYSTEM_KEY, "aws-apigateway"); - headers.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - headers.put(PROXY_PATH_KEY, "/foo"); - headers.put(PROXY_HTTP_METHOD_KEY, "GET"); - headers.put(PROXY_DOMAIN_NAME_KEY, "api.example.com"); - - Context context = this.propagator.extract(root(), headers, MAP_VISITOR); - InferredProxySpan inferredProxySpan = fromContext(context); - assertNotNull(inferredProxySpan); - assertTrue(inferredProxySpan.isValid()); - } - - @ParameterizedTest(name = "{0}") - @MethodSource("invalidOrMissingHeadersProviderForPropagator") - @DisplayName("Should not create InferredProxySpan if some critical headers are missing") - void testExtractionWithMissingCriticalHeaders(String description, Map headers) { - Context rootContext = root(); - Context extractedOuterContext = this.propagator.extract(rootContext, headers, MAP_VISITOR); - InferredProxySpan inferredProxySpan = fromContext(extractedOuterContext); - assertNull(inferredProxySpan, "Invalid inferred proxy span should not be extracted"); - } - - static Stream invalidOrMissingHeadersProviderForPropagator() { // Renamed - Map missingSystem = new HashMap<>(); - missingSystem.put(PROXY_REQUEST_TIME_MS_KEY, "12345"); - missingSystem.put(PROXY_PATH_KEY, "/foo"); - - Map emptyValue = new HashMap<>(); - emptyValue.put(PROXY_SYSTEM_KEY, ""); - - Map nullValue = new HashMap<>(); - nullValue.put(PROXY_SYSTEM_KEY, null); - - Map missingTime = new HashMap<>(); - missingTime.put(PROXY_SYSTEM_KEY, "aws-apigw"); - missingTime.put(PROXY_PATH_KEY, "/foo"); - - return Stream.of( - of("PROXY_SYSTEM_KEY missing", missingSystem), - of("PROXY_SYSTEM_KEY empty", emptyValue), - of("PROXY_SYSTEM_KEY null", nullValue), - of("PROXY_REQUEST_TIME_MS_KEY missing", missingTime)); - } - - @ParametersAreNonnullByDefault - private static class MapVisitor implements CarrierVisitor> { - @Override - public void forEachKeyValue(Map carrier, BiConsumer visitor) { - carrier.forEach(visitor); - } - } -} 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 5622717bd89..36a165ee1c3 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -598,7 +598,6 @@ import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_RESOURCE_REMOVE_TRAILING_SLASH; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_ERROR_STATUSES; import static datadog.trace.api.config.TracerConfig.TRACE_HTTP_SERVER_PATH_RESOURCE_NAME_MAPPING; -import static datadog.trace.api.config.TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_KEEP_LATENCY_THRESHOLD_MS; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_ENABLED; import static datadog.trace.api.config.TracerConfig.TRACE_LONG_RUNNING_FLUSH_INTERVAL; @@ -828,7 +827,6 @@ public static String getHostName() { private final int traceBaggageMaxItems; private final int traceBaggageMaxBytes; private final List traceBaggageTagKeys; - private final boolean traceInferredProxyEnabled; private final int clockSyncPeriod; private final boolean logsInjectionEnabled; @@ -1732,8 +1730,6 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins tracePropagationExtractFirst = configProvider.getBoolean( TRACE_PROPAGATION_EXTRACT_FIRST, DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST); - traceInferredProxyEnabled = - configProvider.getBoolean(TRACE_INFERRED_PROXY_SERVICES_ENABLED, false); clockSyncPeriod = configProvider.getInteger(CLOCK_SYNC_PERIOD, DEFAULT_CLOCK_SYNC_PERIOD); @@ -3122,10 +3118,6 @@ public boolean isTracePropagationExtractFirst() { return tracePropagationExtractFirst; } - public boolean isInferredProxyPropagationEnabled() { - return traceInferredProxyEnabled; - } - public boolean isBaggageExtract() { return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); } @@ -5420,8 +5412,6 @@ public String toString() { + tracePropagationBehaviorExtract + ", tracePropagationExtractFirst=" + tracePropagationExtractFirst - + ", traceInferredProxyEnabled=" - + traceInferredProxyEnabled + ", clockSyncPeriod=" + clockSyncPeriod + ", jmxFetchEnabled=" diff --git a/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java b/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java deleted file mode 100644 index a8256c1b38a..00000000000 --- a/internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java +++ /dev/null @@ -1,108 +0,0 @@ -package datadog.trace.api.gateway; - -import static datadog.context.ContextKey.named; -import static datadog.trace.api.DDTags.RESOURCE_NAME; -import static datadog.trace.api.DDTags.SERVICE_NAME; -import static datadog.trace.api.DDTags.SPAN_TYPE; -import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; -import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_METHOD; -import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL; - -import datadog.context.Context; -import datadog.context.ContextKey; -import datadog.context.ImplicitContextKeyed; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; -import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; -import datadog.trace.bootstrap.instrumentation.api.AgentTracer; -import java.util.Collections; -import java.util.HashMap; -import java.util.Map; - -public class InferredProxySpan implements ImplicitContextKeyed { - private static final ContextKey CONTEXT_KEY = named("inferred-proxy-key"); - static final String PROXY_SYSTEM = "x-dd-proxy"; - static final String PROXY_START_TIME_MS = "x-dd-proxy-request-time-ms"; - static final String PROXY_PATH = "x-dd-proxy-path"; - static final String PROXY_HTTP_METHOD = "x-dd-proxy-httpmethod"; - static final String PROXY_DOMAIN_NAME = "x-dd-proxy-domain-name"; - static final String STAGE = "x-dd-proxy-stage"; - static final Map SUPPORTED_PROXIES; - static final String INSTRUMENTATION_NAME = "inferred_proxy"; - - static { - SUPPORTED_PROXIES = new HashMap<>(); - SUPPORTED_PROXIES.put("aws-apigateway", "aws.apigateway"); - } - - private final Map headers; - private AgentSpan span; - - public static InferredProxySpan fromHeaders(Map values) { - return new InferredProxySpan(values); - } - - public static InferredProxySpan fromContext(Context context) { - return context.get(CONTEXT_KEY); - } - - private InferredProxySpan(Map headers) { - this.headers = headers == null ? Collections.emptyMap() : headers; - } - - public boolean isValid() { - String startTimeStr = header(PROXY_START_TIME_MS); - String proxySystem = header(PROXY_SYSTEM); - return startTimeStr != null - && proxySystem != null - && SUPPORTED_PROXIES.containsKey(proxySystem); - } - - public AgentSpanContext start(AgentSpanContext extracted) { - if (this.span != null || !isValid()) { - return extracted; - } - - long startTime; - try { - startTime = Long.parseLong(header(PROXY_START_TIME_MS)) * 1000; // Convert to microseconds - } catch (NumberFormatException e) { - return extracted; // Invalid timestamp - } - - String proxySystem = header(PROXY_SYSTEM); - String proxy = SUPPORTED_PROXIES.get(proxySystem); - AgentSpan span = AgentTracer.get().startSpan(INSTRUMENTATION_NAME, proxy, extracted, startTime); - - span.setTag(COMPONENT, proxySystem); - span.setTag(RESOURCE_NAME, header(PROXY_HTTP_METHOD) + " " + header(PROXY_PATH)); - span.setTag(SERVICE_NAME, header(PROXY_DOMAIN_NAME)); - span.setTag(SPAN_TYPE, "web"); - span.setTag(HTTP_METHOD, header(PROXY_HTTP_METHOD)); - span.setTag(HTTP_URL, header(PROXY_DOMAIN_NAME) + header(PROXY_PATH)); - span.setTag("stage", header(STAGE)); - span.setTag("_dd.inferred_span", 1); - - // Free collected headers - this.headers.clear(); - // Store inferred span - this.span = span; - // Return inferred span as new parent context - return this.span.context(); - } - - private String header(String name) { - return this.headers.get(name); - } - - public void finish() { - if (this.span != null) { - this.span.finish(); - this.span = null; - } - } - - @Override - public Context storeInto(Context context) { - return context.with(CONTEXT_KEY, this); - } -} 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 0fbb4484353..f8339f6ed94 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 @@ -18,7 +18,6 @@ public final class AgentPropagation { // TODO: remove this priority once we have a story for replacing TagContext with the Context API public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105); public static final Concern XRAY_TRACING_CONCERN = named("tracing-xray"); - public static final Concern INFERRED_PROXY_CONCERN = named("inferred-proxy"); // 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); diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java b/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java deleted file mode 100644 index dc61d4a23aa..00000000000 --- a/internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java +++ /dev/null @@ -1,95 +0,0 @@ -package datadog.trace.api.gateway; - -import static datadog.context.Context.root; -import static datadog.trace.api.gateway.InferredProxySpan.PROXY_START_TIME_MS; -import static datadog.trace.api.gateway.InferredProxySpan.PROXY_SYSTEM; -import static datadog.trace.api.gateway.InferredProxySpan.fromContext; -import static datadog.trace.api.gateway.InferredProxySpan.fromHeaders; -import static java.util.Collections.emptyMap; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.params.provider.Arguments.of; - -import datadog.context.Context; -import java.util.HashMap; -import java.util.Map; -import java.util.stream.Stream; -import org.junit.jupiter.api.DisplayName; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; - -@DisplayName("InferredProxyContext Tests") -class InferredProxySpanTests { - @Test - @DisplayName("Valid headers should make valid span") - void testMapConstructor() { - Map headers = new HashMap<>(); - headers.put(PROXY_START_TIME_MS, "12345"); - headers.put(PROXY_SYSTEM, "aws-apigateway"); - - InferredProxySpan inferredProxySpan = InferredProxySpan.fromHeaders(headers); - assertTrue(inferredProxySpan.isValid()); - assertNotNull( - inferredProxySpan.start(null), "inferred proxy span start and return new parent context"); - assertNull(inferredProxySpan.start(null), "inferred proxy span not should start twice"); - inferredProxySpan.finish(); - } - - @ParameterizedTest(name = "{0}") - @DisplayName("Invalid headers should make invalid span") - @MethodSource("invalidHeaders") - void testInvalidHeaders(String useCase, Map headers) { - InferredProxySpan inferredProxySpan = fromHeaders(headers); - assertFalse(inferredProxySpan.isValid(), useCase + " should not be valid"); - assertNull(inferredProxySpan.start(null), "Invalid inferred proxy span should not start"); - } - - static Stream invalidHeaders() { // Renamed - Map missingSystem = new HashMap<>(); - missingSystem.put(PROXY_START_TIME_MS, "12345"); - - Map missingTime = new HashMap<>(); - missingTime.put(PROXY_SYSTEM, "aws-apigateway"); - Map invalidSystem = new HashMap<>(); - invalidSystem.put(PROXY_START_TIME_MS, "12345"); - invalidSystem.put(PROXY_SYSTEM, "invalidSystem"); - - return Stream.of( - of("Missing system headers", missingSystem), - of("Missing start time headers", missingTime), - of("Invalid system headers", invalidSystem)); - } - - @Test - @DisplayName("Constructor with null should not crash") - void testNullMapConstructor() { - InferredProxySpan inferredProxySpan = fromHeaders(null); - assertNotNull(inferredProxySpan); - assertFalse(inferredProxySpan.isValid()); - } - - @Test - @DisplayName("Constructor with empty map should be invalid") - void testEmptyMapConstructor() { - InferredProxySpan inferredProxySpan = fromHeaders(emptyMap()); - assertNotNull(inferredProxySpan); - assertFalse(inferredProxySpan.isValid()); - } - - @Test - @DisplayName("storeInto and fromContext should correctly attach and retrieve the context") - void testStoreAndFromContext() { - InferredProxySpan inferredProxySpan = fromHeaders(null); - Context context = inferredProxySpan.storeInto(root()); - assertNotNull(context); - - InferredProxySpan retrieved = fromContext(context); - assertNotNull(retrieved); - - assertNull(fromContext(root()), "fromContext on empty context should be null"); - } -} From c8c9ece91ca5dde6997f8accc3601176836fc32c Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 12:22:13 -0400 Subject: [PATCH 14/27] fixes part 1 --- .github/CODEOWNERS | 6 +++--- .../decorator/BaseDecorator.java | 2 +- .../akkahttp/DatadogAsyncHandlerWrapper.java | 8 +++---- ...tadogServerRequestResponseFlowWrapper.java | 6 +++--- .../akkahttp/DatadogWrapperHelper.java | 12 +++++------ .../pekkohttp/DatadogAsyncHandlerWrapper.java | 6 +++--- ...tadogServerRequestResponseFlowWrapper.java | 4 ++-- .../pekkohttp/DatadogWrapperHelper.java | 12 +++++------ .../undertow/ExchangeEndSpanListener.java | 21 ++++++++++--------- .../instrumentation/api/AgentPropagation.java | 1 + 10 files changed, 40 insertions(+), 38 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index e0cd6d19f9d..ca8633372e8 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -28,8 +28,8 @@ /internal-api/src/test/groovy/datadog/trace/api/sampling @DataDog/apm-sdk-api-java # @DataDog/apm-serverless -/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless -/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/main/java/datadog/trace/lambda/ @DataDog/apm-serverless +/dd-trace-core/src/test/groovy/datadog/trace/lambda/ @DataDog/apm-serverless # @DataDog/apm-lang-platform-java /.circleci/ @DataDog/apm-lang-platform-java @@ -127,4 +127,4 @@ dd-trace-api/src/main/java/datadog/trace/api/llmobs/ @DataDog/ml-observability dd-java-agent/agent-llmobs/ @DataDog/ml-observability dd-trace-core/src/main/java/datadog/trace/llmobs/ @DataDog/ml-observability -dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability +dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability \ No newline at end of file diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java index 7472b660f54..9cde0e0e244 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java @@ -77,7 +77,7 @@ public AgentSpan afterStart(final AgentSpan span) { return span; } - public AgentScope beforeFinish(final AgentScope scope) { + public ContextScope beforeFinish(final ContextScope scope) { beforeFinish(scope.context()); return scope; } diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java index 324993dd75c..20df10cd40f 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java @@ -38,7 +38,7 @@ public Future apply(final HttpRequest request) { request.discardEntityBytes(materializer); HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request); span.getRequestContext().getTraceSegment().effectivelyBlocked(); - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); return FastFuture$.MODULE$.successful().apply(response); } @@ -46,7 +46,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(scope, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); throw t; } @@ -67,14 +67,14 @@ public HttpResponse apply(HttpResponse response) { response = newResponse; } - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(scope, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); return t; } }, diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java index 2a23d586ac0..29f0cb497d6 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java @@ -78,7 +78,7 @@ public void onPush() throws Exception { skipNextPull[0] = true; requestContext.getTraceSegment().effectivelyBlocked(); emit(responseOutlet, response); - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); pull(requestInlet); scope.close(); return; @@ -142,7 +142,7 @@ public void onPush() throws Exception { response.discardEntityBytes(materializer()); response = newResponse; } - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); // Check if the active span matches the scope from when the request came in, // and close it. If it's not, then it will be cleaned up actor message // processing instrumentation that drives this state machine @@ -172,7 +172,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception { if (scope != null) { // Mark the span as failed AgentSpan span = fromContext(scope.context()); - DatadogWrapperHelper.finishSpan(scope, ex); + DatadogWrapperHelper.finishSpan(scope.context(), ex); } // We will not receive any more responses from the user code, so clean up any // remaining spans diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index 571936a6966..2acf49c3ad6 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java @@ -20,19 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) { return context.attach(); } - public static void finishSpan(final ContextScope scope, final HttpResponse response) { - final AgentSpan span = fromContext(scope.context()); + public static void finishSpan(final Context context, final HttpResponse response) { + final AgentSpan span = fromContext(context); DECORATE.onResponse(span, response); - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); span.finish(); } - public static void finishSpan(final ContextScope scope, final Throwable t) { - final AgentSpan span = fromContext(scope.context()); + public static void finishSpan(final Context context, final Throwable t) { + final AgentSpan span = fromContext(context); DECORATE.onError(span, t); span.setHttpStatusCode(500); - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); span.finish(); } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java index 90cc371e59c..0578eda8aba 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java @@ -32,7 +32,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(scope, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); throw t; } final Future wrapped = @@ -40,14 +40,14 @@ public Future apply(final HttpRequest request) { new AbstractFunction1() { @Override public HttpResponse apply(final HttpResponse response) { - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(scope, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); return t; } }, diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java index 81166d91ab0..00c1057b3f0 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java @@ -114,7 +114,7 @@ public void onPush() throws Exception { final HttpResponse response = grab(responseInlet); final ContextScope scope = scopes.poll(); if (scope != null) { - DatadogWrapperHelper.finishSpan(scope, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); // Check if the active span matches the scope from when the request came in, // and close it. If it's not, then it will be cleaned up actor message // processing instrumentation that drives this state machine @@ -144,7 +144,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception { ContextScope scope = scopes.poll(); if (scope != null) { // Mark the span as failed - DatadogWrapperHelper.finishSpan(scope, ex); + DatadogWrapperHelper.finishSpan(scope.context(), ex); } // We will not receive any more responses from the user code, so clean up any // remaining spans diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index a29ab47768d..d4565c68bfc 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -20,19 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) { return context.attach(); } - public static void finishSpan(final ContextScope scope, final HttpResponse response) { - final AgentSpan span = fromContext(scope.context()); + public static void finishSpan(final Context context, final HttpResponse response) { + final AgentSpan span = fromContext(context); DECORATE.onResponse(span, response); - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); span.finish(); } - public static void finishSpan(final ContextScope scope, final Throwable t) { - final AgentSpan span = fromContext(scope.context()); + public static void finishSpan(final Context context, final Throwable t) { + final AgentSpan span = fromContext(context); DECORATE.onError(span, t); span.setHttpStatusCode(500); - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); span.finish(); } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index c3c555e4990..e6573f5cf4b 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.undertow; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.undertow.server.DefaultResponseListener; @@ -21,18 +23,17 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener return; } - try (AgentScope scope = continuation.activate()) { - AgentSpan span = scope.span(); + Context context = continuation.context(); + AgentSpan span = fromContext(context); - Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION); - if (throwable != null) { - DECORATE.onError(span, throwable); - } - - DECORATE.onResponse(span, exchange); - DECORATE.beforeFinish(scope.context()); - span.finish(); + Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION); + if (throwable != null) { + DECORATE.onError(span, throwable); } + + DECORATE.onResponse(span, exchange); + DECORATE.beforeFinish(context); + span.finish(); nextListener.proceed(); } } 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 f8339f6ed94..fba659bea74 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 @@ -18,6 +18,7 @@ public final class AgentPropagation { // TODO: remove this priority once we have a story for replacing TagContext with the Context API public static final Concern BAGGAGE_CONCERN = withPriority("baggage", 105); 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 f715d628956d0164df414956e89ddccbe3b11024 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 13:47:44 -0400 Subject: [PATCH 15/27] fix undertow tests --- .github/CODEOWNERS | 2 +- .../trace/instrumentation/undertow/ExchangeEndSpanListener.java | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ca8633372e8..ef437eebae2 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -127,4 +127,4 @@ dd-trace-api/src/main/java/datadog/trace/api/llmobs/ @DataDog/ml-observability dd-java-agent/agent-llmobs/ @DataDog/ml-observability dd-trace-core/src/main/java/datadog/trace/llmobs/ @DataDog/ml-observability -dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability \ No newline at end of file +dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index e6573f5cf4b..2b57a5b14ae 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -33,6 +33,7 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener DECORATE.onResponse(span, exchange); DECORATE.beforeFinish(context); + continuation.cancel(); span.finish(); nextListener.proceed(); } From 7fe9c70c4f19e64ed72ad2fcd809142b07943e76 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 16:45:45 -0400 Subject: [PATCH 16/27] remove span within request --- .../decorator/HttpServerDecorator.java | 1 - .../cxf/InvokerInstrumentation.java | 15 +++-- .../GrizzlyHttpHandlerInstrumentation.java | 4 +- .../grizzly/SpanClosingListener.java | 2 - .../DefaultFilterChainInstrumentation.java | 16 +++-- .../grizzlyhttp232/FilterAdvice.java | 17 ++++-- .../grizzlyhttp232/GrizzlyDecorator.java | 36 ++++++----- .../v1/JaxRsClientV1Instrumentation.java | 12 ++-- .../jetty11/JettyServerAdvice.java | 2 - .../jetty11/SetContextPathAdvice.java | 34 ++++++----- .../jetty12/JettyServerAdvice.java | 2 - .../jetty12/SetContextPathAdvice.java | 58 ++++++++++-------- .../JettyCommitResponseInstrumentation.java | 15 +++-- .../jetty70/JettyServerInstrumentation.java | 2 - .../jetty70/RequestInstrumentation.java | 28 ++++++--- .../JettyCommitResponseInstrumentation.java | 15 +++-- .../jetty76/JettyServerInstrumentation.java | 2 - .../jetty76/RequestInstrumentation.java | 28 ++++++--- .../JettyCommitResponseInstrumentation.java | 16 +++-- .../jetty9/JettyServerInstrumentation.java | 2 - .../jetty9/RequestInstrumentation.java | 28 ++++++--- .../jetty9/ServerHandleInstrumentation.java | 2 +- .../instrumentation/jetty10/HandleAdvice.java | 2 - .../jetty10/JettyCommitResponseHelper.java | 16 +++-- .../jetty10/ServerHandleAdvice.java | 2 +- .../jetty10/SetContextPathAdvice.java | 16 +++-- .../jetty10/SetServletPathAdvice.java | 16 +++-- .../jetty904/JettyCommitResponseHelper.java | 16 +++-- .../jetty93/JettyCommitResponseHelper.java | 16 +++-- .../jetty9421/JettyCommitResponseHelper.java | 16 +++-- .../LibertyServerInstrumentation.java | 2 - .../LibertyServerInstrumentation.java | 2 - .../netty40/AttributeKeys.java | 6 +- ...yChannelHandlerContextInstrumentation.java | 7 ++- .../HttpClientRequestTracingHandler.java | 4 +- .../HttpClientResponseTracingHandler.java | 28 +++++++-- .../HttpServerRequestTracingHandler.java | 20 ++++--- .../HttpServerResponseTracingHandler.java | 11 ++-- .../server/MaybeBlockResponseHandler.java | 7 ++- .../netty41/AttributeKeys.java | 6 +- ...exHandlerStreamChannelInstrumentation.java | 8 +-- ...yChannelHandlerContextInstrumentation.java | 7 ++- .../HttpClientRequestTracingHandler.java | 4 +- .../HttpClientResponseTracingHandler.java | 32 ++++++++-- .../HttpServerRequestTracingHandler.java | 23 ++++---- .../HttpServerResponseTracingHandler.java | 11 ++-- .../server/MaybeBlockResponseHandler.java | 7 ++- .../ratpack/TracingHandler.java | 14 +++-- .../servlet2/Servlet2Advice.java | 35 ++++++----- .../servlet3/AsyncContextInstrumentation.java | 59 ++++++++++--------- .../servlet3/Servlet3Advice.java | 36 ++++++----- .../JakartaServletInstrumentation.java | 17 ++++-- .../RequestDispatcherInstrumentation.java | 28 +++++---- .../test/groovy/RequestDispatcherTest.groovy | 28 +++++---- .../ZuulSendForwardFilterInstrumentation.java | 11 ++-- ...MappingResourceNameFilterForkedTest.groovy | 6 +- .../HandlerAdapterInstrumentation.java | 14 +++-- .../HandlerMappingResourceNameFilter.java | 30 ++++++---- ...MappingResourceNameFilterForkedTest.groovy | 6 +- .../springweb6/ControllerAdvice.java | 14 +++-- .../HandlerMappingResourceNameFilter.java | 30 ++++++---- .../tomcat/TomcatServerInstrumentation.java | 32 +++++----- .../tomcat7/CommitActionInstrumentation.java | 9 ++- .../undertow/ServletInstrumentation.java | 4 +- .../JakartaServletInstrumentation.java | 4 +- .../HttpClientRequestBaseInstrumentation.java | 7 ++- 66 files changed, 634 insertions(+), 372 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 43e534de194..f5804e4106b 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -55,7 +55,6 @@ public abstract class HttpServerDecorator>> MAPS = GenericClassValue.constructing(ConcurrentHashMap.class); - public static final AttributeKey SPAN_ATTRIBUTE_KEY = attributeKey(DD_SPAN_ATTRIBUTE); + public static final AttributeKey CONTEXT_ATTRIBUTE_KEY = + attributeKey(DD_CONTEXT_ATTRIBUTE); public static final AttributeKey CLIENT_PARENT_ATTRIBUTE_KEY = attributeKey("datadog.client.parent.span"); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java index 7fe652004ff..76ca4218b06 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java @@ -6,13 +6,15 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.NettyChannelPipelineInstrumentation.ADDITIONAL_INSTRUMENTATION_NAMES; import static datadog.trace.instrumentation.netty40.NettyChannelPipelineInstrumentation.INSTRUMENTATION_NAME; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -69,7 +71,8 @@ public void methodAdvice(MethodTransformer transformer) { public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { - final AgentSpan channelSpan = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan channelSpan = storedContext != null ? spanFromContext(storedContext) : null; if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index fbbf2309d32..541b06c0f96 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -6,7 +6,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.netty40.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE_SECURE; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; @@ -90,7 +90,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(scope.context()); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index b0cc383f5e1..6819f817fe1 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -2,10 +2,12 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -24,7 +26,13 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } final boolean finishSpan = msg instanceof HttpResponse; @@ -47,7 +55,13 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -68,7 +82,13 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java index 6908510ea2d..5bb70729845 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java @@ -3,8 +3,8 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; @@ -27,7 +27,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { Channel channel = ctx.channel(); if (!(msg instanceof HttpRequest)) { - final AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; if (span == null) { ctx.fireChannelRead(msg); // superclass does not throw } else { @@ -51,7 +52,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); - channel.attr(SPAN_ATTRIBUTE_KEY).set(span); + channel.attr(CONTEXT_ATTRIBUTE_KEY).set(context); channel.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).set(request.headers()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -69,7 +70,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { DECORATE.onError(span, throwable); DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } } @@ -81,10 +82,13 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { super.channelInactive(ctx); } finally { try { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndRemove(); - if (span != null && span.phasedFinish()) { - // at this point we can just publish this span to avoid loosing the rest of the trace - span.publish(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); + if (storedContext != null) { + final AgentSpan span = spanFromContext(storedContext); + if (span != null && span.phasedFinish()) { + // at this point we can just publish this span to avoid loosing the rest of the trace + span.publish(); + } } } catch (final Throwable ignored) { } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java index 82ac4439eed..462e1996cd4 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java @@ -1,11 +1,13 @@ package datadog.trace.instrumentation.netty40.server; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CHANNEL_ID; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; @@ -24,7 +26,8 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); return; @@ -39,7 +42,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.onError(span, throwable); span.setHttpStatusCode(500); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } final boolean isWebsocketUpgrade = @@ -59,7 +62,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); DECORATE.beforeFinish(scope.context()); - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java index 6e3ec443840..a0c01ab5944 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.netty40.server; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; import static io.netty.handler.codec.http.HttpHeaders.setContentLength; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; @@ -56,7 +58,8 @@ private static void markBlockedResponse(Channel ch) { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) throws Exception { Channel channel = ctx.channel(); - AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null diff --git a/dd-java-agent/instrumentation/netty-4.1-shared/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java b/dd-java-agent/instrumentation/netty-4.1-shared/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java index 41ab0146d3d..b6af132f173 100644 --- a/dd-java-agent/instrumentation/netty-4.1-shared/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java +++ b/dd-java-agent/instrumentation/netty-4.1-shared/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.netty41; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.api.GenericClassValue; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -16,7 +17,8 @@ public final class AttributeKeys { private static final ClassValue>> MAPS = GenericClassValue.constructing(ConcurrentHashMap.class); - public static final AttributeKey SPAN_ATTRIBUTE_KEY = attributeKey(DD_SPAN_ATTRIBUTE); + public static final AttributeKey CONTEXT_ATTRIBUTE_KEY = + attributeKey(DD_CONTEXT_ATTRIBUTE); public static final AttributeKey CLIENT_PARENT_ATTRIBUTE_KEY = attributeKey("datadog.client.parent.span"); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java index d532a5e1606..261c05662c0 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java @@ -37,10 +37,10 @@ public static class PropagateContextAdvice { @Advice.OnMethodExit(suppress = Throwable.class) public static void afterCreate(@Advice.This Channel self) { if (self.parent() != null - && self.parent().hasAttr(AttributeKeys.SPAN_ATTRIBUTE_KEY) - && !self.hasAttr(AttributeKeys.SPAN_ATTRIBUTE_KEY)) { - self.attr(AttributeKeys.SPAN_ATTRIBUTE_KEY) - .set(self.parent().attr(AttributeKeys.SPAN_ATTRIBUTE_KEY).getAndRemove()); + && self.parent().hasAttr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY) + && !self.hasAttr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY)) { + self.attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY) + .set(self.parent().attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).getAndRemove()); } } } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java index 66f89dc58d9..075bde792ba 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java @@ -6,13 +6,15 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.NettyChannelPipelineInstrumentation.ADDITIONAL_INSTRUMENTATION_NAMES; import static datadog.trace.instrumentation.netty41.NettyChannelPipelineInstrumentation.INSTRUMENTATION_NAME; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -69,7 +71,8 @@ public void methodAdvice(MethodTransformer transformer) { public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { - final AgentSpan channelSpan = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan channelSpan = storedContext != null ? spanFromContext(storedContext) : null; if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index c944965fd34..24a253f773e 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -6,7 +6,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.netty41.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE_SECURE; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; @@ -91,7 +91,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(scope.context()); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 408c3530e20..32c5408e9d9 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -2,10 +2,12 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -26,7 +28,13 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { final boolean finishSpan = @@ -41,7 +49,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + try (final AgentScope spanScope = activateSpan(span)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanScope.context()); + } } } @@ -56,7 +66,13 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -77,7 +93,13 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + + // Set parent context back to maintain the same functionality as getAndSet(parent) + try (final AgentScope parentScope = activateSpan(parent)) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java index b5a9fa2c95e..3d6b324ea6e 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java @@ -3,8 +3,8 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; @@ -26,11 +26,11 @@ public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapte public void channelRead(final ChannelHandlerContext ctx, final Object msg) { Channel channel = ctx.channel(); if (!(msg instanceof HttpRequest)) { - final AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); - if (span == null) { + final Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + if (storedContext == null) { ctx.fireChannelRead(msg); // superclass does not throw } else { - try (final ContextScope scope = span.attach()) { + try (final ContextScope scope = storedContext.attach()) { ctx.fireChannelRead(msg); // superclass does not throw } } @@ -50,7 +50,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); - channel.attr(SPAN_ATTRIBUTE_KEY).set(span); + channel.attr(CONTEXT_ATTRIBUTE_KEY).set(context); channel.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).set(request.headers()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -74,7 +74,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { DECORATE.onError(span, throwable); DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } } @@ -86,10 +86,13 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { super.channelInactive(ctx); } finally { try { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndRemove(); - if (span != null && span.phasedFinish()) { - // at this point we can just publish this span to avoid loosing the rest of the trace - span.publish(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); + if (storedContext != null) { + final AgentSpan span = spanFromContext(storedContext); + if (span != null && span.phasedFinish()) { + // at this point we can just publish this span to avoid loosing the rest of the trace + span.publish(); + } } } catch (final Throwable ignored) { } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java index ba1b389b69c..3290ed1c02b 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java @@ -1,10 +1,12 @@ package datadog.trace.instrumentation.netty41.server; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; @@ -22,7 +24,8 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); @@ -38,7 +41,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.onError(span, throwable); span.setHttpStatusCode(500); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } final boolean isWebsocketUpgrade = @@ -54,7 +57,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.onResponse(span, response); DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); } } } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java index 9139e318703..17ab2f327ca 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.netty41.server; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; import static io.netty.handler.codec.http.HttpHeaders.setContentLength; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.blocking.BlockingActionHelper; @@ -55,7 +57,8 @@ private static void markBlockedResponse(Channel ch) { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) throws Exception { Channel channel = ctx.channel(); - AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null) { super.write(ctx, msg, prm); diff --git a/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java b/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java index ad00fa0041d..1043cb5b20b 100644 --- a/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java +++ b/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java @@ -1,9 +1,10 @@ package datadog.trace.instrumentation.ratpack; import static datadog.context.Context.root; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.ratpack.RatpackServerDecorator.DECORATE; import com.google.common.reflect.TypeToken; @@ -24,16 +25,17 @@ public final class TracingHandler implements Handler { Types.token(Flow.Action.RequestBlockingAction.class); /** This constant must stay in sync with datadog.trace.instrumentation.netty41.AttributeKeys. */ - public static final AttributeKey SERVER_ATTRIBUTE_KEY = - AttributeKey.valueOf(DD_SPAN_ATTRIBUTE); + public static final AttributeKey SERVER_CONTEXT_ATTRIBUTE_KEY = + AttributeKey.valueOf(DD_CONTEXT_ATTRIBUTE); @Override public void handle(final Context ctx) { final Request request = ctx.getRequest(); - final Attribute spanAttribute = - ctx.getDirectChannelAccess().getChannel().attr(SERVER_ATTRIBUTE_KEY); - final AgentSpan nettySpan = spanAttribute.get(); + final Attribute contextAttribute = + ctx.getDirectChannelAccess().getChannel().attr(SERVER_CONTEXT_ATTRIBUTE_KEY); + final datadog.context.Context nettyContext = contextAttribute.get(); + final AgentSpan nettySpan = nettyContext != null ? fromContext(nettyContext) : null; // Relying on executor instrumentation to assume the netty span is in context as the parent. final AgentSpan ratpackSpan = startSpan(DECORATE.spanName()).setMeasured(true); diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index 1299b6d0bb6..ae67ab27c27 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.servlet2; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE; import datadog.context.Context; @@ -37,13 +37,16 @@ public static boolean onEnter( } final HttpServletRequest httpServletRequest = (HttpServletRequest) request; - Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - final boolean hasServletTrace = spanAttr instanceof AgentSpan; - if (hasServletTrace) { - final AgentSpan span = (AgentSpan) spanAttr; - ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - // Tracing might already be applied by the FilterChain or a parent request (forward/include). - return false; + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context existingContext = (Context) contextAttr; + final AgentSpan span = spanFromContext(existingContext); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + // Tracing might already be applied by the FilterChain or a parent request + // (forward/include). + return false; + } } if (response instanceof HttpServletResponse) { @@ -58,7 +61,7 @@ public static boolean onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, httpServletRequest, httpServletRequest, parentContext); - httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); + httpServletRequest.setAttribute(DD_CONTEXT_ATTRIBUTE, context); httpServletRequest.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); httpServletRequest.setAttribute( @@ -85,13 +88,17 @@ public static void stopSpan( @Advice.Local("contextScope") final ContextScope scope, @Advice.Thrown final Throwable throwable) { // Set user.principal regardless of who created this span. - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); if (Config.get().isServletPrincipalEnabled() - && spanAttr instanceof AgentSpan + && contextAttr instanceof Context && request instanceof HttpServletRequest) { - final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); - if (principal != null) { - ((AgentSpan) spanAttr).setTag(DDTags.USER_NAME, principal.getName()); + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); + if (principal != null) { + span.setTag(DDTags.USER_NAME, principal.getName()); + } } } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java index 39ddb3734aa..eeec7e5bf93 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java @@ -6,8 +6,9 @@ import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_DISPATCH; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.AsyncDispatcherDecorator.DECORATE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DD_SERVLET_PATH_ATTRIBUTE; @@ -15,6 +16,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.CallDepthThreadLocalMap; @@ -84,33 +86,36 @@ public static boolean enter( } final ServletRequest request = context.getRequest(); - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanAttr instanceof AgentSpan) { - final AgentSpan parent = (AgentSpan) spanAttr; - - final AgentSpan span = startSpan(SERVLET_DISPATCH, parent.context()); - // This span should get finished by Servlet3Advice - // However, when using Jetty without servlets (directly org.eclipse.jetty.server.Handler), - // that's not the case (see jetty's HandleAdvice) - DECORATE.afterStart(span); - - // These are pulled from attributes because jetty clears them from the request too early. - span.setTag(SERVLET_CONTEXT, request.getAttribute(DD_CONTEXT_PATH_ATTRIBUTE)); - span.setTag(SERVLET_PATH, request.getAttribute(DD_SERVLET_PATH_ATTRIBUTE)); - - synchronized (request) { - request.setAttribute(DD_DISPATCH_SPAN_ATTRIBUTE, span); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context ctxt = (Context) contextAttr; + final AgentSpan parent = spanFromContext(ctxt); + if (parent != null) { + + final AgentSpan span = startSpan(SERVLET_DISPATCH, parent.context()); + // This span should get finished by Servlet3Advice + // However, when using Jetty without servlets (directly org.eclipse.jetty.server.Handler), + // that's not the case (see jetty's HandleAdvice) + DECORATE.afterStart(span); + + // These are pulled from attributes because jetty clears them from the request too early. + span.setTag(SERVLET_CONTEXT, request.getAttribute(DD_CONTEXT_PATH_ATTRIBUTE)); + span.setTag(SERVLET_PATH, request.getAttribute(DD_SERVLET_PATH_ATTRIBUTE)); + + synchronized (request) { + request.setAttribute(DD_DISPATCH_SPAN_ATTRIBUTE, span); + } + + if (args.length == 1 && args[0] instanceof String) { + span.setResourceName((String) args[0]); + } else if (args.length == 2 && args[1] instanceof String) { + span.setResourceName((String) args[1]); + } + + // We can't register FinishAsyncDispatchListener here. + // The dispatch may happen on an onTimeout/onError, and adding listeners + // when listeners are being iterated on causes a ConcurrentModificationException on jetty } - - if (args.length == 1 && args[0] instanceof String) { - span.setResourceName((String) args[0]); - } else if (args.length == 2 && args[1] instanceof String) { - span.setResourceName((String) args[1]); - } - - // We can't register FinishAsyncDispatchListener here. - // The dispatch may happen on an onTimeout/onError, and adding listeners - // when listeners are being iterated on causes a ConcurrentModificationException on jetty } return true; } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java index 078b1ca6e45..e8030aa367f 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java @@ -2,10 +2,10 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE; import datadog.context.Context; @@ -77,14 +77,16 @@ public static boolean onEnter( finishSpan = true; - Object spanAttrValue = request.getAttribute(DD_SPAN_ATTRIBUTE); - final boolean hasServletTrace = spanAttrValue instanceof AgentSpan; - if (hasServletTrace) { - final AgentSpan span = (AgentSpan) spanAttrValue; - ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - // Tracing might already be applied by other instrumentation, - // the FilterChain or a parent request (forward/include). - return false; + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context existingContext = (Context) contextAttr; + final AgentSpan span = spanFromContext(existingContext); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + // Tracing might already be applied by other instrumentation, + // the FilterChain or a parent request (forward/include). + return false; + } } final Context parentContext = DECORATE.extract(httpServletRequest); @@ -95,7 +97,7 @@ public static boolean onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, httpServletRequest, httpServletRequest, parentContext); - httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); + httpServletRequest.setAttribute(DD_CONTEXT_ATTRIBUTE, context); httpServletRequest.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); httpServletRequest.setAttribute( @@ -125,13 +127,17 @@ public static void stopSpan( rumServletWrapper.commit(); } // Set user.principal regardless of who created this span. - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); if (Config.get().isServletPrincipalEnabled() - && spanAttr instanceof AgentSpan + && contextAttr instanceof Context && request instanceof HttpServletRequest) { - final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); - if (principal != null) { - ((AgentSpan) spanAttr).setTag(DDTags.USER_NAME, principal.getName()); + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); + if (principal != null) { + span.setTag(DDTags.USER_NAME, principal.getName()); + } } } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java index 6b82e220519..f9c1e49c30d 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java @@ -3,14 +3,16 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.ClassloaderConfigurationOverrides; @@ -97,12 +99,15 @@ public static AgentSpan before( } } - Object span = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (span instanceof AgentSpan + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context && CallDepthThreadLocalMap.incrementCallDepth(HttpServletRequest.class) == 0) { - final AgentSpan agentSpan = (AgentSpan) span; - ClassloaderConfigurationOverrides.maybeEnrichSpan(agentSpan); - return agentSpan; + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + return span; + } } return null; } diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java index 28676b5b8df..444e86a9f2e 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java @@ -8,7 +8,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.servlet.ServletRequestSetter.SETTER; import static datadog.trace.instrumentation.servlet.SpanNameCache.SERVLET_PREFIX; import static datadog.trace.instrumentation.servlet.SpanNameCache.SPAN_NAME_CACHE; @@ -21,6 +22,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.InstrumentationContext; @@ -85,13 +87,13 @@ public static class RequestDispatcherAdvice { public static AgentScope start( @Advice.Origin("#m") final String method, @Advice.This final RequestDispatcher dispatcher, - @Advice.Local("_requestSpan") Object requestSpan, + @Advice.Local("_requestContext") Object requestContext, @Advice.Argument(0) final ServletRequest request) { final AgentSpan parentSpan = activeSpan(); - final Object servletSpanObject = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); final AgentSpan servletSpan = - servletSpanObject instanceof AgentSpan ? (AgentSpan) servletSpanObject : null; + contextAttr instanceof Context ? spanFromContext((Context) contextAttr) : null; if (parentSpan == null && servletSpan == null) { // Don't want to generate a new top-level span @@ -122,16 +124,21 @@ public static AgentScope start( DECORATE.injectContext(span, request, SETTER); // temporarily replace from request to avoid spring resource name bubbling up: - requestSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + requestContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); - return activateSpan(span); + final AgentScope scope = activateSpan(span); + // Set the context after activation so we have the proper Context object + if (scope != null) { + request.setAttribute(DD_CONTEXT_ATTRIBUTE, scope.context()); + } + + return scope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stop( @Advice.Enter final AgentScope scope, - @Advice.Local("_requestSpan") final Object requestSpan, + @Advice.Local("_requestContext") final Object requestContext, @Advice.Argument(0) final ServletRequest request, @Advice.Argument(1) final ServletResponse response, @Advice.Thrown final Throwable throwable) { @@ -139,9 +146,8 @@ public static void stop( return; } - if (requestSpan != null) { - // now add it back... - request.setAttribute(DD_SPAN_ATTRIBUTE, requestSpan); + if (requestContext != null) { + request.setAttribute(DD_CONTEXT_ATTRIBUTE, requestContext); } DECORATE.onError(scope, throwable); diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy index 6963425837e..4ddcaffdc3c 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy +++ b/dd-java-agent/instrumentation/servlet/src/test/groovy/RequestDispatcherTest.groovy @@ -1,6 +1,6 @@ import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDSpanTypes -import datadog.trace.core.DDSpan +import datadog.context.Context import javax.servlet.ServletException import javax.servlet.http.HttpServletRequest @@ -8,7 +8,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE import static datadog.trace.core.propagation.DatadogHttpCodec.SAMPLING_PRIORITY_KEY import static datadog.trace.core.propagation.DatadogHttpCodec.SPAN_ID_KEY import static datadog.trace.core.propagation.DatadogHttpCodec.TRACE_ID_KEY @@ -17,7 +17,7 @@ class RequestDispatcherTest extends AgentTestRunner { def request = Mock(HttpServletRequest) def response = Mock(HttpServletResponse) - def mockSpan = Stub(DDSpan) + def mockContext = Mock(Context) def dispatcher = new RequestDispatcherUtils(request, response) def "test dispatch no-parent"() { @@ -26,7 +26,7 @@ class RequestDispatcherTest extends AgentTestRunner { dispatcher.include("") then: - 2 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 2 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(2) { trace(1) { basicSpan(it, "forward-child") @@ -44,7 +44,7 @@ class RequestDispatcherTest extends AgentTestRunner { } then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(1) { trace(3) { basicSpan(it, "parent") @@ -67,11 +67,13 @@ class RequestDispatcherTest extends AgentTestRunner { 1 * request.setAttribute(SPAN_ID_KEY, _) 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) >> mockContext then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName.toString() == "servlet.$operation" }) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, { Context ctx -> + ctx != null && ctx != mockContext // Verify it's a new context + }) then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, mockContext) where: operation | method @@ -97,7 +99,7 @@ class RequestDispatcherTest extends AgentTestRunner { def th = thrown(ServletException) th == ex - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(1) { trace(3) { basicSpan(it, "parent", null, ex) @@ -122,11 +124,13 @@ class RequestDispatcherTest extends AgentTestRunner { 1 * request.setAttribute(SPAN_ID_KEY, _) 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) >> mockContext then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName.toString() == "servlet.$operation" }) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, { Context ctx -> + ctx != null && ctx != mockContext // Verify it's a new context + }) then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, mockContext) where: operation | method diff --git a/dd-java-agent/instrumentation/spring-cloud-zuul-2/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java b/dd-java-agent/instrumentation/spring-cloud-zuul-2/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java index 77714dabea7..f167301e029 100644 --- a/dd-java-agent/instrumentation/spring-cloud-zuul-2/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-cloud-zuul-2/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.springcloudzuul2; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; import com.netflix.zuul.context.RequestContext; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -48,9 +50,10 @@ public static void onEnter( request = ctx.getRequest(); if (request != null) { // Capture the span from the request before forwarding. - Object span = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (span instanceof AgentSpan) { - parentSpan = (AgentSpan) span; + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + parentSpan = spanFromContext(context); } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy index 17188ac34d6..17c98e5f112 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy @@ -15,7 +15,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE @WebAppConfiguration // TODO this doesn't exist in 3.1.0 so will need rework if we want to test that version // https://mvnrepository.com/artifact/org.springframework/spring-test-mvc?repo=springio-plugins-release @@ -39,12 +39,12 @@ class HandlerMappingResourceNameFilterForkedTest extends AgentTestRunner { when: runUnderTrace("test-servlet", { - request.setAttribute(DD_SPAN_ATTRIBUTE, activeSpan()) + request.setAttribute(DD_CONTEXT_ATTRIBUTE, activeSpan().context()) filter.doFilterInternal(request, response, filterChain) }) then: - assert request.getAttributeNames().toList() == [DD_SPAN_ATTRIBUTE] + assert request.getAttributeNames().toList() == [DD_CONTEXT_ATTRIBUTE] 0 * response._ where: diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index 123fbc0a638..c9583a8eff8 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -7,7 +7,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; @@ -17,6 +18,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -73,9 +75,13 @@ public static AgentScope nameResourceAndStartSpan( handlerSpanKey = ""; // Name the parent span based on the matching pattern - Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - DECORATE.onRequest((AgentSpan) parentSpan, request, request, getRootContext()); + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + DECORATE.onRequest(parentSpan, request, request, getRootContext()); + } } if (activeSpan() == null) { diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java index f78268e61fc..073bce859f1 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.springweb; import static datadog.context.Context.root; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; import java.util.List; @@ -33,18 +35,22 @@ protected void doFilterInternal( final FilterChain filterChain) throws ServletException, IOException { - final Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - PathMatchingHttpServletRequestWrapper wrappedRequest = - new PathMatchingHttpServletRequestWrapper(request); - try { - if (findMapping(wrappedRequest)) { - // Name the parent span based on the matching pattern - // Let the parent span resource name be set with the attribute set in findMapping. - DECORATE.onRequest((AgentSpan) parentSpan, wrappedRequest, wrappedRequest, root()); + final Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + PathMatchingHttpServletRequestWrapper wrappedRequest = + new PathMatchingHttpServletRequestWrapper(request); + try { + if (findMapping(wrappedRequest)) { + // Name the parent span based on the matching pattern + // Let the parent span resource name be set with the attribute set in findMapping. + DECORATE.onRequest(parentSpan, wrappedRequest, wrappedRequest, root()); + } + } catch (final Exception ignored) { + // mapping.getHandler() threw exception. Ignore } - } catch (final Exception ignored) { - // mapping.getHandler() threw exception. Ignore } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy index 17188ac34d6..17c98e5f112 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy @@ -15,7 +15,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE @WebAppConfiguration // TODO this doesn't exist in 3.1.0 so will need rework if we want to test that version // https://mvnrepository.com/artifact/org.springframework/spring-test-mvc?repo=springio-plugins-release @@ -39,12 +39,12 @@ class HandlerMappingResourceNameFilterForkedTest extends AgentTestRunner { when: runUnderTrace("test-servlet", { - request.setAttribute(DD_SPAN_ATTRIBUTE, activeSpan()) + request.setAttribute(DD_CONTEXT_ATTRIBUTE, activeSpan().context()) filter.doFilterInternal(request, response, filterChain) }) then: - assert request.getAttributeNames().toList() == [DD_SPAN_ATTRIBUTE] + assert request.getAttributeNames().toList() == [DD_CONTEXT_ATTRIBUTE] 0 * response._ where: diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index 8d56521fce0..78f3b2388e6 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -4,11 +4,13 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.http.HttpServletRequest; @@ -24,9 +26,13 @@ public static AgentScope nameResourceAndStartSpan( @Advice.Local("handlerSpanKey") String handlerSpanKey) { handlerSpanKey = ""; // Name the parent span based on the matching pattern - Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - DECORATE.onRequest((AgentSpan) parentSpan, request, request, root()); + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + DECORATE.onRequest(parentSpan, request, request, root()); + } } if (activeSpan() == null) { diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java index 0b5008f5e00..eac11f69ca2 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.springweb6; import static datadog.context.Context.root; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.FilterChain; import jakarta.servlet.ServletException; @@ -33,18 +35,22 @@ protected void doFilterInternal( final FilterChain filterChain) throws ServletException, IOException { - final Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - PathMatchingHttpServletRequestWrapper wrappedRequest = - new PathMatchingHttpServletRequestWrapper(request); - try { - if (findMapping(wrappedRequest)) { - // Name the parent span based on the matching pattern - // Let the parent span resource name be set with the attribute set in findMapping. - DECORATE.onRequest((AgentSpan) parentSpan, wrappedRequest, wrappedRequest, root()); + final Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + PathMatchingHttpServletRequestWrapper wrappedRequest = + new PathMatchingHttpServletRequestWrapper(request); + try { + if (findMapping(wrappedRequest)) { + // Name the parent span based on the matching pattern + // Let the parent span resource name be set with the attribute set in findMapping. + DECORATE.onRequest(parentSpan, wrappedRequest, wrappedRequest, root()); + } + } catch (final Exception ignored) { + // mapping.getHandler() threw exception. Ignore } - } catch (final Exception ignored) { - // mapping.getHandler() threw exception. Ignore } } diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index f4a849c9371..ba3e5ad848f 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -6,7 +6,6 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; @@ -132,7 +131,6 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque final AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); @@ -163,19 +161,23 @@ public static void afterParse( @Advice.Argument(1) Request req, @Advice.Argument(3) Response resp, @Advice.Return(readOnly = false) Boolean ret) { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - AgentSpan span = (AgentSpan) spanObj; - req.setAttribute(CorrelationIdentifier.getTraceIdKey(), AgentTracer.get().getTraceId(span)); - req.setAttribute(CorrelationIdentifier.getSpanIdKey(), AgentTracer.get().getSpanId(span)); - Object ctxObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE); - Context parentContext = ctxObj instanceof Context ? (Context) ctxObj : getRootContext(); - DECORATE.onRequest(span, req, req, parentContext); - Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); - if (rba != null) { - TomcatBlockingHelper.commitBlockingResponse( - span.getRequestContext().getTraceSegment(), req, resp, rba); - ret = false; // skip pipeline + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + req.setAttribute( + CorrelationIdentifier.getTraceIdKey(), AgentTracer.get().getTraceId(span)); + req.setAttribute(CorrelationIdentifier.getSpanIdKey(), AgentTracer.get().getSpanId(span)); + Object ctxObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE); + Context parentContext = ctxObj instanceof Context ? (Context) ctxObj : getRootContext(); + DECORATE.onRequest(span, req, req, parentContext); + Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); + if (rba != null) { + TomcatBlockingHelper.commitBlockingResponse( + span.getRequestContext().getTraceSegment(), req, resp, rba); + ret = false; // skip pipeline + } } } } diff --git a/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java b/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java index 670c877ffad..8583b92156c 100644 --- a/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-appsec-7/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java @@ -1,11 +1,13 @@ package datadog.trace.instrumentation.tomcat7; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.gateway.BlockResponseFunction; @@ -92,7 +94,12 @@ static class ProcessCommitActionAdvice { request.setAttribute(HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE, Boolean.TRUE); // on async requests AgentTracer.activeSpan() may return null - AgentSpan agentSpan = (AgentSpan) request.getAttribute(HttpServerDecorator.DD_SPAN_ATTRIBUTE); + Object contextObj = request.getAttribute(HttpServerDecorator.DD_CONTEXT_ATTRIBUTE); + if (!(contextObj instanceof Context)) { + return false; + } + Context context = (Context) contextObj; + AgentSpan agentSpan = spanFromContext(context); if (agentSpan == null) { return false; } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java index 9a63d2c947c..2787580733f 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java @@ -3,7 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.SERVLET_REQUEST; @@ -68,7 +68,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_SPAN_ATTRIBUTE, undertowSpan); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, undertowSpan.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java index 58fb7553e76..6f3dcc66af4 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java @@ -3,7 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.SERVLET_REQUEST; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -60,7 +60,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_SPAN_ATTRIBUTE, undertowSpan); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, undertowSpan.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java index b2ffddc5bd1..4c43b14357d 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java @@ -2,6 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.httpurlconnection.HttpUrlConnectionDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPackagePrivate; @@ -9,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -57,8 +59,9 @@ public static void onExit( @Advice.FieldValue("stream") final HttpClientStream stream, @Advice.Return boolean result) { if (result) { - AgentSpan nettySpan = - stream.connection().channel().attr(AttributeKeys.SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = + stream.connection().channel().attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan nettySpan = storedContext != null ? spanFromContext(storedContext) : null; if (nettySpan != null) { try (final AgentScope scope = activateSpan(nettySpan)) { DECORATE.onError(scope, cause); From 0f0f228bacdc4a0335200c9164bdde6475476126 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 17:50:52 -0400 Subject: [PATCH 17/27] fix tests --- .../trace/instrumentation/jetty11/SetContextPathAdvice.java | 4 ++-- .../trace/instrumentation/jetty12/SetContextPathAdvice.java | 4 ++-- .../instrumentation/jetty70/RequestInstrumentation.java | 6 +++--- .../instrumentation/jetty76/RequestInstrumentation.java | 6 +++--- .../instrumentation/jetty9/RequestInstrumentation.java | 6 +++--- .../trace/instrumentation/jetty10/SetContextPathAdvice.java | 4 ++-- .../trace/instrumentation/jetty10/SetServletPathAdvice.java | 4 ++-- .../instrumentation/undertow/ServletInstrumentation.java | 2 +- .../undertow/JakartaServletInstrumentation.java | 2 +- 9 files changed, 19 insertions(+), 19 deletions(-) diff --git a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java index dfb65e520ac..e8007c38d23 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java @@ -1,8 +1,8 @@ package datadog.trace.instrumentation.jetty11; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -28,7 +28,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context ctx = (Context) contextObj; - AgentSpan span = fromContext(ctx); + AgentSpan span = spanFromContext(ctx); if (span != null) { if (context != null && context.getContextPath() != null) { final String servletContext = context.getContextPath(); diff --git a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java index 8dddbb0e8bc..6b1eb97e6b5 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java @@ -1,8 +1,8 @@ package datadog.trace.instrumentation.jetty12; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty12.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -27,7 +27,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { if (contextHandler != null && contextHandler.getContextPath() != null) { final String servletContext = contextHandler.getContextPath(); diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java index b7d56329c93..a79b766affd 100644 --- a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -68,7 +68,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_CONTEXT, contextPath); req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); @@ -91,7 +91,7 @@ public static void updateServletPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_PATH, servletPath); req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java index eea7c531767..346ec1b0e84 100644 --- a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -68,7 +68,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_CONTEXT, contextPath); req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); @@ -91,7 +91,7 @@ public static void updateServletPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_PATH, servletPath); req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java index d44c4cc52ab..996596c4008 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -75,7 +75,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_CONTEXT, contextPath); req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); @@ -98,7 +98,7 @@ public static void updateServletPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_PATH, servletPath); req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java index 2f2522a5bdd..6c0540843d9 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; @@ -24,7 +24,7 @@ public static void updateContextPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_CONTEXT, contextPath); req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java index 94b657b22ca..af307de6144 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; @@ -27,7 +27,7 @@ public static void updateServletPath( // Don't want to update while being dispatched to new servlet if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { Context context = (Context) contextObj; - AgentSpan span = fromContext(context); + AgentSpan span = spanFromContext(context); if (span != null) { span.setTag(SERVLET_PATH, servletPath); req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java index 2787580733f..b97a0a4c3b8 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java @@ -68,7 +68,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_CONTEXT_ATTRIBUTE, undertowSpan.context()); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, continuation.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java index 6f3dcc66af4..ab5b64fc3b5 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java @@ -60,7 +60,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_CONTEXT_ATTRIBUTE, undertowSpan.context()); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, continuation.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); From ecc5049ac696b4f8e5b31f478db9cb5170af63c6 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 27 Aug 2025 18:57:05 -0400 Subject: [PATCH 18/27] fix tests --- .../request-3/src/test/groovy/JettyServlet3Test.groovy | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/JettyServlet3Test.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/JettyServlet3Test.groovy index 11bf3eb20e7..312caadca48 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/JettyServlet3Test.groovy +++ b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/JettyServlet3Test.groovy @@ -9,6 +9,10 @@ import datadog.trace.instrumentation.servlet3.AsyncDispatcherDecorator import datadog.trace.instrumentation.servlet3.HtmlRumServlet import datadog.trace.instrumentation.servlet3.TestServlet3 import datadog.trace.instrumentation.servlet3.XmlRumServlet +import datadog.context.Context + +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE import groovy.servlet.AbstractHttpServlet import org.eclipse.jetty.server.Request import org.eclipse.jetty.server.Server @@ -526,7 +530,8 @@ class ServeFromOnAsyncTimeout extends AbstractHttpServlet { @Override void onTimeout(AsyncEvent event) throws IOException { - AgentSpan span = event.getSuppliedRequest().getAttribute('datadog.span') + Context ddContext = (Context) event.getSuppliedRequest().getAttribute(DD_CONTEXT_ATTRIBUTE) + AgentSpan span = spanFromContext(ddContext) activateSpan(span).withCloseable { try { delegateServlet.service(req, resp) From 4bd005dc97cb74b57e5df890412fb4e03a6e7ebc Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 28 Aug 2025 14:17:54 -0400 Subject: [PATCH 19/27] implement fixes part 1 --- .../decorator/BaseDecorator.java | 4 --- .../decorator/WebsocketDecorator.java | 4 +++ .../cxf/InvokerInstrumentation.java | 14 +++------ .../DefaultFilterChainInstrumentation.java | 15 +++------- .../grizzlyhttp232/FilterAdvice.java | 18 ++++------- .../grizzlyhttp232/GrizzlyDecorator.java | 30 ++++++++----------- .../v1/JaxRsClientV1Instrumentation.java | 8 ++--- .../instrumentation/mule4/MuleDecorator.java | 4 +++ 8 files changed, 35 insertions(+), 62 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java index 9cde0e0e244..fc13585c1d1 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java @@ -82,10 +82,6 @@ public ContextScope beforeFinish(final ContextScope scope) { return scope; } - public AgentSpan beforeFinish(final AgentSpan span) { - return span; - } - public Context beforeFinish(final Context context) { return context; } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java index 2123caa395d..7d9b3ba8a15 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -61,6 +61,10 @@ public AgentSpan afterStart(AgentSpan span) { return super.afterStart(span).setMeasured(true); } + public AgentSpan beforeFinish(final AgentSpan span) { + return span; + } + @Nonnull public AgentSpan onReceiveFrameStart( final HandlerContext.Receiver handlerContext, final Object data, boolean partialDelivery) { diff --git a/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java b/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java index 2893f6ac80f..775fe5fced3 100644 --- a/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java +++ b/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java @@ -2,15 +2,13 @@ import com.google.auto.service.AutoService; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers; import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers; import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; -import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; @@ -57,7 +55,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class PropagateSpanAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) { + public static ContextScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) { if (exchange == null || exchange.getInMessage() == null || AgentTracer.activeSpan() != null) { return null; } @@ -66,17 +64,13 @@ public static AgentScope beforeInvoke(@Advice.Argument(0) final Exchange exchang exchange.getInMessage().get("HTTP.REQUEST"), HttpServerDecorator.DD_CONTEXT_ATTRIBUTE); if (contextObj instanceof Context) { - Context context = (Context) contextObj; - AgentSpan span = Java8BytecodeBridge.spanFromContext(context); - if (span != null) { - return AgentTracer.activateSpan(span); - } + return ((Context) contextObj).attach(); } return null; } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void afterInvoke(@Advice.Enter final AgentScope scope) { + public static void afterInvoke(@Advice.Enter final ContextScope scope) { if (scope != null) { scope.close(); } diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java index 48beb78cfdb..1e97c55b7de 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java @@ -1,9 +1,7 @@ package datadog.trace.instrumentation.grizzlyhttp232; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPrivate; @@ -11,10 +9,10 @@ import com.google.auto.service.AutoService; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.Collections; import net.bytebuddy.asm.Advice; @@ -69,7 +67,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class PropagateServerSpanAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(2) final FilterChainContext ctx) { + public static ContextScope onEnter(@Advice.Argument(2) final FilterChainContext ctx) { final AgentSpan active = activeSpan(); // don't activate a span if already one is active if (active != null) { @@ -77,18 +75,13 @@ public static AgentScope onEnter(@Advice.Argument(2) final FilterChainContext ct } final Object contextObj = ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); if (contextObj instanceof Context) { - final Context context = (Context) contextObj; - final AgentSpan span = spanFromContext(context); - if (span != null) { - // activate the http server span when nothing is already active - return activateSpan(span); - } + return ((Context) contextObj).attach(); } return null; } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void onExit(@Advice.Enter final AgentScope scope) { + public static void onExit(@Advice.Enter final ContextScope scope) { if (scope != null) { scope.close(); } diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java index 4a404e4106c..3b8946984c8 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java @@ -1,36 +1,28 @@ package datadog.trace.instrumentation.grizzlyhttp232; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import datadog.context.Context; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.context.ContextScope; import net.bytebuddy.asm.Advice; import org.glassfish.grizzly.filterchain.FilterChainContext; public class FilterAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) { - if (activeSpan() != null) { + public static ContextScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) { + if (Context.current() != Context.root()) { return null; } Object contextObj = ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); if (contextObj instanceof Context) { - final Context context = (Context) contextObj; - final AgentSpan span = spanFromContext(context); - if (span != null) { - return activateSpan(span); - } + return ((Context) contextObj).attach(); } return null; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void onExit(@Advice.Enter final AgentScope scope) { + public static void onExit(@Advice.Enter final ContextScope scope) { if (scope != null) { scope.close(); } diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index f6386a0a118..9c75322a1aa 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -89,23 +89,19 @@ protected CharSequence component() { public static void onHttpServerFilterPrepareResponseEnter( FilterChainContext ctx, HttpResponsePacket responsePacket) { Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); - if (context != null) { - AgentSpan span = spanFromContext(context); - if (span != null) { - DECORATE.onResponse(span, responsePacket); - } + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { + DECORATE.onResponse(span, responsePacket); } } public static void onHttpServerFilterPrepareResponseExit( FilterChainContext ctx, HttpResponsePacket responsePacket) { Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); - if (context != null) { - AgentSpan span = spanFromContext(context); - if (span != null) { - DECORATE.beforeFinish(context); - span.finish(); - } + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { + DECORATE.beforeFinish(context); + span.finish(); } ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); @@ -153,13 +149,11 @@ public static NextAction onHttpCodecFilterExit( public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable) { Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); - if (context != null) { - AgentSpan span = spanFromContext(context); - if (span != null) { - DECORATE.onError(span, throwable); - DECORATE.beforeFinish(context); - span.finish(); - } + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(context); + span.finish(); } ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); diff --git a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java index 415305b2b4e..63f1e9a1de5 100644 --- a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java +++ b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java @@ -72,13 +72,9 @@ public static AgentScope onEnter( final AgentSpan span = startSpan(JAX_RS_CLIENT_CALL); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - final AgentScope scope = activateSpan(span); - // Store the context after activation to ensure we have the proper Context object - if (scope != null) { - request.getProperties().put(DD_CONTEXT_ATTRIBUTE, scope.context()); - } + request.getProperties().put(DD_CONTEXT_ATTRIBUTE, span); DECORATE.injectContext(getCurrentContext().with(span), request.getHeaders(), SETTER); - return scope; + return activateSpan(span); } return null; } diff --git a/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java b/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java index ecd17b30a15..b06e15277f6 100644 --- a/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java +++ b/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java @@ -57,6 +57,10 @@ public AgentSpan afterStart(final AgentSpan span) { return super.afterStart(span); } + public AgentSpan beforeFinish(final AgentSpan span) { + return span; + } + public AgentSpan onMuleSpan( AgentSpan parentSpan, InitialSpanInfo spanInfo, CoreEvent event, Component component) { // we stick with the same level of detail of OTEL exporter. From 9b4d887cb8229018f3e9ec4c76780d3866ba0145 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 28 Aug 2025 16:18:37 -0400 Subject: [PATCH 20/27] fix tests --- .../bootstrap/instrumentation/decorator/BaseDecorator.java | 4 ++++ .../instrumentation/decorator/WebsocketDecorator.java | 4 ---- .../trace/instrumentation/grizzlyhttp232/FilterAdvice.java | 4 +++- .../datadog/trace/instrumentation/mule4/MuleDecorator.java | 4 ---- 4 files changed, 7 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java index fc13585c1d1..9cde0e0e244 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java @@ -82,6 +82,10 @@ public ContextScope beforeFinish(final ContextScope scope) { return scope; } + public AgentSpan beforeFinish(final AgentSpan span) { + return span; + } + public Context beforeFinish(final Context context) { return context; } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java index 7d9b3ba8a15..2123caa395d 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/WebsocketDecorator.java @@ -61,10 +61,6 @@ public AgentSpan afterStart(AgentSpan span) { return super.afterStart(span).setMeasured(true); } - public AgentSpan beforeFinish(final AgentSpan span) { - return span; - } - @Nonnull public AgentSpan onReceiveFrameStart( final HandlerContext.Receiver handlerContext, final Object data, boolean partialDelivery) { diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java index 3b8946984c8..9e39339d732 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java @@ -1,5 +1,7 @@ package datadog.trace.instrumentation.grizzlyhttp232; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import datadog.context.Context; @@ -11,7 +13,7 @@ public class FilterAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) { - if (Context.current() != Context.root()) { + if (getCurrentContext() != getRootContext()) { return null; } Object contextObj = ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); diff --git a/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java b/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java index b06e15277f6..ecd17b30a15 100644 --- a/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java +++ b/dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/MuleDecorator.java @@ -57,10 +57,6 @@ public AgentSpan afterStart(final AgentSpan span) { return super.afterStart(span); } - public AgentSpan beforeFinish(final AgentSpan span) { - return span; - } - public AgentSpan onMuleSpan( AgentSpan parentSpan, InitialSpanInfo spanInfo, CoreEvent event, Component component) { // we stick with the same level of detail of OTEL exporter. From ee71b23e6fc2efa16afd9f559288ea46b924eb71 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Fri, 3 Oct 2025 22:28:37 -0400 Subject: [PATCH 21/27] address jetty comments --- .../decorator/HttpServerDecorator.java | 2 +- .../jetty70/ServerHandleInstrumentation.java | 6 ----- .../jetty76/ServerHandleInstrumentation.java | 6 ----- .../JettyCommitResponseInstrumentation.java | 27 ++++++------------- .../jetty9/ServerHandleInstrumentation.java | 6 ----- .../jetty10/JettyCommitResponseHelper.java | 27 ++++++------------- .../jetty10/ServerHandleAdvice.java | 6 ----- .../jetty904/JettyCommitResponseHelper.java | 27 ++++++------------- .../jetty93/JettyCommitResponseHelper.java | 27 ++++++------------- .../jetty9421/JettyCommitResponseHelper.java | 27 ++++++------------- 10 files changed, 41 insertions(+), 120 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 1701a4a31f1..e2a53ee7e8c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -542,7 +542,7 @@ public Context beforeFinish(Context context) { onRequestEndForInstrumentationGateway(span); // Close Serverless Gateway Inferred Span if any - // finishInferredProxySpan(context); + finishInferredProxySpan(context); return super.beforeFinish(context); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java index 426c856a589..0acc3649f3f 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.jetty70; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; @@ -95,11 +94,6 @@ public static void onExit( return; } - // Extract the span from the context if we didn't get it from @Local - if (span == null) { - span = spanFromContext(scope.context()); - } - if (t != null) { DECORATE.onError(span, t); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java index 5faff27ff31..ae8515aec21 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.jetty76; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; @@ -95,11 +94,6 @@ public static void onExit( return; } - // Extract the span from the context if we didn't get it from @Local - if (span == null) { - span = spanFromContext(scope.context()); - } - if (t != null) { DECORATE.onError(span, t); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java index 061db8894b9..a5be4c07b16 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java @@ -100,25 +100,14 @@ static class CommitResponseAdvice { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (!(contextObj instanceof Context)) { - _committed.set(false); - return false; - } - Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - _committed.set(false); - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java index 0e79a99ab19..590390d8404 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.jetty9; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; @@ -124,11 +123,6 @@ public static void onExit( return; } - // Extract the span from the context if we didn't get it from @Local - if (span == null) { - span = spanFromContext(scope.context()); - } - boolean registeredFinishListener; synchronized (req) { req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java index 836ee29aae7..4dd2494355a 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java @@ -51,25 +51,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - state.partialResponse(); - return false; - } - - Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (!(contextObj instanceof Context)) { - state.partialResponse(); - return false; - } - Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - state.partialResponse(); - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { state.partialResponse(); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java index 033751b9418..b121af3f8eb 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java @@ -1,6 +1,5 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; @@ -62,11 +61,6 @@ public static void onExit( return; } - // Extract the span from the context if we didn't get it from @Local - if (span == null) { - span = spanFromContext(scope.context()); - } - if (t != null) { DECORATE.onError(span, t); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java index 92bae744f84..22548b59291 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty904/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java @@ -50,25 +50,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (!(contextObj instanceof Context)) { - _committed.set(false); - return false; - } - Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - _committed.set(false); - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java index bd512c5d406..5e4dc5fa078 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty93/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java @@ -50,25 +50,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (!(contextObj instanceof Context)) { - _committed.set(false); - return false; - } - Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - _committed.set(false); - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java index 61c9815b3be..ba34bd6412d 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java_jetty9421/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java @@ -53,25 +53,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - state.partialResponse(); - return false; - } - - Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (!(contextObj instanceof Context)) { - state.partialResponse(); - return false; - } - Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - state.partialResponse(); - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { state.partialResponse(); return false; } From 2c27431643e077c91da820495e12715466d569f2 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Sat, 25 Oct 2025 14:40:10 -0400 Subject: [PATCH 22/27] use context scope --- .../akkahttp/DatadogAsyncHandlerWrapper.java | 12 +++++---- .../axway/HTTPPluginAdvice.java | 16 +++++++----- .../instrumentation/axway/StateAdvice.java | 4 ++- .../finatra/FinatraInstrumentation.java | 12 ++++----- .../instrumentation/finatra/Listener.java | 22 +++++++++------- .../JettyCommitResponseInstrumentation.java | 10 +++---- .../JettyCommitResponseInstrumentation.java | 10 +++---- .../WriteFinalNettyResponseAdvice.java | 5 ++-- .../ChannelFutureListenerInstrumentation.java | 4 +-- .../HttpServerResponseTracingHandler.java | 5 ++-- .../ChannelFutureListenerInstrumentation.java | 4 +-- ...yChannelHandlerContextInstrumentation.java | 2 +- .../HttpClientRequestTracingHandler.java | 2 +- .../HttpClientResponseTracingHandler.java | 20 ++++++-------- .../HttpServerRequestTracingHandler.java | 12 ++++----- .../HttpServerResponseTracingHandler.java | 7 +++-- .../server/MaybeBlockResponseHandler.java | 2 +- .../ChannelFutureListenerInstrumentation.java | 4 +-- ...yChannelHandlerContextInstrumentation.java | 2 +- .../HttpClientRequestTracingHandler.java | 2 +- .../HttpClientResponseTracingHandler.java | 26 +++++++------------ .../HttpServerRequestTracingHandler.java | 10 +++---- .../HttpServerResponseTracingHandler.java | 7 +++-- .../server/MaybeBlockResponseHandler.java | 2 +- .../DispatcherServletInstrumentation.java | 13 +++++----- .../HandlerAdapterInstrumentation.java | 20 +++++++------- .../springweb6/ControllerAdvice.java | 20 +++++++------- .../springweb6/RenderAdvice.java | 13 +++++----- .../HttpClientRequestBaseInstrumentation.java | 2 +- 29 files changed, 130 insertions(+), 140 deletions(-) diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java index 20df10cd40f..91850d3621f 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java @@ -6,6 +6,7 @@ import akka.http.scaladsl.model.HttpResponse; import akka.http.scaladsl.util.FastFuture$; import akka.stream.Materializer; +import datadog.context.Context; import datadog.context.ContextScope; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -29,7 +30,8 @@ public DatadogAsyncHandlerWrapper( @Override public Future apply(final HttpRequest request) { final ContextScope scope = DatadogWrapperHelper.createSpan(request); - final AgentSpan span = fromContext(scope.context()); + final Context context = scope.context(); + final AgentSpan span = fromContext(context); Future futureResponse; // handle blocking in the beginning of the request @@ -38,7 +40,7 @@ public Future apply(final HttpRequest request) { request.discardEntityBytes(materializer); HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request); span.getRequestContext().getTraceSegment().effectivelyBlocked(); - DatadogWrapperHelper.finishSpan(scope.context(), response); + DatadogWrapperHelper.finishSpan(context, response); return FastFuture$.MODULE$.successful().apply(response); } @@ -46,7 +48,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(scope.context(), t); + DatadogWrapperHelper.finishSpan(context, t); throw t; } @@ -67,14 +69,14 @@ public HttpResponse apply(HttpResponse response) { response = newResponse; } - DatadogWrapperHelper.finishSpan(scope.context(), response); + DatadogWrapperHelper.finishSpan(context, response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(scope.context(), t); + DatadogWrapperHelper.finishSpan(context, t); return t; } }, diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java index d03f8e16dba..e6cdc9c5d71 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java @@ -1,36 +1,38 @@ package datadog.trace.instrumentation.axway; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE; import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.SERVER_TRANSACTION_CLASS; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; public class HTTPPluginAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(value = 2) final Object serverTransaction) { + public static ContextScope onEnter(@Advice.Argument(value = 2) final Object serverTransaction) { final AgentSpan span = startSpan("axway-api", DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); // serverTransaction is like request + connection in one object: DECORATE.onRequest(span, serverTransaction, serverTransaction, getRootContext()); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Argument(value = 2) final Object serverTransaction, @Advice.Thrown final Throwable throwable) { if (scope == null) { return; } - final AgentSpan span = scope.span(); + final Context context = scope.context(); + final AgentSpan span = fromContext(context); try { if (null != serverTransaction) { // manual DECORATE.onResponse(span, serverTransaction): @@ -45,7 +47,7 @@ public static void onExit( if (throwable != null) { DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); } finally { scope.close(); span.finish(); diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java index fabff3d1346..9c63a5d31d9 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java @@ -5,6 +5,7 @@ import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.AXWAY_TRY_TRANSACTION; import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; @@ -34,11 +35,12 @@ public static void onExit( return; } final AgentSpan span = scope.span(); + final Context context = scope.context(); try { if (throwable != null) { DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(scope.context()); + DECORATE.beforeFinish(context); } finally { scope.close(); span.finish(); diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java index da9f9ce6426..372aa0f5fc5 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java @@ -3,9 +3,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE; import static datadog.trace.instrumentation.finatra.FinatraDecorator.FINATRA_CONTROLLER; @@ -17,9 +17,9 @@ import com.twitter.finagle.http.Request; import com.twitter.finagle.http.Response; import com.twitter.util.Future; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.Tags; import net.bytebuddy.asm.Advice; @@ -61,7 +61,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class RouteAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope nameSpan( + public static ContextScope nameSpan( @Advice.Argument(0) final Request request, @Advice.FieldValue("path") final String path, @Advice.FieldValue("clazz") final Class clazz) { @@ -76,12 +76,12 @@ public static AgentScope nameSpan( DECORATE.afterStart(span); span.setResourceName(DECORATE.className(clazz)); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void setupCallback( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable, @Advice.Return final Some> responseOption) { @@ -89,7 +89,7 @@ public static void setupCallback( return; } - final AgentSpan span = scope.span(); + final AgentSpan span = spanFromContext(scope.context()); if (throwable != null) { DECORATE.onError(span, throwable); DECORATE.beforeFinish(scope.context()); diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java index 677ede1884f..6d948f86532 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java @@ -1,36 +1,40 @@ package datadog.trace.instrumentation.finatra; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE; import com.twitter.finagle.http.Response; import com.twitter.util.FutureEventListener; +import datadog.context.ContextScope; import datadog.trace.api.Config; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; public class Listener implements FutureEventListener { - private final AgentScope scope; + private final ContextScope scope; - public Listener(final AgentScope scope) { + public Listener(final ContextScope scope) { this.scope = scope; } @Override public void onSuccess(final Response response) { + final AgentSpan span = spanFromContext(scope.context()); // Don't use DECORATE.onResponse because this is the controller span if (Config.get().getHttpServerErrorStatuses().get(DECORATE.status(response))) { - scope.span().setError(true); + span.setError(true); } - DECORATE.beforeFinish(scope.span()); - scope.span().finish(); + DECORATE.beforeFinish(scope.context()); + span.finish(); scope.close(); } @Override public void onFailure(final Throwable cause) { - DECORATE.onError(scope.span(), cause); - DECORATE.beforeFinish(scope.span()); - scope.span().finish(); + final AgentSpan span = spanFromContext(scope.context()); + DECORATE.onError(span, cause); + DECORATE.beforeFinish(scope.context()); + span.finish(); scope.close(); } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java index fc44ad0d208..f5311a224ec 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java @@ -77,13 +77,9 @@ static class CommitResponseAdvice { return false; } Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + AgentSpan span = spanFromContext(context); + RequestContext requestContext; + if (span == null || (requestContext = span.getRequestContext()) == null) { return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java index ad9e1d94ed8..7f6dbac9c40 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java @@ -88,13 +88,9 @@ static class CommitResponseAdvice { return false; } Context context = (Context) contextObj; - AgentSpan existingSpan = spanFromContext(context); - if (existingSpan == null) { - return false; - } - AgentSpan span = existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + AgentSpan span = spanFromContext(context); + RequestContext requestContext; + if (span == null || (requestContext = span.getRequestContext()) == null) { return false; } diff --git a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java index 33dbfef94ac..00b11ea54d6 100644 --- a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java +++ b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java @@ -1,10 +1,9 @@ package datadog.trace.instrumentation.micronaut; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.DECORATE; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.SPAN_ATTRIBUTE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.micronaut.http.HttpRequest; import io.micronaut.http.MutableHttpResponse; @@ -20,7 +19,7 @@ public static void beginRequest( return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = span.attach()) { DECORATE.onResponse(span, message); DECORATE.beforeFinish(scope.context()); span.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 6836a3bb72e..1f9dcc906cd 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -2,7 +2,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.NETTY; @@ -11,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.ContextStore; @@ -109,7 +109,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = errorSpan.attach()) { DECORATE.onError(errorSpan, cause); DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index 2b77f94c660..176602faa03 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -1,10 +1,9 @@ package datadog.trace.instrumentation.netty38.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.ContextScope; import datadog.trace.bootstrap.ContextStore; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import datadog.trace.instrumentation.netty38.ChannelTraceContext; @@ -36,7 +35,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = span.attach()) { final HttpResponse response = (HttpResponse) msg.getMessage(); try { diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java index 255fb792bab..27e83d32c74 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java @@ -2,7 +2,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.NETTY; @@ -11,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -94,7 +94,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = errorSpan.attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java index 76ca4218b06..7cc2884896d 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java @@ -72,7 +72,7 @@ public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan channelSpan = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan channelSpan = spanFromContext(storedContext); if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index 541b06c0f96..affab2e3783 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -90,7 +90,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(scope.context()); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index 6819f817fe1..fd2a31b9654 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -27,12 +27,10 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); final boolean finishSpan = msg instanceof HttpResponse; @@ -56,12 +54,11 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -83,12 +80,11 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java index 5bb70729845..d627282b43b 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java @@ -28,7 +28,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { Channel channel = ctx.channel(); if (!(msg instanceof HttpRequest)) { final Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); if (span == null) { ctx.fireChannelRead(msg); // superclass does not throw } else { @@ -83,12 +83,10 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { } finally { try { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); - if (storedContext != null) { - final AgentSpan span = spanFromContext(storedContext); - if (span != null && span.phasedFinish()) { - // at this point we can just publish this span to avoid loosing the rest of the trace - span.publish(); - } + final AgentSpan span = spanFromContext(storedContext); + if (span != null && span.phasedFinish()) { + // at this point we can just publish this span to avoid loosing the rest of the trace + span.publish(); } } catch (final Throwable ignored) { } diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java index 462e1996cd4..6244bea495f 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java @@ -1,6 +1,5 @@ package datadog.trace.instrumentation.netty40.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CHANNEL_ID; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; @@ -8,7 +7,7 @@ import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import io.netty.channel.ChannelHandler; @@ -27,13 +26,13 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = storedContext.attach()) { final HttpResponse response = (HttpResponse) msg; try { diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java index a0c01ab5944..c5abbe0c807 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java @@ -59,7 +59,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr Channel channel = ctx.channel(); Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); - AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + AgentSpan span = spanFromContext(storedContext); RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java index a552db7704b..491cc796855 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java @@ -2,7 +2,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.NETTY; @@ -11,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -94,7 +94,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = errorSpan.attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java index 075bde792ba..62363fda0cc 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java @@ -72,7 +72,7 @@ public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan channelSpan = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan channelSpan = spanFromContext(storedContext); if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index 24a253f773e..6c02c7311cd 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -91,7 +91,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(scope.context()); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 32c5408e9d9..318caed9d02 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -29,12 +29,10 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); if (span != null) { final boolean finishSpan = @@ -49,9 +47,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - try (final AgentScope spanScope = activateSpan(span)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); } } @@ -67,12 +63,11 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -94,12 +89,11 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); - // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final AgentScope parentScope = activateSpan(parent)) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); - } if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java index 3d6b324ea6e..58892004518 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java @@ -87,12 +87,10 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { } finally { try { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); - if (storedContext != null) { - final AgentSpan span = spanFromContext(storedContext); - if (span != null && span.phasedFinish()) { - // at this point we can just publish this span to avoid loosing the rest of the trace - span.publish(); - } + final AgentSpan span = spanFromContext(storedContext); + if (span != null && span.phasedFinish()) { + // at this point we can just publish this span to avoid loosing the rest of the trace + span.publish(); } } catch (final Throwable ignored) { } diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java index 3290ed1c02b..1cf31f91eae 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java @@ -1,13 +1,12 @@ package datadog.trace.instrumentation.netty41.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import io.netty.channel.ChannelHandler; @@ -25,14 +24,14 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); - final AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + final AgentSpan span = spanFromContext(storedContext); if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = storedContext.attach()) { final HttpResponse response = (HttpResponse) msg; try { diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java index 17ab2f327ca..7503e871d65 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java @@ -58,7 +58,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr Channel channel = ctx.channel(); Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); - AgentSpan span = storedContext != null ? spanFromContext(storedContext) : null; + AgentSpan span = spanFromContext(storedContext); RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null) { super.write(ctx, msg, prm); diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java index b0945179d19..6f3b9a3209e 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE_RENDER; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.RESPONSE_RENDER; @@ -14,9 +14,9 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import java.util.List; @@ -95,20 +95,21 @@ public static void afterRefresh( public static class RenderAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { + public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(RESPONSE_RENDER); DECORATE_RENDER.afterStart(span); DECORATE_RENDER.onRender(span, mv); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + final AgentSpan span = spanFromContext(scope.context()); DECORATE_RENDER.onError(scope, throwable); DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index c9583a8eff8..fa3ac4ddd8a 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -3,7 +3,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; @@ -19,9 +18,9 @@ import com.google.auto.service.AutoService; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import javax.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; @@ -68,7 +67,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope nameResourceAndStartSpan( + public static ContextScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, @Advice.Argument(2) final Object handler, @Advice.Local("handlerSpanKey") String handlerSpanKey) { @@ -97,22 +96,25 @@ public static AgentScope nameResourceAndStartSpan( handlerKey = handler.getClass().getName(); } handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; - final Object existingSpan = request.getAttribute(handlerSpanKey); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + + // If the context already exists, return it + final Object existingContext = request.getAttribute(handlerSpanKey); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); DECORATE.onHandle(span, handler); + request.setAttribute(handlerSpanKey, span); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( @Advice.Argument(0) final HttpServletRequest request, - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable, @Advice.Local("handlerSpanKey") String handlerSpanKey) { if (scope == null) { @@ -121,7 +123,7 @@ public static void stopSpan( boolean finish = !Boolean.TRUE.equals( request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); - final AgentSpan span = scope.span(); + final AgentSpan span = spanFromContext(scope.context()); scope.close(); if (throwable != null) { DECORATE.onError(span, throwable); diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index 78f3b2388e6..e61f56273f1 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.springweb6; import static datadog.context.Context.root; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; @@ -11,7 +10,7 @@ import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; import datadog.context.Context; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; @@ -20,7 +19,7 @@ public class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope nameResourceAndStartSpan( + public static ContextScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, @Advice.Argument(2) final Object handler, @Advice.Local("handlerSpanKey") String handlerSpanKey) { @@ -48,21 +47,24 @@ public static AgentScope nameResourceAndStartSpan( handlerKey = handler.getClass().getName(); } handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; - final Object existingSpan = request.getAttribute(handlerSpanKey); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + + // If the context already exists, return it + final Object existingContext = request.getAttribute(handlerSpanKey); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); DECORATE.onHandle(span, handler); + request.setAttribute(handlerSpanKey, span); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Argument(0) final HttpServletRequest request, @Advice.Thrown final Throwable throwable, @Advice.Local("handlerSpanKey") String handlerSpanKey) { @@ -72,7 +74,7 @@ public static void stopSpan( boolean finish = !Boolean.TRUE.equals( request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); - final AgentSpan span = scope.span(); + final AgentSpan span = spanFromContext(scope.context()); scope.close(); if (throwable != null) { DECORATE.onError(span, throwable); diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java index a829cfdbeba..7f73c369b84 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java @@ -1,9 +1,9 @@ package datadog.trace.instrumentation.springweb6; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.springframework.web.servlet.ModelAndView; @@ -11,19 +11,20 @@ public class RenderAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { + public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(SpringWebHttpServerDecorator.RESPONSE_RENDER); SpringWebHttpServerDecorator.DECORATE_RENDER.afterStart(span); SpringWebHttpServerDecorator.DECORATE_RENDER.onRender(span, mv); - return activateSpan(span); + return span.attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + final AgentSpan span = spanFromContext(scope.context()); SpringWebHttpServerDecorator.DECORATE_RENDER.onError(scope, throwable); SpringWebHttpServerDecorator.DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java index 4c43b14357d..5c617ce0275 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java @@ -61,7 +61,7 @@ public static void onExit( if (result) { Context storedContext = stream.connection().channel().attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).get(); - AgentSpan nettySpan = storedContext != null ? spanFromContext(storedContext) : null; + AgentSpan nettySpan = spanFromContext(storedContext); if (nettySpan != null) { try (final AgentScope scope = activateSpan(nettySpan)) { DECORATE.onError(scope, cause); From 6e950a0d727902f8f81b0e7f05c8791ab73b2e8c Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Sat, 25 Oct 2025 15:19:22 -0400 Subject: [PATCH 23/27] fix netty tests --- .../client/HttpClientRequestTracingHandler.java | 8 +++++++- .../HttpClientResponseTracingHandler.java | 13 ++++++++++--- .../client/HttpClientRequestTracingHandler.java | 8 +++++++- .../HttpClientResponseTracingHandler.java | 17 +++++++++++++---- 4 files changed, 37 insertions(+), 9 deletions(-) diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index affab2e3783..89b59e990ba 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -12,6 +12,8 @@ import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty40.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -76,6 +78,10 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); + final Context spanContext; + try (final ContextScope contextScope = span.attach()) { + spanContext = contextScope.context(); + } try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -90,7 +96,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index fd2a31b9654..e272d4db846 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -30,7 +31,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } final boolean finishSpan = msg instanceof HttpResponse; @@ -57,7 +60,9 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the @@ -83,7 +88,9 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index 6c02c7311cd..1dbb653a420 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -12,6 +12,8 @@ import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty41.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -77,6 +79,10 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); + final Context spanContext; + try (final ContextScope contextScope = span.attach()) { + spanContext = contextScope.context(); + } try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -91,7 +97,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 318caed9d02..0d3ca11c4b6 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -32,7 +33,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { final boolean finishSpan = @@ -47,7 +50,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(span.context()); + try (final ContextScope contextScope = span.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(contextScope.context()); + } } } @@ -66,7 +71,9 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the @@ -92,7 +99,9 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parent.context()); + try (final ContextScope parentScope = parent.attach()) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + } if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { From fd11e7a1cb3971240afd01c5e501e6b7b690e305 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 30 Oct 2025 14:02:57 -0400 Subject: [PATCH 24/27] fix baggage context usage --- .../datadog/trace/instrumentation/axway/HTTPPluginAdvice.java | 3 ++- .../trace/instrumentation/finatra/FinatraInstrumentation.java | 3 ++- .../springweb/DispatcherServletInstrumentation.java | 3 ++- .../springweb/HandlerAdapterInstrumentation.java | 3 ++- .../trace/instrumentation/springweb6/ControllerAdvice.java | 3 ++- .../datadog/trace/instrumentation/springweb6/RenderAdvice.java | 3 ++- 6 files changed, 12 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java index e6cdc9c5d71..e887ffec7f1 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java @@ -2,6 +2,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE; import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.SERVER_TRANSACTION_CLASS; @@ -20,7 +21,7 @@ public static ContextScope onEnter(@Advice.Argument(value = 2) final Object serv DECORATE.afterStart(span); // serverTransaction is like request + connection in one object: DECORATE.onRequest(span, serverTransaction, serverTransaction, getRootContext()); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java index 1beae4b5c83..e93d0858dc6 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java @@ -5,6 +5,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE; @@ -78,7 +79,7 @@ public static ContextScope nameSpan( DECORATE.afterStart(span); span.setResourceName(DECORATE.className(clazz)); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java index 6f3b9a3209e..8fd5cb59219 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java @@ -4,6 +4,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE_RENDER; @@ -99,7 +100,7 @@ public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(RESPONSE_RENDER); DECORATE_RENDER.afterStart(span); DECORATE_RENDER.onRender(span, mv); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index fa3ac4ddd8a..6558878f8ad 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -5,6 +5,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; @@ -108,7 +109,7 @@ public static ContextScope nameResourceAndStartSpan( DECORATE.onHandle(span, handler); request.setAttribute(handlerSpanKey, span); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index e61f56273f1..bb95ff99f8f 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -3,6 +3,7 @@ import static datadog.context.Context.root; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; @@ -59,7 +60,7 @@ public static ContextScope nameResourceAndStartSpan( DECORATE.onHandle(span, handler); request.setAttribute(handlerSpanKey, span); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java index 7f73c369b84..542af56afda 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.springweb6; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import datadog.context.ContextScope; @@ -15,7 +16,7 @@ public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(SpringWebHttpServerDecorator.RESPONSE_RENDER); SpringWebHttpServerDecorator.DECORATE_RENDER.afterStart(span); SpringWebHttpServerDecorator.DECORATE_RENDER.onRender(span, mv); - return span.attach(); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) From b9ad08ce3c68372c698c5876c01c7f70e70e75c4 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Tue, 4 Nov 2025 17:37:34 -0500 Subject: [PATCH 25/27] fix context propagation --- .../WriteFinalNettyResponseAdvice.java | 3 ++- .../ChannelFutureListenerInstrumentation.java | 3 ++- .../HttpServerResponseTracingHandler.java | 3 ++- .../ChannelFutureListenerInstrumentation.java | 3 ++- .../HttpClientRequestTracingHandler.java | 3 ++- .../HttpClientResponseTracingHandler.java | 6 +++--- .../ChannelFutureListenerInstrumentation.java | 3 ++- .../HttpClientRequestTracingHandler.java | 3 ++- .../HttpClientResponseTracingHandler.java | 9 ++++---- .../RequestDispatcherInstrumentation.java | 21 +++++++++---------- .../SynapseServerWorkerInstrumentation.java | 16 ++++++++++---- 11 files changed, 44 insertions(+), 29 deletions(-) diff --git a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java index 00b11ea54d6..47802519796 100644 --- a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java +++ b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.micronaut; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.DECORATE; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.SPAN_ATTRIBUTE; @@ -19,7 +20,7 @@ public static void beginRequest( return; } - try (final ContextScope scope = span.attach()) { + try (final ContextScope scope = getCurrentContext().with(span).attach()) { DECORATE.onResponse(span, message); DECORATE.beforeFinish(scope.context()); span.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 1f9dcc906cd..34281d33a46 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -109,7 +110,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final ContextScope scope = errorSpan.attach()) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { DECORATE.onError(errorSpan, cause); DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index 176602faa03..984af432114 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.netty38.server; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; import datadog.context.ContextScope; @@ -35,7 +36,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m return; } - try (final ContextScope scope = span.attach()) { + try (final ContextScope scope = getCurrentContext().with(span).attach()) { final HttpResponse response = (HttpResponse) msg.getMessage(); try { diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java index 27e83d32c74..60834870a8c 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -94,7 +95,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final ContextScope scope = errorSpan.attach()) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index 89b59e990ba..df5b670aee4 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; @@ -79,7 +80,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); final Context spanContext; - try (final ContextScope contextScope = span.attach()) { + try (final ContextScope contextScope = getCurrentContext().with(span).attach()) { spanContext = contextScope.context(); } try (final AgentScope scope = activateSpan(span)) { diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index e272d4db846..5b27e6a0804 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -31,7 +31,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } @@ -60,7 +60,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } @@ -88,7 +88,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java index 491cc796855..f0d04a533d2 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -94,7 +95,7 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final ContextScope scope = errorSpan.attach()) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index 1dbb653a420..b4ebab59282 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; @@ -80,7 +81,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); final Context spanContext; - try (final ContextScope contextScope = span.attach()) { + try (final ContextScope contextScope = getCurrentContext().with(span).attach()) { spanContext = contextScope.context(); } try (final AgentScope scope = activateSpan(span)) { diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 0d3ca11c4b6..84d9a811c9a 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.netty41.client; +import static datadog.context.Context.current; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; @@ -33,7 +34,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } @@ -50,7 +51,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - try (final ContextScope contextScope = span.attach()) { + try (final ContextScope contextScope = current().with(span).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(contextScope.context()); } } @@ -71,7 +72,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } @@ -99,7 +100,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = parent.attach()) { + try (final ContextScope parentScope = storedContext.with(parent).attach()) { ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java index 444e86a9f2e..60835aa76b0 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java @@ -3,11 +3,11 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.servlet.ServletRequestSetter.SETTER; @@ -23,10 +23,10 @@ import com.google.auto.service.AutoService; import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -84,7 +84,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class RequestDispatcherAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope start( + public static ContextScope start( @Advice.Origin("#m") final String method, @Advice.This final RequestDispatcher dispatcher, @Advice.Local("_requestContext") Object requestContext, @@ -126,18 +126,16 @@ public static AgentScope start( // temporarily replace from request to avoid spring resource name bubbling up: requestContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); - final AgentScope scope = activateSpan(span); + final ContextScope scope = getCurrentContext().with(span).attach(); // Set the context after activation so we have the proper Context object - if (scope != null) { - request.setAttribute(DD_CONTEXT_ATTRIBUTE, scope.context()); - } + request.setAttribute(DD_CONTEXT_ATTRIBUTE, scope.context()); return scope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stop( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("_requestContext") final Object requestContext, @Advice.Argument(0) final ServletRequest request, @Advice.Argument(1) final ServletResponse response, @@ -150,10 +148,11 @@ public static void stop( request.setAttribute(DD_CONTEXT_ATTRIBUTE, requestContext); } - DECORATE.onError(scope, throwable); - DECORATE.beforeFinish(scope); + final AgentSpan span = spanFromContext(scope.context()); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } } diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java index 12fceba2509..22720024845 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.DECORATE; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.SYNAPSE_CONTINUATION_KEY; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.SYNAPSE_SPAN_KEY; @@ -12,6 +13,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -66,26 +68,32 @@ public static void createWorker(@Advice.Argument(0) final SourceRequest request) public static final class ServerWorkerResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope beginResponse( + public static ContextScope beginResponse( @Advice.FieldValue("request") final SourceRequest request) { AgentScope.Continuation continuation = (AgentScope.Continuation) request.getConnection().getContext().removeAttribute(SYNAPSE_CONTINUATION_KEY); if (null != continuation) { - return continuation.activate(); + // Activate the continuation to get the span, then attach it to current context + AgentScope agentScope = continuation.activate(); + try { + return agentScope.span().attach(); + } finally { + agentScope.close(); + } } return null; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void responseReady( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.FieldValue("request") final SourceRequest request, @Advice.Thrown final Throwable error) { if (null == scope) { return; } - AgentSpan span = scope.span(); + AgentSpan span = spanFromContext(scope.context()); HttpResponse httpResponse = request.getConnection().getHttpResponse(); if (null != httpResponse) { DECORATE.onResponse(span, httpResponse); From b33528c31c91c3bcf671c676ccc16983eddcc3bc Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Wed, 5 Nov 2025 09:43:58 -0500 Subject: [PATCH 26/27] fix netty timeout tests --- .../HttpClientRequestTracingHandler.java | 6 +----- .../HttpClientResponseTracingHandler.java | 13 ++++++------- .../HttpClientRequestTracingHandler.java | 6 +----- .../HttpClientResponseTracingHandler.java | 18 ++++++++---------- 4 files changed, 16 insertions(+), 27 deletions(-) diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index df5b670aee4..763e5e5c8e5 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -14,7 +14,6 @@ import static datadog.trace.instrumentation.netty40.client.NettyResponseInjectAdapter.SETTER; import datadog.context.Context; -import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -79,10 +78,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); - final Context spanContext; - try (final ContextScope contextScope = getCurrentContext().with(span).attach()) { - spanContext = contextScope.context(); - } + final Context spanContext = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index 5b27e6a0804..75ec585739f 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -8,7 +8,6 @@ import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; import datadog.context.Context; -import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -31,8 +30,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } final boolean finishSpan = msg instanceof HttpResponse; @@ -60,8 +59,8 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } if (span != null) { @@ -88,8 +87,8 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } if (span != null && span != parent) { diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index b4ebab59282..688e30f2dc7 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -14,7 +14,6 @@ import static datadog.trace.instrumentation.netty41.client.NettyResponseInjectAdapter.SETTER; import datadog.context.Context; -import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,10 +79,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); - final Context spanContext; - try (final ContextScope contextScope = getCurrentContext().with(span).attach()) { - spanContext = contextScope.context(); - } + final Context spanContext = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 84d9a811c9a..ce764cf7bb9 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -1,6 +1,5 @@ package datadog.trace.instrumentation.netty41.client; -import static datadog.context.Context.current; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; @@ -9,7 +8,6 @@ import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; import datadog.context.Context; -import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -34,8 +32,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } if (span != null) { @@ -51,8 +49,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - try (final ContextScope contextScope = current().with(span).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(contextScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(span)); } } } @@ -72,8 +70,8 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } if (span != null) { @@ -100,8 +98,8 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final AgentSpan span = spanFromContext(storedContext); // Set parent context back to maintain the same functionality as getAndSet(parent) - try (final ContextScope parentScope = storedContext.with(parent).attach()) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(parentScope.context()); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); } if (span != null && span != parent) { From 841b27cac25c803e09535c9c9151a01cd3d16765 Mon Sep 17 00:00:00 2001 From: Zarir Hamza Date: Thu, 6 Nov 2025 17:49:34 -0500 Subject: [PATCH 27/27] address all comments --- .../decorator/HttpServerDecorator.java | 6 +- .../grizzly/SpanClosingListener.java | 2 + .../grizzlyhttp232/GrizzlyDecorator.java | 2 + .../jetty10/JettyCommitResponseHelper.java | 28 ++++++--- .../instrumentation/jetty10/ResetAdvice.java | 2 + .../HttpClientRequestTracingHandler.java | 4 +- .../HttpClientRequestTracingHandler.java | 4 +- .../HttpClientResponseTracingHandler.java | 2 +- .../synapse3/SynapseClientDecorator.java | 2 +- .../SynapseClientInstrumentation.java | 58 +++++++++++-------- .../SynapseClientWorkerInstrumentation.java | 25 +++++--- .../synapse3/SynapseServerDecorator.java | 2 +- .../SynapseServerInstrumentation.java | 10 ++-- .../SynapseServerWorkerInstrumentation.java | 8 +-- .../tomcat/RequestInstrumentation.java | 10 +++- .../tomcat/ResponseInstrumentation.java | 10 +++- .../undertow/ExchangeEndSpanListener.java | 25 ++++---- .../HttpRequestParserInstrumentation.java | 4 +- 18 files changed, 129 insertions(+), 75 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index e2a53ee7e8c..c6c1f747a4a 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -539,10 +539,12 @@ private Flow callIGCallbackURI( @Override public Context beforeFinish(Context context) { AgentSpan span = AgentSpan.fromContext(context); - onRequestEndForInstrumentationGateway(span); + if (span != null) { + onRequestEndForInstrumentationGateway(span); + } // Close Serverless Gateway Inferred Span if any - finishInferredProxySpan(context); + // finishInferredProxySpan(context); return super.beforeFinish(context); } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java index f6ae98af71c..46142eb7771 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java @@ -24,6 +24,8 @@ public void onAfterService(final Request request) { DECORATE.onResponse(span, request.getResponse()); DECORATE.beforeFinish(context); span.finish(); + } else { + DECORATE.beforeFinish(context); } } } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index 9c75322a1aa..adf307e2cbf 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -154,6 +154,8 @@ public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable DECORATE.onError(span, throwable); DECORATE.beforeFinish(context); span.finish(); + } else if (context != null) { + DECORATE.beforeFinish(context); } ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java index 4dd2494355a..14afb30fc2a 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java @@ -51,14 +51,26 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - Object contextObj; - Context context; - AgentSpan span; - RequestContext requestContext; - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null - || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) - || (span = spanFromContext(context = (Context) contextObj)) == null - || (requestContext = span.getRequestContext()) == null) { + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { + state.partialResponse(); + return false; + } + + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (!(contextObj instanceof Context)) { + state.partialResponse(); + return false; + } + + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span == null) { + state.partialResponse(); + return false; + } + + RequestContext requestContext = span.getRequestContext(); + if (requestContext == null) { state.partialResponse(); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java index f398135f14f..d28f86d14dd 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java @@ -26,6 +26,8 @@ public static void stopSpan(@Advice.This final HttpChannel channel) { JettyDecorator.OnResponse.onResponse(span, channel); DECORATE.beforeFinish(context); span.finish(); + } else { + DECORATE.beforeFinish(context); } } } diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index 763e5e5c8e5..cd000aaead4 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -78,7 +78,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); - final Context spanContext = getCurrentContext().with(span); + final Context context = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -93,7 +93,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(context); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index 688e30f2dc7..89e6f151bff 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -79,7 +79,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); - final Context spanContext = getCurrentContext().with(span); + final Context context = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -94,7 +94,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(context); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index ce764cf7bb9..a1973706338 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -50,7 +50,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { } } else { if (storedContext != null) { - ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(span)); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext); } } } diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientDecorator.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientDecorator.java index 98d330bdc39..825241fedd6 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientDecorator.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientDecorator.java @@ -14,7 +14,7 @@ public final class SynapseClientDecorator extends HttpClientDecorator