From 257fe750c3a64e6e4ae38a7cb933663f1e6fb032 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 10:54:35 +0200 Subject: [PATCH 1/3] Fix Java HTTP async continuation lifecycle --- .../httpclient/HttpClientInstrumentation.java | 1 + .../httpclient/BodyHandlerWrapper.java | 52 +++++++- .../httpclient/BodyHandlerWrapperTest.java | 111 ++++++++++++++++++ 3 files changed, 158 insertions(+), 6 deletions(-) create mode 100644 dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpClientInstrumentation.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpClientInstrumentation.java index e30648f61fa..e0a995378de 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/main/java/datadog/trace/instrumentation/httpclient/HttpClientInstrumentation.java @@ -55,6 +55,7 @@ public String[] helperClassNames() { return new String[] { packageName + ".BodyHandlerWrapper", packageName + ".BodyHandlerWrapper$BodySubscriberWrapper", + packageName + ".BodyHandlerWrapper$SubscriptionWrapper", packageName + ".CompletableFutureWrapper", packageName + ".JavaNetClientDecorator", packageName + ".ResponseConsumer" 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 3a8c7fa70fc..4cbf4e01854 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 @@ -10,6 +10,7 @@ import java.util.List; import java.util.concurrent.CompletionStage; import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicBoolean; public class BodyHandlerWrapper implements BodyHandler { private final BodyHandler delegate; @@ -27,12 +28,13 @@ public BodySubscriber apply(ResponseInfo responseInfo) { if (subscriber instanceof BodySubscriberWrapper) { return subscriber; } - return new BodySubscriberWrapper<>(subscriber, span.captureWithContext()); + return new BodySubscriberWrapper<>(subscriber, span.captureWithContext().hold()); } static class BodySubscriberWrapper implements BodySubscriber { private final BodySubscriber delegate; private final ContextContinuation continuation; + private final AtomicBoolean continuationReleased = new AtomicBoolean(); public BodySubscriberWrapper(BodySubscriber delegate, ContextContinuation continuation) { this.delegate = delegate; @@ -50,7 +52,7 @@ public CompletionStage getBody() { @Override public void onSubscribe(Flow.Subscription subscription) { - delegate.onSubscribe(subscription); + delegate.onSubscribe(new SubscriptionWrapper(subscription, this)); } @Override @@ -62,15 +64,53 @@ public void onNext(List item) { @Override public void onError(Throwable throwable) { - try (ContextScope ignore = continuation.resume()) { - delegate.onError(throwable); + try { + try (ContextScope ignore = continuation.resume()) { + delegate.onError(throwable); + } + } finally { + releaseContinuation(); } } @Override public void onComplete() { - try (ContextScope ignore = continuation.resume()) { - delegate.onComplete(); + try { + try (ContextScope ignore = continuation.resume()) { + delegate.onComplete(); + } + } finally { + releaseContinuation(); + } + } + + private void releaseContinuation() { + if (continuationReleased.compareAndSet(false, true)) { + continuation.release(); + } + } + } + + static final class SubscriptionWrapper implements Flow.Subscription { + private final Flow.Subscription delegate; + private final BodySubscriberWrapper subscriber; + + SubscriptionWrapper(Flow.Subscription delegate, BodySubscriberWrapper subscriber) { + this.delegate = delegate; + this.subscriber = subscriber; + } + + @Override + public void request(long count) { + delegate.request(count); + } + + @Override + public void cancel() { + try { + delegate.cancel(); + } finally { + subscriber.releaseContinuation(); } } } diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java new file mode 100644 index 00000000000..4b1dab7682e --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java @@ -0,0 +1,111 @@ +package datadog.trace.instrumentation.httpclient; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.context.Context; +import datadog.context.ContextContinuation; +import datadog.context.ContextScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.lang.reflect.Proxy; +import java.net.http.HttpResponse.BodySubscriber; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.Flow; +import org.junit.jupiter.api.Test; + +class BodyHandlerWrapperTest { + + @Test + void releasesContinuationWhenSubscriptionIsCancelled() { + RecordingContinuation continuation = new RecordingContinuation(); + AgentSpan span = + (AgentSpan) + Proxy.newProxyInstance( + AgentSpan.class.getClassLoader(), + new Class[] {AgentSpan.class}, + (proxy, method, args) -> continuation); + RecordingSubscriber subscriber = new RecordingSubscriber(); + BodySubscriber wrapper = + new BodyHandlerWrapper<>(ignored -> subscriber, span).apply(null); + RecordingSubscription subscription = new RecordingSubscription(); + + wrapper.onSubscribe(subscription); + subscriber.subscription.request(3); + subscriber.subscription.cancel(); + wrapper.onComplete(); + + assertEquals(3, subscription.requested); + assertEquals(1, subscription.cancelled); + assertEquals(1, continuation.released); + } + + private static final class RecordingSubscriber + implements java.net.http.HttpResponse.BodySubscriber { + private final CompletableFuture body = new CompletableFuture<>(); + private Flow.Subscription subscription; + + @Override + public CompletionStage getBody() { + return body; + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + this.subscription = subscription; + } + + @Override + public void onNext(List item) {} + + @Override + public void onError(Throwable throwable) { + body.completeExceptionally(throwable); + } + + @Override + public void onComplete() { + body.complete(null); + } + } + + private static final class RecordingSubscription implements Flow.Subscription { + private long requested; + private int cancelled; + + @Override + public void request(long count) { + requested += count; + } + + @Override + public void cancel() { + cancelled++; + } + } + + private static final class RecordingContinuation implements ContextContinuation { + private int released; + + @Override + public ContextContinuation hold() { + return this; + } + + @Override + public Context context() { + return null; + } + + @Override + public ContextScope resume() { + return null; + } + + @Override + public void release() { + released++; + } + } +} From 72169baf12414606b32b9463e705c3b3fd82e1f8 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 12:42:57 +0200 Subject: [PATCH 2/3] Release Java HTTP context when body callbacks fail --- .../httpclient/BodyHandlerWrapper.java | 44 ++++++++++--- .../httpclient/BodyHandlerWrapperTest.java | 61 ++++++++++++++++--- 2 files changed, 87 insertions(+), 18 deletions(-) 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 4cbf4e01854..2aa22ccde00 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 @@ -10,7 +10,7 @@ import java.util.List; import java.util.concurrent.CompletionStage; import java.util.concurrent.Flow; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; public class BodyHandlerWrapper implements BodyHandler { private final BodyHandler delegate; @@ -32,9 +32,13 @@ public BodySubscriber apply(ResponseInfo responseInfo) { } static class BodySubscriberWrapper implements BodySubscriber { + private static final AtomicReferenceFieldUpdater + CONTINUATION = + AtomicReferenceFieldUpdater.newUpdater( + BodySubscriberWrapper.class, ContextContinuation.class, "continuation"); + private final BodySubscriber delegate; - private final ContextContinuation continuation; - private final AtomicBoolean continuationReleased = new AtomicBoolean(); + private volatile ContextContinuation continuation; public BodySubscriberWrapper(BodySubscriber delegate, ContextContinuation continuation) { this.delegate = delegate; @@ -52,20 +56,36 @@ public CompletionStage getBody() { @Override public void onSubscribe(Flow.Subscription subscription) { - delegate.onSubscribe(new SubscriptionWrapper(subscription, this)); + boolean completed = false; + try { + delegate.onSubscribe(new SubscriptionWrapper(subscription, this)); + completed = true; + } finally { + if (!completed) { + releaseContinuation(); + } + } } @Override public void onNext(List item) { - try (ContextScope ignore = continuation.resume()) { - delegate.onNext(item); + boolean completed = false; + try { + try (ContextScope ignore = resumeContinuation()) { + delegate.onNext(item); + } + completed = true; + } finally { + if (!completed) { + releaseContinuation(); + } } } @Override public void onError(Throwable throwable) { try { - try (ContextScope ignore = continuation.resume()) { + try (ContextScope ignore = resumeContinuation()) { delegate.onError(throwable); } } finally { @@ -76,7 +96,7 @@ public void onError(Throwable throwable) { @Override public void onComplete() { try { - try (ContextScope ignore = continuation.resume()) { + try (ContextScope ignore = resumeContinuation()) { delegate.onComplete(); } } finally { @@ -84,8 +104,14 @@ public void onComplete() { } } + private ContextScope resumeContinuation() { + ContextContinuation continuation = this.continuation; + return continuation == null ? null : continuation.resume(); + } + private void releaseContinuation() { - if (continuationReleased.compareAndSet(false, true)) { + ContextContinuation continuation = CONTINUATION.getAndSet(this, null); + if (continuation != null) { continuation.release(); } } diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java index 4b1dab7682e..8e05b49a5ac 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.httpclient; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import datadog.context.Context; import datadog.context.ContextContinuation; @@ -17,18 +18,40 @@ class BodyHandlerWrapperTest { + @Test + void releasesContinuationWhenOnSubscribeThrows() { + RecordingContinuation continuation = new RecordingContinuation(); + RecordingSubscriber subscriber = new RecordingSubscriber(); + subscriber.throwOnSubscribe = true; + BodySubscriber wrapper = wrap(subscriber, continuation); + + assertThrows( + IllegalStateException.class, () -> wrapper.onSubscribe(new RecordingSubscription())); + assertEquals(1, continuation.released); + wrapper.onComplete(); + + assertEquals(1, continuation.released); + } + + @Test + void releasesContinuationWhenOnNextThrows() { + RecordingContinuation continuation = new RecordingContinuation(); + RecordingSubscriber subscriber = new RecordingSubscriber(); + subscriber.throwOnNext = true; + BodySubscriber wrapper = wrap(subscriber, continuation); + + assertThrows(IllegalStateException.class, () -> wrapper.onNext(List.of())); + assertEquals(1, continuation.released); + wrapper.onComplete(); + + assertEquals(1, continuation.released); + } + @Test void releasesContinuationWhenSubscriptionIsCancelled() { RecordingContinuation continuation = new RecordingContinuation(); - AgentSpan span = - (AgentSpan) - Proxy.newProxyInstance( - AgentSpan.class.getClassLoader(), - new Class[] {AgentSpan.class}, - (proxy, method, args) -> continuation); RecordingSubscriber subscriber = new RecordingSubscriber(); - BodySubscriber wrapper = - new BodyHandlerWrapper<>(ignored -> subscriber, span).apply(null); + BodySubscriber wrapper = wrap(subscriber, continuation); RecordingSubscription subscription = new RecordingSubscription(); wrapper.onSubscribe(subscription); @@ -41,10 +64,23 @@ void releasesContinuationWhenSubscriptionIsCancelled() { assertEquals(1, continuation.released); } + private static BodySubscriber wrap( + RecordingSubscriber subscriber, RecordingContinuation continuation) { + AgentSpan span = + (AgentSpan) + Proxy.newProxyInstance( + AgentSpan.class.getClassLoader(), + new Class[] {AgentSpan.class}, + (proxy, method, args) -> continuation); + return new BodyHandlerWrapper<>(ignored -> subscriber, span).apply(null); + } + private static final class RecordingSubscriber implements java.net.http.HttpResponse.BodySubscriber { private final CompletableFuture body = new CompletableFuture<>(); private Flow.Subscription subscription; + private boolean throwOnSubscribe; + private boolean throwOnNext; @Override public CompletionStage getBody() { @@ -54,10 +90,17 @@ public CompletionStage getBody() { @Override public void onSubscribe(Flow.Subscription subscription) { this.subscription = subscription; + if (throwOnSubscribe) { + throw new IllegalStateException("onSubscribe"); + } } @Override - public void onNext(List item) {} + public void onNext(List item) { + if (throwOnNext) { + throw new IllegalStateException("onNext"); + } + } @Override public void onError(Throwable throwable) { From b3f2a1640329bd4f90e89b46a4e1c617d1782065 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 18 Sep 2026 10:22:38 +0200 Subject: [PATCH 3/3] suggestion --- .../httpclient/BodyHandlerWrapperTest.java | 38 ++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java index 8e05b49a5ac..ddb930c85f6 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-11.0/src/test/java/datadog/trace/instrumentation/httpclient/BodyHandlerWrapperTest.java @@ -1,15 +1,18 @@ package datadog.trace.instrumentation.httpclient; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.context.ContextKey; import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.lang.reflect.Proxy; import java.net.http.HttpResponse.BodySubscriber; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; @@ -18,6 +21,36 @@ class BodyHandlerWrapperTest { + @Test + void holdsContextAcrossCallbacksUntilCompletion() { + Context original = Context.current(); + Context captured = original.with(ContextKey.named("response-body"), new Object()); + ContextContinuation continuation = captured.capture(); + RecordingSubscriber subscriber = new RecordingSubscriber(); + BodySubscriber wrapper = wrap(subscriber, continuation); + + try { + wrapper.onNext(List.of()); + assertSame(captured, subscriber.callbackContexts.get(0)); + assertSame(original, Context.current()); + + wrapper.onNext(List.of()); + assertSame(captured, subscriber.callbackContexts.get(1)); + assertSame(original, Context.current()); + + wrapper.onComplete(); + assertSame(captured, subscriber.callbackContexts.get(2)); + assertSame(original, Context.current()); + + // A released continuation must no longer reactivate the captured context. + try (ContextScope ignored = continuation.resume()) { + assertSame(original, Context.current()); + } + } finally { + continuation.release(); + } + } + @Test void releasesContinuationWhenOnSubscribeThrows() { RecordingContinuation continuation = new RecordingContinuation(); @@ -65,7 +98,7 @@ void releasesContinuationWhenSubscriptionIsCancelled() { } private static BodySubscriber wrap( - RecordingSubscriber subscriber, RecordingContinuation continuation) { + RecordingSubscriber subscriber, ContextContinuation continuation) { AgentSpan span = (AgentSpan) Proxy.newProxyInstance( @@ -78,6 +111,7 @@ private static BodySubscriber wrap( private static final class RecordingSubscriber implements java.net.http.HttpResponse.BodySubscriber { private final CompletableFuture body = new CompletableFuture<>(); + private final List callbackContexts = new ArrayList<>(); private Flow.Subscription subscription; private boolean throwOnSubscribe; private boolean throwOnNext; @@ -97,6 +131,7 @@ public void onSubscribe(Flow.Subscription subscription) { @Override public void onNext(List item) { + callbackContexts.add(Context.current()); if (throwOnNext) { throw new IllegalStateException("onNext"); } @@ -109,6 +144,7 @@ public void onError(Throwable throwable) { @Override public void onComplete() { + callbackContexts.add(Context.current()); body.complete(null); } }