diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java index 70587135ac5..85598459f26 100644 --- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/AsyncPropagatingDisableInstrumentation.java @@ -47,6 +47,8 @@ public AsyncPropagatingDisableInstrumentation() { namedOneOf("reactor.core.scheduler.SchedulerTask", "reactor.core.scheduler.WorkerTask"); private static final ElementMatcher RXJAVA2_DISABLED_TYPE_INITIALIZERS = named("io.reactivex.internal.schedulers.AbstractDirectTask"); + private static final ElementMatcher JAVA_HTTP_CLIENT = + extendsClass(named("java.net.http.HttpClient")); @Override public boolean onlyMatchKnownTypes() { @@ -80,7 +82,8 @@ public String[] knownMatchingTypes() { "org.springframework.jms.listener.DefaultMessageListenerContainer", "org.apache.activemq.broker.TransactionBroker", "com.mongodb.internal.connection.DefaultConnectionPool$AsyncWorkManager", - "io.reactivex.internal.schedulers.AbstractDirectTask" + "io.reactivex.internal.schedulers.AbstractDirectTask", + "jdk.internal.net.http.HttpClientImpl" }; } @@ -94,7 +97,8 @@ public ElementMatcher hierarchyMatcher() { return RX_WORKERS .or(GRPC_MANAGED_CHANNEL) .or(REACTOR_DISABLED_TYPE_INITIALIZERS) - .or(RXJAVA2_DISABLED_TYPE_INITIALIZERS); + .or(RXJAVA2_DISABLED_TYPE_INITIALIZERS) + .or(JAVA_HTTP_CLIENT); } @Override @@ -180,6 +184,7 @@ public void methodAdvice(MethodTransformer transformer) { isTypeInitializer().and(isDeclaredBy(REACTOR_DISABLED_TYPE_INITIALIZERS)), advice); transformer.applyAdvice( isTypeInitializer().and(isDeclaredBy(RXJAVA2_DISABLED_TYPE_INITIALIZERS)), advice); + transformer.applyAdvice(namedOneOf("sendAsync").and(isDeclaredBy(JAVA_HTTP_CLIENT)), advice); } public static class DisableAsyncAdvice { diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/BodyHandlerWrapper.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/BodyHandlerWrapper.java index e0c2ce126a4..569160917e1 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/BodyHandlerWrapper.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/BodyHandlerWrapper.java @@ -1,6 +1,9 @@ package datadog.trace.instrumentation.httpclient; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan; + import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.net.http.HttpResponse.BodyHandler; import java.net.http.HttpResponse.BodySubscriber; import java.net.http.HttpResponse.ResponseInfo; @@ -11,20 +14,21 @@ public class BodyHandlerWrapper implements BodyHandler { private final BodyHandler delegate; - private final AgentScope.Continuation continuation; + private final AgentSpan span; - public BodyHandlerWrapper(BodyHandler delegate, AgentScope.Continuation context) { + public BodyHandlerWrapper(BodyHandler delegate, AgentSpan span) { this.delegate = delegate; - this.continuation = context; + this.span = span; } @Override public BodySubscriber apply(ResponseInfo responseInfo) { + // Capture the continuation lazily here rather than at sendAsync() call time. BodySubscriber subscriber = delegate.apply(responseInfo); if (subscriber instanceof BodySubscriberWrapper) { return subscriber; } - return new BodySubscriberWrapper<>(subscriber, continuation); + return new BodySubscriberWrapper<>(subscriber, captureSpan(span)); } static class BodySubscriberWrapper implements BodySubscriber { diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java index c58b1e5c562..33b2ed9db21 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java11/datadog/trace/instrumentation/httpclient/SendAsyncAdvice.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.httpclient; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.DECORATE; import static datadog.trace.instrumentation.httpclient.JavaNetClientDecorator.INSTRUMENTATION_NAME; @@ -38,7 +37,10 @@ public static AgentScope methodEnter( final AgentSpan span = startSpan(INSTRUMENTATION_NAME, OPERATION_NAME); final AgentScope scope = activateSpan(span); if (bodyHandler != null) { - bodyHandler = new BodyHandlerWrapper<>(bodyHandler, captureSpan(span)); + // Pass span directly — BodyHandlerWrapper captures the continuation lazily in apply(), + // only once response headers arrive. This avoids leaking a continuation when the + // connection fails before headers are received. + bodyHandler = new BodyHandlerWrapper<>(bodyHandler, span); } DECORATE.afterStart(span); diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy index c666e0d0efa..c07838d941f 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/groovy/datadog/trace/instrumentation/httpclient/JavaHttpClientTest.groovy @@ -10,11 +10,6 @@ import java.net.http.HttpResponse import java.time.Duration abstract class JavaHttpClientTest extends HttpClientTest { - @Override - boolean useStrictTraceWrites() { - // TODO fix this by making sure that spans get closed properly - return false - } def client = HttpClient.newBuilder() .connectTimeout(Duration.ofMillis(CONNECT_TIMEOUT_MS))