handled case with recycling when client closed before executing

elastic · Nov 9, 2023 · 2140554 · 2140554
1 parent 63e40db
commit 2140554
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 12 deletions.
diff --git a/...main/java/co/elastic/apm/agent/httpclient/common/AbstractApacheHttpAsyncClientAdvice.java b/...main/java/co/elastic/apm/agent/httpclient/common/AbstractApacheHttpAsyncClientAdvice.java
@@ -38,10 +38,12 @@ public static <PRODUCER, WRAPPER extends PRODUCER, CALLBACK, CALLBACK_WRAPPER ex
             span.deactivate();
             // End the span if the method terminated with an exception.
             // The exception means that the listener who normally does the ending will not be invoked.
+            WRAPPER wrapper = (WRAPPER) enter[0];
             if (t != null) {
-                WRAPPER wrapper = (WRAPPER) enter[0];
                 CALLBACK_WRAPPER cb = (CALLBACK_WRAPPER) enter[1];
+                // only for apachehttpclient_v4
                 asyncHelper.failedWithoutException(cb, t);
+
                 asyncHelper.recycle(wrapper);
             }
         }

diff --git a/...c/main/java/co/elastic/apm/agent/httpclient/v5/ApacheHttpAsyncClient5Instrumentation.java b/...c/main/java/co/elastic/apm/agent/httpclient/v5/ApacheHttpAsyncClient5Instrumentation.java
@@ -6,7 +6,6 @@
 import net.bytebuddy.description.NamedElement;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.hc.core5.concurrent.FutureCallback;
 import org.apache.hc.core5.http.nio.AsyncRequestProducer;
@@ -15,6 +14,7 @@
 import javax.annotation.Nullable;
 
 import static co.elastic.apm.agent.sdk.bytebuddy.CustomElementMatchers.classLoaderCanLoadClass;
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isBootstrapClassLoader;
 import static net.bytebuddy.matcher.ElementMatchers.nameContains;
@@ -55,7 +55,10 @@ public static class ApacheHttpAsyncClient5Advice extends AbstractApacheHttpAsync
         private static ApacheHttpAsyncClientHelper asyncHelper = new ApacheHttpAsyncClientHelper();
 
         @Nullable
-        @Advice.AssignReturned.ToArguments({@Advice.AssignReturned.ToArguments.ToArgument(index = 0, value = 0, typing = Assigner.Typing.DYNAMIC), @Advice.AssignReturned.ToArguments.ToArgument(index = 1, value = 3, typing = Assigner.Typing.DYNAMIC)})
+        @Advice.AssignReturned.ToArguments({
+            @Advice.AssignReturned.ToArguments.ToArgument(index = 0, value = 0, typing = DYNAMIC),
+            @Advice.AssignReturned.ToArguments.ToArgument(index = 1, value = 3, typing = DYNAMIC)
+        })
         @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
         public static Object[] onBeforeExecute(@Advice.Argument(value = 0) AsyncRequestProducer asyncRequestProducer, @Advice.Argument(value = 2) HttpContext context, @Advice.Argument(value = 3) FutureCallback<?> futureCallback) {
             return startSpan(tracer, asyncHelper, asyncRequestProducer, context, futureCallback);

diff --git a/.../src/main/java/co/elastic/apm/agent/httpclient/v5/helper/ApacheHttpAsyncClientHelper.java b/.../src/main/java/co/elastic/apm/agent/httpclient/v5/helper/ApacheHttpAsyncClientHelper.java
@@ -23,7 +23,6 @@ public class ApacheHttpAsyncClientHelper implements AbstractApacheHttpAsyncClien
     private final Tracer tracer;
     private final ObjectPool<AsyncRequestProducerWrapper> requestProducerWrapperObjectPool;
     private final ObjectPool<FutureCallbackWrapper<?>> futureCallbackWrapperObjectPool;
-
     private final ObjectPool<RequestChannelWrapper> requestChannelWrapperObjectPool;
 
     public ApacheHttpAsyncClientHelper() {
@@ -68,7 +67,7 @@ public FutureCallbackWrapper<?> wrapFutureCallback(FutureCallback futureCallback
 
     @Override
     public void failedWithoutException(FutureCallbackWrapper<?> cb, Throwable t) {
-        cb.failedWithoutExecution(t);
+        // ignore
     }
 
     public RequestChannelWrapper wrapRequestChannel(RequestChannel requestChannel, @Nullable Span<?> span, @Nullable ElasticContext<?> toPropagate) {

diff --git a/.../src/main/java/co/elastic/apm/agent/httpclient/v5/helper/AsyncRequestProducerWrapper.java b/.../src/main/java/co/elastic/apm/agent/httpclient/v5/helper/AsyncRequestProducerWrapper.java
@@ -5,7 +5,6 @@
 import co.elastic.apm.agent.tracer.Span;
 import co.elastic.apm.agent.tracer.pooling.Recyclable;
 import org.apache.hc.core5.http.HttpException;
-import org.apache.hc.core5.http.HttpHost;
 import org.apache.hc.core5.http.nio.AsyncRequestProducer;
 import org.apache.hc.core5.http.nio.DataStreamChannel;
 import org.apache.hc.core5.http.nio.RequestChannel;
@@ -42,12 +41,14 @@ public AsyncRequestProducerWrapper with(AsyncRequestProducer delegate, @Nullable
 
     @Override
     public void sendRequest(RequestChannel requestChannel, HttpContext httpContext) throws HttpException, IOException {
-        RequestChannelWrapper wrappedRequestChannel = null;
+        RequestChannelWrapper wrappedRequestChannel = asyncClientHelper.wrapRequestChannel(requestChannel, span, toPropagate);
+        boolean isNotNullWrappedRequestChannel = null != wrappedRequestChannel;
         try {
-            wrappedRequestChannel = asyncClientHelper.wrapRequestChannel(requestChannel, span, toPropagate);
-        } finally {
-            boolean isNotNullWrappedRequestChannel = null != wrappedRequestChannel;
             delegate.sendRequest(isNotNullWrappedRequestChannel ? wrappedRequestChannel : requestChannel, httpContext);
+        } catch (HttpException | IOException | IllegalStateException e) {
+            asyncClientHelper.recycle(this);
+            throw e;
+        } finally {
             if (isNotNullWrappedRequestChannel) {
                 asyncClientHelper.recycle(wrappedRequestChannel);
             }

diff --git a/...est/java/co/elastic/apm/agent/httpclient/v5/ApacheHttpAsyncClientInstrumentationTest.java b/...est/java/co/elastic/apm/agent/httpclient/v5/ApacheHttpAsyncClientInstrumentationTest.java
@@ -18,6 +18,7 @@
 import org.junit.Test;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 
 import java.io.IOException;
 import java.net.URI;
@@ -73,15 +74,14 @@ public void cancelled() {
 
 
     @Test
-    @Ignore
     public void testSpanFinishOnEarlyException() throws Exception {
 
         client.close(); //this forces execute to immediately exit with an exception
 
         reporter.disableCheckServiceTarget();
         reporter.disableCheckDestinationAddress();
         try {
-            assertThatThrownBy(() -> performGet(getBaseUrl() + "/")).isInstanceOf(IllegalStateException.class);
+            assertThatThrownBy(() -> performGet(getBaseUrl() + "/")).cause().isInstanceOf(IllegalStateException.class);
         } finally {
             //Reset state for other tests
             setUp();