Fix HttpClient 5.x callback span handling

CHANGES.md

@@ -4,6 +4,7 @@ Release Notes.
 
 9.8.0
 ------------------
+* Fix `httpclient-5.x-plugin` closing the caller thread's active span when `FutureCallback` executes on the caller thread (apache/skywalking#14097).
 
 * Fix the `NullPointerException` thrown by the `spring-webflux-5.x-webclient` and
   `spring-webflux-6.x-webclient` plugins when `DefaultClientRequestBuilder$BodyInserterRequest#writeTo` runs

apm-sniffer/apm-sdk-plugin/httpclient-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/httpclient/v5/AsyncExitSpan.java

@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under 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 KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.apm.plugin.httpclient.v5;
+
+import org.apache.hc.core5.http.HttpHost;
+import org.apache.skywalking.apm.agent.core.context.tag.Tags;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
+
+public class AsyncExitSpan {
+    private final HttpHost target;
+    private volatile Thread creator = Thread.currentThread();
+    private AbstractSpan span;
+
+    public AsyncExitSpan(HttpHost target) {
+        this.target = target;
+    }
+
+    public HttpHost getTarget() {
+        return target;
+    }
+
+    public boolean claimCreation() {
+        if (creator != Thread.currentThread()) {
+            return false;
+        }
+        creator = null;
+        return true;
+    }
+
+    public void callerReturned() {
+        creator = null;
+    }
+
+    public synchronized void start(AbstractSpan span) {
+        this.span = span;
+    }
+
+    public synchronized void onResponse(int statusCode) {
+        if (span != null) {
+            Tags.HTTP_RESPONSE_STATUS_CODE.set(span, statusCode);
+            if (statusCode >= 400) {
+                span.errorOccurred();
+            }
+        }
+    }
+
+    public synchronized void finish() {
+        end(false, null);
+    }
+
+    public synchronized void fail(Throwable cause) {
+        end(true, cause);
+    }
+
+    public synchronized void abort() {
+        end(true, null);
+    }
+
+    private void end(boolean error, Throwable cause) {
+        if (span == null) {
+            return;
+        }
+
+        if (error) {
+            span.errorOccurred();
+        }
+
+        if (cause != null) {
+            span.log(cause);
+        }
+
+        span.asyncFinish();
+        span = null;
+    }
+}

apm-sniffer/apm-sdk-plugin/httpclient-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/httpclient/v5/HttpAsyncClientDoExecuteInterceptor.java

@@ -18,42 +18,53 @@
 
 package org.apache.skywalking.apm.plugin.httpclient.v5;
 
+import java.lang.reflect.Method;
 import org.apache.hc.core5.concurrent.FutureCallback;
+import org.apache.hc.core5.http.HttpHost;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
 import org.apache.hc.core5.http.nio.AsyncResponseConsumer;
-import org.apache.hc.core5.http.protocol.HttpContext;
 import org.apache.skywalking.apm.agent.core.context.ContextManager;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
 import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.MethodInterceptResult;
+import org.apache.skywalking.apm.plugin.httpclient.v5.wrapper.AsyncRequestProducerWrapper;
 import org.apache.skywalking.apm.plugin.httpclient.v5.wrapper.AsyncResponseConsumerWrapper;
 import org.apache.skywalking.apm.plugin.httpclient.v5.wrapper.FutureCallbackWrapper;
 
-import java.lang.reflect.Method;
-
 public class HttpAsyncClientDoExecuteInterceptor implements InstanceMethodsAroundInterceptor {
 
     @Override
     public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
             MethodInterceptResult result) throws Throwable {
+        if (!ContextManager.isActive()) {
+            return;
+        }
+
+        AsyncExitSpan exitSpan = new AsyncExitSpan((HttpHost) allArguments[0]);
+
+        AsyncRequestProducer producer = (AsyncRequestProducer) allArguments[1];
         AsyncResponseConsumer consumer = (AsyncResponseConsumer) allArguments[2];
-        HttpContext context = (HttpContext) allArguments[4];
         FutureCallback callback = (FutureCallback) allArguments[5];
-        allArguments[2] = new AsyncResponseConsumerWrapper(consumer);
-        allArguments[5] = new FutureCallbackWrapper(callback);
-        if (ContextManager.isActive()) {
-            context.setAttribute(Constants.SKYWALKING_CONTEXT_SNAPSHOT, ContextManager.capture());
-        }
+
+        allArguments[1] = new AsyncRequestProducerWrapper(producer, exitSpan);
+        allArguments[2] = new AsyncResponseConsumerWrapper(consumer, exitSpan);
+        allArguments[5] = new FutureCallbackWrapper(callback, exitSpan);
     }
 
     @Override
     public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
             Object ret) throws Throwable {
+        if (allArguments[1] instanceof AsyncRequestProducerWrapper) {
+            ((AsyncRequestProducerWrapper) allArguments[1]).getExitSpan().callerReturned();
+        }
         return ret;
     }
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
             Class<?>[] argumentsTypes, Throwable t) {
-
+        if (allArguments[1] instanceof AsyncRequestProducerWrapper) {
+            ((AsyncRequestProducerWrapper) allArguments[1]).getExitSpan().fail(t);
+        }
     }
 }

apm-sniffer/apm-sdk-plugin/httpclient-5.x-plugin/src/main/java/org/apache/skywalking/apm/plugin/httpclient/v5/wrapper/AsyncRequestProducerWrapper.java

@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under 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 KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.skywalking.apm.plugin.httpclient.v5.wrapper;
+
+import java.io.IOException;
+import org.apache.hc.core5.http.HttpException;
+import org.apache.hc.core5.http.HttpRequest;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
+import org.apache.hc.core5.http.nio.DataStreamChannel;
+import org.apache.hc.core5.http.nio.RequestChannel;
+import org.apache.hc.core5.http.protocol.HttpContext;
+import org.apache.skywalking.apm.agent.core.context.CarrierItem;
+import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
+import org.apache.skywalking.apm.agent.core.context.ContextManager;
+import org.apache.skywalking.apm.agent.core.context.tag.Tags;
+import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
+import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
+import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
+import org.apache.skywalking.apm.plugin.httpclient.v5.AsyncExitSpan;
+
+public class AsyncRequestProducerWrapper implements AsyncRequestProducer {
+
+    private final AsyncRequestProducer producer;
+    private final AsyncExitSpan exitSpan;
+
+    public AsyncRequestProducerWrapper(AsyncRequestProducer producer, AsyncExitSpan exitSpan) {
+        this.producer = producer;
+        this.exitSpan = exitSpan;
+    }
+
+    public AsyncExitSpan getExitSpan() {
+        return exitSpan;
+    }
+
+    @Override
+    public void sendRequest(RequestChannel channel, HttpContext context) throws IOException, HttpException {
+        producer.sendRequest((request, entityDetails, requestContext) -> {
+            if (exitSpan.claimCreation()) {
+                try {
+                    startExitSpan(request);
+                } catch (Throwable ignored) {
+                    // Never let tracing instrumentation break the user's HTTP request.
+                }
+            }
+
+            channel.sendRequest(request, entityDetails, requestContext);
+        }, context);
+    }
+
+    private void startExitSpan(HttpRequest request) {
+        String operationName = request.getRequestUri();
+        String remotePeer = exitSpan.getTarget().toHostString();
+
+        ContextCarrier contextCarrier = new ContextCarrier();
+        AbstractSpan span = ContextManager.createExitSpan(
+                operationName,
+                contextCarrier,
+                remotePeer
+        );
+
+        boolean nested = ContextManager.activeSpan().isExit();
+
+        if (!nested) {
+            span.setComponent(ComponentsDefine.HTTP_ASYNC_CLIENT);
+            Tags.URL.set(span, request.getRequestUri());
+            SpanLayer.asHttp(span);
+        }
+
+        CarrierItem next = contextCarrier.items();
+        while (next.hasNext()) {
+            request.setHeader(next.getHeadKey(), next.getHeadValue());
+            next = next.next();
+        }
+
+        if (!nested) {
+            span.prepareForAsync();
+        }
+
+        ContextManager.stopSpan(span);
+
+        if (!nested) {
+            exitSpan.start(span);
+        }
+    }
+
+    @Override
+    public boolean isRepeatable() {
+        return producer.isRepeatable();
+    }
+
+    @Override
+    public void produce(DataStreamChannel channel) throws IOException {
+        producer.produce(channel);
+    }
+
+    @Override
+    public int available() {
+        return producer.available();
+    }
+
+    @Override
+    public void failed(Exception cause) {
+        producer.failed(cause);
+    }
+
+    @Override
+    public void releaseResources() {
+        producer.releaseResources();
+    }
+}