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/AsyncRequestSpans.java

@@ -0,0 +1,140 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicReference;
+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;
+
+/**
+ * Per-request async exit span, owned by the request itself rather than by whatever thread happens to be running
+ * when a callback fires.
+ *
+ * <p>The span is created once, on the caller thread inside {@code doExecute}, while the caller's tracing context is
+ * still active. It is then immediately detached via {@link AbstractSpan#prepareForAsync()} +
+ * {@code ContextManager.stopSpan(span)} so it never sits on any thread's active-span stack while the request is in
+ * flight. From that point on it is finished exactly once, by reference, from whichever lifecycle callback gets
+ * there first (I/O thread response consumer, or the future callback on the caller/business thread) — never by a
+ * parameterless {@code ContextManager.stopSpan()} that would blindly pop whatever span is currently active on that
+ * thread.
+ *
+ * <p>All mutating operations are synchronized: {@link #onResponse(int)} (tagging, typically the I/O thread) can
+ * otherwise race with {@link #finish()} / {@link #fail(Throwable)} (typically the response-consumer or callback
+ * thread) finishing and clearing the span in the same window. An {@link AtomicReference} alone would prevent a
+ * double-finish but not a tag-write racing a finish.
+ */
+public class AsyncRequestSpans {
+
+    private final HttpHost target;
+
+    /**
+     * Only true, and only once, on the thread that is still inside {@code doExecute} when the request producer
+     * hands the concrete request to the channel. Any other thread (a custom {@code AsyncRequestProducer} that
+     * defers sending) has no relationship to the caller's context, so it must not create a span.
+     */
+    private final AtomicReference<Thread> creator = new AtomicReference<>(Thread.currentThread());
+
+    private AbstractSpan span;
+    private boolean finished;
+
+    public AsyncRequestSpans(HttpHost target) {
+        this.target = target;
+    }
+
+    public HttpHost getTarget() {
+        return target;
+    }
+
+    /**
+     * Claims the right to create the span. Returns {@code true} at most once, and only for the thread that
+     * constructed this holder (the {@code doExecute} caller thread).
+     */
+    public boolean claimCreation() {
+    Thread current = Thread.currentThread();
+    return creator.compareAndSet(current, null);
+    }
+
+    /**
+     * Called at the end of {@code doExecute} (success or failure) so a late/duplicate send from the same thread
+     * cannot still claim creation after the caller has moved on.
+     */
+    public void callerReturned() {
+        creator.set(null);
+    }
+
+    /**
+     * Stores the span. Must be called only after the span has already been detached with
+     * {@code prepareForAsync()} + {@code ContextManager.stopSpan(span)} — this class never touches the active-span
+     * stack itself.
+     */
+    public synchronized void start(AbstractSpan span) {
+        this.span = span;
+    }
+
+    public synchronized void onResponse(int statusCode) {
+        if (span == null || finished) {
+            return;
+        }
+
+        Tags.HTTP_RESPONSE_STATUS_CODE.set(span, statusCode);
+
+        if (statusCode >= 400) {
+            span.errorOccurred();
+        }
+    }
+
+    /** The whole response completed successfully. */
+    public synchronized void finish() {
+        end(false, null);
+    }
+
+    /** The exchange failed with an exception. */
+    public synchronized void fail(Throwable cause) {
+        end(true, cause);
+    }
+
+    /**
+     * Cancelled, or resources released before the response ever completed (e.g. a redirect exec that declines to
+     * resend a non-repeatable entity and never invokes {@code completed()}). Only takes effect if the span is
+     * still open — the normal-completion paths already finished it earlier, so this is then a no-op.
+     */
+    public synchronized void abort() {
+        end(true, null);
+    }
+
+    private void end(boolean error, Throwable cause) {
+        if (span == null || finished) {
+            return;
+        }
+
+        finished = true;
+
+        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,93 @@
 
 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;
-
+/**
+ * Intercepts the internal {@code doExecute(HttpHost, AsyncRequestProducer, AsyncResponseConsumer, ..., FutureCallback)}
+ * overload shared by every async client implementation (Internal*AsyncClient, Minimal*AsyncClient, and the
+ * classic-facade adapter), whose argument order/types are identical across HttpClient 5.0 through 5.6.
+ *
+ * <p>Unlike the previous implementation, this interceptor never stores anything in the {@code HttpContext} and
+ * never wraps a callback purely to call a parameterless {@code ContextManager.stopSpan()}. It only:
+ * <ol>
+ *   <li>creates a per-request {@link AsyncRequestSpans} holder, while the caller's context is still active;</li>
+ *   <li>wraps the request producer so the exit span is created on the caller thread, synchronously, the moment the
+ *   concrete {@code HttpRequest} becomes available;</li>
+ *   <li>wraps the response consumer and future callback so the retained span is finished by reference.</li>
+ * </ol>
+ * Because span creation no longer depends on the {@code HttpContext}, this also fixes HttpClient 5.4+, where the
+ * context argument passed by the classic facade and by {@code execute(SimpleHttpRequest, FutureCallback)} is
+ * {@code null}.
+ */
 public class HttpAsyncClientDoExecuteInterceptor implements InstanceMethodsAroundInterceptor {
 
+    private static final int TARGET_INDEX = 0;
+    private static final int REQUEST_PRODUCER_INDEX = 1;
+    private static final int RESPONSE_CONSUMER_INDEX = 2;
+    private static final int CALLBACK_INDEX = 5;
+
     @Override
-    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
-            MethodInterceptResult result) throws Throwable {
-        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());
+    public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
+        Class<?>[] argumentsTypes, MethodInterceptResult result) throws Throwable {
+        if (!ContextManager.isActive()) {
+            return;
         }
+        if (!(allArguments[REQUEST_PRODUCER_INDEX] instanceof AsyncRequestProducer)
+            || !(allArguments[RESPONSE_CONSUMER_INDEX] instanceof AsyncResponseConsumer)) {
+            return;
+        }
+
+        final HttpHost target = allArguments[TARGET_INDEX] instanceof HttpHost
+            ? (HttpHost) allArguments[TARGET_INDEX] : null;
+        final AsyncRequestSpans spans = new AsyncRequestSpans(target);
+
+        allArguments[REQUEST_PRODUCER_INDEX] = new AsyncRequestProducerWrapper(
+            (AsyncRequestProducer) allArguments[REQUEST_PRODUCER_INDEX], spans);
+        allArguments[RESPONSE_CONSUMER_INDEX] = new AsyncResponseConsumerWrapper<>(
+            (AsyncResponseConsumer<?>) allArguments[RESPONSE_CONSUMER_INDEX], spans);
+        // Wrap even when the caller passed null: it's the only lifecycle hook that sees cancellation and the
+        // synchronous-failure-before-consumer-runs path for callers who supplied no callback of their own.
+        allArguments[CALLBACK_INDEX] = new FutureCallbackWrapper<>(
+            (FutureCallback<?>) allArguments[CALLBACK_INDEX], spans);
     }
 
     @Override
-    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
-            Object ret) throws Throwable {
+    public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments,
+        Class<?>[] argumentsTypes, Object ret) throws Throwable {
+        releaseCreationClaim(allArguments);
         return ret;
     }
 
     @Override
     public void handleMethodException(EnhancedInstance objInst, Method method, Object[] allArguments,
-            Class<?>[] argumentsTypes, Throwable t) {
+        Class<?>[] argumentsTypes, Throwable t) {
+        if (allArguments[REQUEST_PRODUCER_INDEX] instanceof AsyncRequestProducerWrapper) {
+            AsyncRequestProducerWrapper wrapper = (AsyncRequestProducerWrapper) allArguments[REQUEST_PRODUCER_INDEX];
+            wrapper.getSpans().fail(t);
+        }
+        releaseCreationClaim(allArguments);
+    }
 
+    /**
+     * Once {@code doExecute} has returned (or thrown), no thread other than a genuinely deferred custom producer
+     * has any business claiming span creation — clearing this here keeps {@link AsyncRequestSpans#claimCreation()}
+     * honest even if the same thread somehow re-enters.
+     */
+    private void releaseCreationClaim(Object[] allArguments) {
+        if (allArguments[REQUEST_PRODUCER_INDEX] instanceof AsyncRequestProducerWrapper) {
+            ((AsyncRequestProducerWrapper) allArguments[REQUEST_PRODUCER_INDEX]).getSpans().callerReturned();
+        }
     }
 }