Akka http: propagate context to callbacks (open-telemetry#3263)

Co-authored-by: Trask Stalnaker <[email protected]>

instrumentation/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/AkkaHttpClientInstrumentationModule.java

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.akkahttp.client;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -20,6 +20,6 @@ public AkkaHttpClientInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new HttpExtClientInstrumentation());
+    return asList(new HttpExtClientInstrumentation(), new PoolMasterActorInstrumentation());
   }
 }

instrumentation/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/FutureWrapper.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.akkahttp.client;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.concurrent.impl.Promise;
+import scala.runtime.AbstractFunction1;
+import scala.util.Try;
+
+public class FutureWrapper {
+  public static <T> Future<T> wrap(
+      Future<T> future, ExecutionContext executionContext, Context context) {
+    Promise.DefaultPromise<T> promise = new Promise.DefaultPromise<>();
+    future.onComplete(
+        new AbstractFunction1<Try<T>, Object>() {
+
+          @Override
+          public Object apply(Try<T> result) {
+            try (Scope ignored = context.makeCurrent()) {
+              return promise.complete(result);
+            }
+          }
+        },
+        executionContext);
+
+    return promise;
+  }
+}

instrumentation/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/HttpExtClientInstrumentation.java

@@ -70,7 +70,7 @@ public static void methodEnter(
     public static void methodExit(
         @Advice.Argument(0) HttpRequest request,
         @Advice.This HttpExt thiz,
-        @Advice.Return Future<HttpResponse> responseFuture,
+        @Advice.Return(readOnly = false) Future<HttpResponse> responseFuture,
         @Advice.Thrown Throwable throwable,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
@@ -84,6 +84,10 @@ public static void methodExit(
       } else {
         tracer().endExceptionally(context, throwable);
       }
+      if (responseFuture != null) {
+        responseFuture =
+            FutureWrapper.wrap(responseFuture, thiz.system().dispatcher(), currentContext());
+      }
     }
   }
 }

instrumentation/akka-http-10.0/javaagent/src/test/groovy/AkkaHttpClientInstrumentationTest.groovy

@@ -48,30 +48,18 @@ class AkkaHttpClientInstrumentationTest extends HttpClientTest<HttpRequest> impl
   @Override
   void sendRequestWithCallback(HttpRequest request, String method, URI uri, Map<String, String> headers, RequestResult requestResult) {
     Http.get(system).singleRequest(request, materializer).whenComplete {response, throwable ->
-      response.discardEntityBytes(materializer)
+      if (throwable == null) {
+        response.discardEntityBytes(materializer)
+      }
       requestResult.complete({ response.status().intValue() }, throwable)
     }
   }
 
-  // TODO(anuraaga): Context leak seems to prevent us from running asynchronous tests in a row.
-  // Disable for now.
-  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/2639
-  @Override
-  boolean testCallback() {
-    false
-  }
-
   @Override
   boolean testRedirects() {
     false
   }
 
-  @Override
-  boolean testRemoteConnection() {
-    // Not sure how to properly set timeouts...
-    return false
-  }
-
   @Override
   SingleConnection createSingleConnection(String host, int port) {
     // singleConnection test would require instrumentation to support requests made through pools

instrumentation/akka-http-10.0/javaagent/src/test/resources/application.conf

@@ -2,5 +2,10 @@ akka.http {
   host-connection-pool {
     // Limit maximum http backoff for tests
     max-connection-backoff = 100ms
+    max-open-requests = 1024
+    max-retries = 0
+    client {
+      connecting-timeout = 5s
+    }
   }
 }