Keep spans active during AsyncHandler methods in AsyncHttpClient

CHANGELOG.asciidoc

@@ -53,6 +53,7 @@ This is not a breaking change, so former
 * Fix OpenTracing error tag handling (set transaction error result when tag value is `true`) {pull}1159[#1159]
 * Due to a bug in the build we didn't include the gRPC plugin in the build so far
 * `java.lang.ClassNotFoundException: Unable to load class 'jdk.internal...'` is thrown when tracing specific versions of Atlassian systems {pull}1168[#1168]
+* Make sure spans are kept active during `AsyncHandler` methods in the `AsyncHttpClient`
 
 
 [[release-notes-1.x]]

apm-agent-plugins/apm-asynchttpclient-plugin/src/main/java/co/elastic/apm/agent/asynchttpclient/AbstractAsyncHttpClientInstrumentation.java

@@ -178,10 +178,18 @@ public AsyncHandlerOnCompletedInstrumentation() {
         }
 
         @Advice.OnMethodEnter(suppress = Throwable.class)
-        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler) {
-            final Span span = handlerSpanMap.remove(asyncHandler);
+        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span) {
+            span = handlerSpanMap.remove(asyncHandler);
+            if (span != null) {
+                span.activate();
+            }
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        private static void onMethodExit(@Advice.This AsyncHandler<?> asyncHandler, @Nullable @Advice.Local("span") Span span) {
             if (span != null) {
                 span.end();
+                span.deactivate();
             }
         }
     }
@@ -193,10 +201,18 @@ public AsyncHandlerOnThrowableInstrumentation() {
         }
 
         @Advice.OnMethodEnter(suppress = Throwable.class)
-        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Argument(0) Throwable t) {
-            final Span span = handlerSpanMap.remove(asyncHandler);
+        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span) {
+            span = handlerSpanMap.remove(asyncHandler);
+            if (span != null) {
+                span.activate();
+            }
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        private static void onMethodExit(@Advice.This AsyncHandler<?> asyncHandler, @Nullable @Advice.Local("span") Span span, @Advice.Argument(0) Throwable t) {
             if (span != null) {
                 span.captureException(t).end();
+                span.deactivate();
             }
         }
     }
@@ -208,10 +224,18 @@ public AsyncHandlerOnStatusReceivedInstrumentation() {
         }
 
         @Advice.OnMethodEnter(suppress = Throwable.class)
-        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Argument(0) HttpResponseStatus status) {
-            final Span span = handlerSpanMap.get(asyncHandler);
+        private static void onMethodEnter(@Advice.This AsyncHandler<?> asyncHandler, @Advice.Local("span") Span span, @Advice.Argument(0) HttpResponseStatus status) {
+            span = handlerSpanMap.get(asyncHandler);
+            if (span != null) {
+                span.activate();
+            }
+        }
+
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        private static void onMethodExit(@Advice.This AsyncHandler<?> asyncHandler, @Nullable @Advice.Local("span") Span span, @Advice.Argument(0) HttpResponseStatus status) {
             if (span != null) {
                 span.getContext().getHttp().withStatusCode(status.getStatusCode());
+                span.deactivate();
             }
         }
     }

apm-agent-plugins/apm-asynchttpclient-plugin/src/test/java/co/elastic/apm/agent/asynchttpclient/AsyncHttpClientInstrumentationTest.java

@@ -11,9 +11,9 @@
  * the Apache License, Version 2.0 (the "License"); you may
  * not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -25,10 +25,7 @@
 package co.elastic.apm.agent.asynchttpclient;
 
 import co.elastic.apm.agent.httpclient.AbstractHttpClientInstrumentationTest;
-import org.asynchttpclient.AsyncCompletionHandlerBase;
-import org.asynchttpclient.AsyncHttpClient;
-import org.asynchttpclient.Dsl;
-import org.asynchttpclient.RequestBuilder;
+import org.asynchttpclient.*;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.runner.RunWith;
@@ -38,6 +35,7 @@
 import java.util.Arrays;
 
 import static org.asynchttpclient.Dsl.asyncHttpClient;
+import static org.assertj.core.api.Assertions.assertThat;
 
 @RunWith(Parameterized.class)
 public class AsyncHttpClientInstrumentationTest extends AbstractHttpClientInstrumentationTest {
@@ -49,11 +47,35 @@ public AsyncHttpClientInstrumentationTest(RequestExecutor requestExecutor) {
         this.requestExecutor = requestExecutor;
     }
 
+    public static AsyncHandler<Response> customAsyncHandler = new AsyncCompletionHandler<Response>() {
+        @Override
+        public State onStatusReceived(HttpResponseStatus responseStatus) {
+            assertThat(tracer.getActive()).isNotNull();
+            assertThat(tracer.getActive().isExit()).isTrue();
+            return State.CONTINUE;
+        }
+
+        @Override
+        public void onThrowable(Throwable t) {
+            assertThat(tracer.getActive()).isNotNull();
+            assertThat(tracer.getActive().isExit()).isTrue();
+        }
+
+        @Override
+        public Response onCompleted(Response response) {
+            assertThat(tracer.getActive()).isNotNull();
+            assertThat(tracer.getActive().isExit()).isTrue();
+            return response;
+        }
+
+    };
+
     @Parameterized.Parameters()
     public static Iterable<RequestExecutor> data() {
         return Arrays.asList(
             (client, path) -> client.executeRequest(new RequestBuilder().setUrl(path).build()).get(),
             (client, path) -> client.executeRequest(new RequestBuilder().setUrl(path).build(), new AsyncCompletionHandlerBase()).get(),
+            (client, path) -> client.executeRequest(new RequestBuilder().setUrl(path).build(), customAsyncHandler).get(),
             (client, path) -> client.prepareGet(path).execute(new AsyncCompletionHandlerBase()).get(),
             (client, path) -> client.prepareGet(path).execute().get()
         );