Propagate context into jdk http client callback (#3719)

* Propage context into jdk http client callback * address review comment * remove null check
open-telemetry · Jul 30, 2021 · 78a4126 · 78a4126
1 parent a04a7a6
commit 78a4126
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 7 deletions.
diff --git a/.../java/io/opentelemetry/javaagent/instrumentation/httpclient/CompletableFutureWrapper.java b/.../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 <T> CompletableFuture<T> wrap(CompletableFuture<T> future, Context context) {
+    if (context == Context.root()) {
+      return future;
+    }
+
+    CompletableFuture<T> 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/...java/io/opentelemetry/javaagent/instrumentation/httpclient/HttpClientInstrumentation.java b/...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<HttpResponse<?>> 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
@@ -57,10 +57,4 @@ class JdkHttpClientTest extends HttpClientTest<HttpRequest> implements AgentTest
   boolean testWithClientParent() {
     false
   }
-
-  // TODO: context not propagated to callback
-  @Override
-  boolean testErrorWithCallback() {
-    return false
-  }
 }