diff --git a/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/CompletableFutureWrapper.java b/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/CompletableFutureWrapper.java new file mode 100644 index 000000000000..902f02c90637 --- /dev/null +++ b/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/CompletableFutureWrapper.java @@ -0,0 +1,35 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.httpclient; + +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import java.util.concurrent.CompletableFuture; + +public final class CompletableFutureWrapper { + + private CompletableFutureWrapper() {} + + public static CompletableFuture wrap(CompletableFuture future, Context context) { + if (context == Context.root()) { + return future; + } + + CompletableFuture result = new CompletableFuture<>(); + future.whenComplete( + (T value, Throwable throwable) -> { + try (Scope ignored = context.makeCurrent()) { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + result.complete(value); + } + } + }); + + return result; + } +} diff --git a/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/HttpClientInstrumentation.java b/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/HttpClientInstrumentation.java index 7125b013e519..5017021409b9 100644 --- a/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/HttpClientInstrumentation.java +++ b/instrumentation/java-http-client/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/httpclient/HttpClientInstrumentation.java @@ -106,8 +106,9 @@ public static void methodEnter( @Advice.Argument(value = 0) HttpRequest httpRequest, @Advice.Argument(value = 1, readOnly = false) HttpResponse.BodyHandler bodyHandler, @Advice.Local("otelContext") Context context, + @Advice.Local("otelParentContext") Context parentContext, @Advice.Local("otelScope") Scope scope) { - Context parentContext = currentContext(); + parentContext = currentContext(); if (bodyHandler != null) { bodyHandler = new BodyHandlerWrapper(bodyHandler, parentContext); } @@ -124,6 +125,7 @@ public static void methodExit( @Advice.Return(readOnly = false) CompletableFuture> future, @Advice.Thrown Throwable throwable, @Advice.Local("otelContext") Context context, + @Advice.Local("otelParentContext") Context parentContext, @Advice.Local("otelScope") Scope scope) { if (scope == null) { return; @@ -134,6 +136,7 @@ public static void methodExit( tracer().endExceptionally(context, null, throwable); } else { future = future.whenComplete(new ResponseConsumer(context)); + future = CompletableFutureWrapper.wrap(future, parentContext); } } } diff --git a/instrumentation/java-http-client/javaagent/src/test/groovy/JdkHttpClientTest.groovy b/instrumentation/java-http-client/javaagent/src/test/groovy/JdkHttpClientTest.groovy index 20053b8f54ff..97a6c40d02f8 100644 --- a/instrumentation/java-http-client/javaagent/src/test/groovy/JdkHttpClientTest.groovy +++ b/instrumentation/java-http-client/javaagent/src/test/groovy/JdkHttpClientTest.groovy @@ -57,10 +57,4 @@ class JdkHttpClientTest extends HttpClientTest implements AgentTest boolean testWithClientParent() { false } - - // TODO: context not propagated to callback - @Override - boolean testErrorWithCallback() { - return false - } }