From 7f44867b5363107484586c3f50573c792e7086f1 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 1 Jul 2026 17:09:24 +0800 Subject: [PATCH 1/4] Fix java.net.http async trace propagation --- .../tracing/context/ITraceSpan.java | 8 + .../tracing/context/impl/TracingContext.java | 13 + .../tracing/context/impl/TracingSpan.java | 18 ++ .../tracing/context/TraceSpanDetachTest.java | 94 ++++++++ ...lientThreadPropagationInterceptorTest.java | 222 ++++++++++++++++++ .../javanethttp/JavaNetHttpClientPlugin.java | 10 +- .../interceptor/HttpClientImpl$SendAsync.java | 25 +- ...va => HttpRequestImpl$GetUserHeaders.java} | 10 +- 8 files changed, 376 insertions(+), 24 deletions(-) create mode 100644 agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java create mode 100644 agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java rename agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/{HttpRequest$Headers.java => HttpRequestImpl$GetUserHeaders.java} (85%) diff --git a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java index 2f69c4ce7d..f747002c80 100644 --- a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java +++ b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java @@ -146,6 +146,14 @@ default ITraceSpan method(Class clazz, String method) { ITraceSpan start(); + /** + * Detach this span from its current trace context and return a span with the same identity + * in a copied context so it can be finished asynchronously. + */ + default ITraceSpan detach() { + return this; + } + /** * Finish a span. * The implementation should guarantee that multiple calls on this method are safe. diff --git a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java index 8046037fe1..0aea925bfd 100644 --- a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java +++ b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java @@ -153,6 +153,19 @@ void onSpanStarted(TracingSpan span) { TraceContextListener.getInstance().onSpanStarted(span); } + boolean detach(TracingSpan span) { + if (!spanStack.isEmpty() && spanStack.peek().equals(span)) { + spanStack.pop(); + return true; + } + + LoggerFactory.getLogger(TracingContext.class).warn("Try to detach a span which is not in the stack. This IS a bug.\n" + + "Current span: \n{}, Unfinished Spans:\n{}", + span, + spanStack); + return false; + } + void onSpanFinished(TracingSpan span) { try { TraceContextListener.getInstance().onSpanFinished(span); diff --git a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingSpan.java b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingSpan.java index 73c4766d84..61d5bb1763 100644 --- a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingSpan.java +++ b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingSpan.java @@ -178,6 +178,24 @@ public TracingSpan start() { return this; } + @Override + public ITraceSpan detach() { + tracingContext.detach(this); + TracingContext detachedContext = (TracingContext) tracingContext.copy(); + TracingSpan detachedSpan = (TracingSpan) detachedContext.newSpan(parentSpanId, spanId); + + detachedSpan.startTime = this.startTime; + detachedSpan.endTime = this.endTime; + detachedSpan.kind = this.kind; + detachedSpan.component = this.component; + detachedSpan.parentApplication = this.parentApplication; + detachedSpan.clazz = this.clazz; + detachedSpan.method = this.method; + detachedSpan.tags.putAll(this.tags); + + return detachedSpan; + } + @Override public void finish() { if (this.endTime != 0) { diff --git a/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java b/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java new file mode 100644 index 0000000000..fd186017ec --- /dev/null +++ b/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java @@ -0,0 +1,94 @@ +/* + * Copyright 2020 bithon.org + * + * Licensed 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.bithon.agent.observability.tracing.context; + +import org.bithon.agent.configuration.ConfigurationManager; +import org.bithon.agent.observability.tracing.Tracer; +import org.bithon.agent.observability.tracing.config.TraceConfig; +import org.bithon.agent.observability.tracing.id.impl.DefaultSpanIdGenerator; +import org.bithon.agent.observability.tracing.reporter.ITraceReporter; +import org.bithon.agent.observability.tracing.reporter.ReporterConfig; +import org.bithon.agent.observability.tracing.sampler.SamplingMode; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Field; +import java.util.List; + +public class TraceSpanDetachTest { + + private Tracer originalTracer; + + @BeforeEach + public void setUp() throws Exception { + ConfigurationManager.createForTesting(); + originalTracer = setTracer(new Tracer("test", "test") + .traceConfig(new TraceConfig()) + .spanIdGenerator(new DefaultSpanIdGenerator()) + .reporter(new ITraceReporter() { + @Override + public ReporterConfig getReporterConfig() { + return new ReporterConfig(); + } + + @Override + public void report(List spans) { + } + })); + } + + @AfterEach + public void tearDown() throws Exception { + TraceContextHolder.detach(); + setTracer(originalTracer); + } + + @Test + public void testDetachReturnsCopiedContextEvenWhenSpanIsNotStackTop() { + ITraceContext context = TraceContextFactory.newContext(SamplingMode.FULL); + ITraceSpan rootSpan = context.currentSpan().name("root"); + ITraceSpan childSpan = rootSpan.newChildSpan("child"); + ITraceSpan grandchildSpan = childSpan.newChildSpan("grandchild"); + + ITraceSpan detachedSpan = childSpan.detach(); + + Assertions.assertNotSame(childSpan, detachedSpan); + Assertions.assertNotSame(childSpan.context(), detachedSpan.context()); + Assertions.assertEquals(childSpan.traceId(), detachedSpan.traceId()); + Assertions.assertEquals(childSpan.spanId(), detachedSpan.spanId()); + Assertions.assertEquals(childSpan.parentSpanId(), detachedSpan.parentSpanId()); + Assertions.assertSame(detachedSpan, detachedSpan.context().currentSpan()); + + detachedSpan.finish(); + detachedSpan.context().finish(); + + grandchildSpan.finish(); + childSpan.finish(); + rootSpan.finish(); + context.finish(); + } + + private static Tracer setTracer(Tracer tracer) throws Exception { + Field field = Tracer.class.getDeclaredField("INSTANCE"); + field.setAccessible(true); + Tracer previous = (Tracer) field.get(null); + field.set(null, tracer); + return previous; + } +} diff --git a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java new file mode 100644 index 0000000000..4893f6da79 --- /dev/null +++ b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java @@ -0,0 +1,222 @@ +/* + * Copyright 2020 bithon.org + * + * Licensed 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.bithon.agent.plugins.test.httpclient; + +import com.sun.net.httpserver.HttpServer; +import net.bytebuddy.agent.ByteBuddyAgent; +import org.bithon.agent.instrumentation.aop.interceptor.plugin.IPlugin; +import org.bithon.agent.java.adaptor.Java9Adaptor; +import org.bithon.agent.java.adaptor.JavaAdaptorFactory; +import org.bithon.agent.observability.tracing.context.ITraceContext; +import org.bithon.agent.observability.tracing.context.ITraceSpan; +import org.bithon.agent.observability.tracing.context.TraceContextFactory; +import org.bithon.agent.observability.tracing.context.TraceContextHolder; +import org.bithon.agent.observability.tracing.context.propagation.w3c.W3CTraceContextHeader; +import org.bithon.agent.observability.tracing.sampler.SamplingMode; +import org.bithon.agent.plugin.httpclient.javanethttp.JavaNetHttpClientPlugin; +import org.bithon.agent.plugins.test.AbstractPluginInterceptorTest; +import org.bithon.component.commons.tracing.Tags; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; + +import java.io.Closeable; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; + +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +public class JavaNetHttpClientThreadPropagationInterceptorTest extends AbstractPluginInterceptorTest { + + private static final Object INTERCEPTOR_INSTALL_LOCK = new Object(); + private static boolean interceptorInstalled; + + @Override + protected IPlugin[] getPlugins() { + return new IPlugin[]{new JavaNetHttpClientPlugin()}; + } + + @Test + @Order(0) + @Override + public void testInterceptorInstallation() { + installInterceptor(); + } + + private void installInterceptor() { + synchronized (INTERCEPTOR_INSTALL_LOCK) { + if (interceptorInstalled) { + return; + } + + JavaAdaptorFactory.setAdaptor(new Java9Adaptor(ByteBuddyAgent.getInstrumentation())); + super.testInterceptorInstallation(); + interceptorInstalled = true; + } + } + + @Test + public void testSendAsyncPropagatesWorkerAsyncTaskSpanContext() throws Exception { + installInterceptor(); + + AtomicReference traceParent = new AtomicReference<>(); + HttpServer server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + server.createContext("/trace", exchange -> { + traceParent.set(exchange.getRequestHeaders().getFirst(W3CTraceContextHeader.TRACE_HEADER_PARENT)); + byte[] response = "OK".getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, response.length); + exchange.getResponseBody().write(response); + exchange.close(); + }); + server.start(); + + ITraceContext ctx = TraceContextFactory.newContext(SamplingMode.FULL); + ctx.currentSpan().name("ROOT"); + String rootSpanId = ctx.currentSpan().spanId(); + + try (TracingExecutor executor = new TracingExecutor()) { + try (Closeable ignored = TraceContextHolder::detach) { + TraceContextHolder.attach(ctx); + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/trace")) + .GET() + .build(); + + CompletableFuture> responseFuture = HttpClient.newBuilder() + .executor(executor) + .build() + .sendAsync(request, HttpResponse.BodyHandlers.ofString()); + + Assertions.assertEquals(rootSpanId, TraceContextHolder.current().currentSpan().spanId()); + + HttpResponse response = responseFuture.get(); + + Assertions.assertEquals(200, response.statusCode()); + } + ctx.currentSpan().finish(); + ctx.finish(); + } finally { + server.stop(0); + } + + Assertions.assertNotNull(traceParent.get()); + String propagatedParentSpanId = parentSpanId(traceParent.get()); + Assertions.assertNotEquals(rootSpanId, propagatedParentSpanId); + + ITraceSpan asyncTaskSpan = awaitSpan(span -> Objects.equals(propagatedParentSpanId, span.spanId())); + ITraceSpan httpClientSpan = awaitSpan(span -> "http-client".equals(span.name()) && "sendAsync".equals(span.method())); + ITraceSpan threadPoolSpan = awaitSpan(span -> Objects.equals(asyncTaskSpan.parentSpanId(), span.spanId())); + + Assertions.assertEquals("async-task", asyncTaskSpan.name()); + Assertions.assertEquals("thread-pool", threadPoolSpan.name()); + Assertions.assertEquals("execute", threadPoolSpan.method()); + Assertions.assertEquals(rootSpanId, httpClientSpan.parentSpanId()); + Assertions.assertEquals(httpClientSpan.spanId(), threadPoolSpan.parentSpanId()); + Assertions.assertEquals("HttpClient-test-Worker", asyncTaskSpan.tags().get(Tags.Thread.NAME)); + Assertions.assertNotEquals(httpClientSpan.spanId(), propagatedParentSpanId); + } + + private static ITraceSpan awaitSpan(Predicate predicate) throws InterruptedException { + long deadline = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadline) { + synchronized (REPORTED_SPANS) { + for (ITraceSpan span : REPORTED_SPANS) { + if (predicate.test(span)) { + return span; + } + } + } + Thread.sleep(10); + } + + Assertions.fail("Expected span not reported: " + REPORTED_SPANS); + return null; + } + + private static String parentSpanId(String traceParent) { + return traceParent.split("-")[2]; + } + + private static class TracingExecutor implements Executor, Closeable { + private final ExecutorService delegate = Executors.newFixedThreadPool(1, runnable -> new Thread(runnable, "HttpClient-test-Worker")); + + @Override + public void execute(Runnable command) { + ITraceSpan threadPoolSpan = TraceContextFactory.newSpan("thread-pool"); + if (threadPoolSpan != null) { + threadPoolSpan.method(ExecutorService.class.getName(), "execute") + .start(); + } + + ITraceSpan taskSpan; + try { + taskSpan = TraceContextFactory.newAsyncSpan("async-task"); + if (taskSpan != null) { + taskSpan.method(command.getClass().getName(), "run"); + } + } finally { + if (threadPoolSpan != null) { + threadPoolSpan.finish(); + } + } + + delegate.execute(() -> { + if (taskSpan == null) { + command.run(); + return; + } + + Throwable exception = null; + TraceContextHolder.attach(taskSpan.context()); + taskSpan.start(); + try { + command.run(); + } catch (Throwable e) { + exception = e; + throw e; + } finally { + Thread currentThread = Thread.currentThread(); + taskSpan.tag(Tags.Thread.NAME, currentThread.getName()) + .tag(Tags.Thread.ID, currentThread.getId()) + .tag(exception) + .finish(); + taskSpan.context().finish(); + TraceContextHolder.detach(); + } + }); + } + + @Override + public void close() { + delegate.shutdownNow(); + } + } +} diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java index 8b4ab320fe..d65dca947a 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java @@ -57,14 +57,10 @@ public List getInterceptors() { .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpClientImpl$SendAsync") .build(), - // Intercept HttpRequest.headers() to propagate tracing headers + // Intercept the JDK's send-time user headers access to propagate tracing headers forClass("jdk.internal.net.http.HttpRequestImpl") - .onMethod("headers") - .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequest$Headers") - .build(), - forClass("jdk.internal.net.http.ImmutableHttpRequest") - .onMethod("headers") - .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequest$Headers") + .onMethod("getUserHeaders") + .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequestImpl$GetUserHeaders") .build() ); } diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java index aefed8eba4..fda84a66c0 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java @@ -35,10 +35,9 @@ import java.net.http.HttpResponse; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.Executor; /** - * {@link jdk.internal.net.http.HttpClientImpl#sendAsync(HttpRequest, HttpResponse.BodyHandler, HttpResponse.PushPromiseHandler, Executor)} + * {@link jdk.internal.net.http.HttpClientImpl#sendAsync(HttpRequest, HttpResponse.BodyHandler, HttpResponse.PushPromiseHandler, java.util.concurrent.Executor)} * * @author frank.chen021@outlook.com * @date 2024/12/19 @@ -60,7 +59,7 @@ public InterceptionDecision before(AopContext aopContext) { HttpRequest request = aopContext.getArgAs(0); // Start tracing span - ITraceSpan span = TraceContextFactory.newAsyncSpan("http-client"); + ITraceSpan span = TraceContextFactory.newSpan("http-client"); if (span != null) { span.method(aopContext.getTargetClass(), aopContext.getMethod()) .kind(SpanKind.CLIENT) @@ -85,7 +84,8 @@ public void after(AopContext aopContext) { final String url = request.uri().toString(); final String httpMethod = request.method(); - final ITraceSpan span = aopContext.getUserContext(); + ITraceSpan span = aopContext.getUserContext(); + if (span != null) { span.tag(Tags.Http.URL, url) .tag(Tags.Http.METHOD, httpMethod); @@ -99,11 +99,12 @@ public void after(AopContext aopContext) { // Finish span with exception if (span != null) { span.finish(); - span.context().finish(); } return; } + final ITraceSpan httpClientSpan = span == null ? null : span.detach(); + // For async operations, we need to hook into the returned CompletableFuture CompletableFuture> returnFuture = aopContext.getReturningAs(); @@ -126,23 +127,23 @@ public void after(AopContext aopContext) { .updateIOMetrics(HttpClientUtils.getRequestSize(request), HttpClientUtils.getResponseSize(response)); // Add response headers to trace - if (span != null) { + if (httpClientSpan != null) { if (traceConfig.getHeaders() != null) { for (String headerName : traceConfig.getHeaders().getResponse()) { Optional headerValue = response.headers().firstValue(headerName); - headerValue.ifPresent(s -> span.tag(Tags.Http.RESPONSE_HEADER_PREFIX + headerName, s)); + headerValue.ifPresent(s -> httpClientSpan.tag(Tags.Http.RESPONSE_HEADER_PREFIX + headerName, s)); } } - span.tag(Tags.Http.STATUS, Integer.toString(statusCode)); + httpClientSpan.tag(Tags.Http.STATUS, Integer.toString(statusCode)); } } // Finish span - if (span != null) { - span.tag(throwable) - .finish(); - span.context().finish(); + if (httpClientSpan != null) { + httpClientSpan.tag(throwable) + .finish(); + httpClientSpan.context().finish(); } } catch (Throwable t) { // Catch all to prevent any exception from thrown from agent to user code diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequest$Headers.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java similarity index 85% rename from agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequest$Headers.java rename to agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java index a196d0080e..3434e3ff64 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequest$Headers.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java @@ -24,18 +24,18 @@ import org.bithon.agent.observability.tracing.context.TraceMode; import java.net.http.HttpHeaders; -import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; /** - * {@link jdk.internal.net.http.HttpRequestImpl#headers()} - * {@link jdk.internal.net.http.ImmutableHttpRequest#headers()} + * {@link jdk.internal.net.http.HttpRequestImpl#getUserHeaders()} * * @author frank.chen021@outlook.com * @date 11/9/25 6:22 pm */ -public class HttpRequest$Headers extends AfterInterceptor { +public class HttpRequestImpl$GetUserHeaders extends AfterInterceptor { + @Override public void after(AopContext aopContext) { if (aopContext.hasException()) { @@ -47,7 +47,7 @@ public void after(AopContext aopContext) { } HttpHeaders headers = aopContext.getReturningAs(); - final Map> headerMap = new HashMap<>(headers.map()); + final Map> headerMap = new LinkedHashMap<>(headers.map()); ctx.propagate(headerMap, (carrier, key, value) -> carrier.put(key, List.of(value))); aopContext.setReturning(HttpHeaders.of(headerMap, (a, b) -> true)); From e1097813f928bd7bbd4b4addce52401d9bbec5ac Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 1 Jul 2026 17:34:03 +0800 Subject: [PATCH 2/4] Skip java.net.http thread propagation test before JDK 11 --- .../JavaNetHttpClientThreadPropagationInterceptorTest.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java index 4893f6da79..60ab71abf0 100644 --- a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java +++ b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java @@ -35,6 +35,8 @@ import org.junit.jupiter.api.Order; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestMethodOrder; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.JRE; import java.io.Closeable; import java.net.InetAddress; @@ -53,6 +55,7 @@ import java.util.function.Predicate; @TestMethodOrder(MethodOrderer.OrderAnnotation.class) +@EnabledForJreRange(min = JRE.JAVA_11) public class JavaNetHttpClientThreadPropagationInterceptorTest extends AbstractPluginInterceptorTest { private static final Object INTERCEPTOR_INSTALL_LOCK = new Object(); From de9e536e4be700e2d957b6a8a944dd952a8269e2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 1 Jul 2026 18:18:43 +0800 Subject: [PATCH 3/4] Address java.net.http async propagation review --- .../tracing/context/ITraceSpan.java | 4 + .../tracing/context/impl/TracingContext.java | 6 +- .../tracing/context/TraceSpanDetachTest.java | 5 +- ...lientThreadPropagationInterceptorTest.java | 81 ++++++++++++++++--- .../javanethttp/JavaNetHttpClientPlugin.java | 5 ++ .../interceptor/HttpClientImpl$SendAsync.java | 14 ++++ .../interceptor/HttpRequestImpl$Ctor.java | 44 ++++++++++ .../HttpRequestImpl$GetUserHeaders.java | 18 ++++- .../HttpRequestPropagationContext.java | 59 ++++++++++++++ pom.xml | 11 ++- 10 files changed, 229 insertions(+), 18 deletions(-) create mode 100644 agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java create mode 100644 agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java diff --git a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java index f747002c80..6951b3784d 100644 --- a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java +++ b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/ITraceSpan.java @@ -149,6 +149,10 @@ default ITraceSpan method(Class clazz, String method) { /** * Detach this span from its current trace context and return a span with the same identity * in a copied context so it can be finished asynchronously. + *

+ * This should be called only when this span is the current span of its context. After + * detaching, callers should finish the returned span and leave the original span reference + * unfinished. */ default ITraceSpan detach() { return this; diff --git a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java index 0aea925bfd..8d141803ab 100644 --- a/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java +++ b/agent/agent-observability/src/main/java/org/bithon/agent/observability/tracing/context/impl/TracingContext.java @@ -159,9 +159,11 @@ boolean detach(TracingSpan span) { return true; } - LoggerFactory.getLogger(TracingContext.class).warn("Try to detach a span which is not in the stack. This IS a bug.\n" - + "Current span: \n{}, Unfinished Spans:\n{}", + ITraceSpan currentSpan = spanStack.isEmpty() ? null : spanStack.peek(); + LoggerFactory.getLogger(TracingContext.class).warn("Try to detach a span which is not the current stack top. This IS a bug.\n" + + "Span to detach: \n{}, Current stack top: \n{}, Unfinished Spans:\n{}", span, + currentSpan, spanStack); return false; } diff --git a/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java b/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java index fd186017ec..18ad3ef6f7 100644 --- a/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java +++ b/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java @@ -60,11 +60,10 @@ public void tearDown() throws Exception { } @Test - public void testDetachReturnsCopiedContextEvenWhenSpanIsNotStackTop() { + public void testDetachReturnsCopiedContextForCurrentSpan() { ITraceContext context = TraceContextFactory.newContext(SamplingMode.FULL); ITraceSpan rootSpan = context.currentSpan().name("root"); ITraceSpan childSpan = rootSpan.newChildSpan("child"); - ITraceSpan grandchildSpan = childSpan.newChildSpan("grandchild"); ITraceSpan detachedSpan = childSpan.detach(); @@ -78,8 +77,6 @@ public void testDetachReturnsCopiedContextEvenWhenSpanIsNotStackTop() { detachedSpan.finish(); detachedSpan.context().finish(); - grandchildSpan.finish(); - childSpan.finish(); rootSpan.finish(); context.finish(); } diff --git a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java index 60ab71abf0..d181a3ac2d 100644 --- a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java +++ b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java @@ -39,6 +39,7 @@ import org.junit.jupiter.api.condition.JRE; import java.io.Closeable; +import java.io.IOException; import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.URI; @@ -90,15 +91,7 @@ public void testSendAsyncPropagatesWorkerAsyncTaskSpanContext() throws Exception installInterceptor(); AtomicReference traceParent = new AtomicReference<>(); - HttpServer server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); - server.createContext("/trace", exchange -> { - traceParent.set(exchange.getRequestHeaders().getFirst(W3CTraceContextHeader.TRACE_HEADER_PARENT)); - byte[] response = "OK".getBytes(StandardCharsets.UTF_8); - exchange.sendResponseHeaders(200, response.length); - exchange.getResponseBody().write(response); - exchange.close(); - }); - server.start(); + HttpServer server = newTraceServer(traceParent); ITraceContext ctx = TraceContextFactory.newContext(SamplingMode.FULL); ctx.currentSpan().name("ROOT"); @@ -147,6 +140,62 @@ public void testSendAsyncPropagatesWorkerAsyncTaskSpanContext() throws Exception Assertions.assertNotEquals(httpClientSpan.spanId(), propagatedParentSpanId); } + @Test + public void testSendAsyncFallsBackToClientSpanContextWithoutWorkerTraceContext() throws Exception { + installInterceptor(); + + AtomicReference traceParent = new AtomicReference<>(); + HttpServer server = newTraceServer(traceParent); + + ITraceContext ctx = TraceContextFactory.newContext(SamplingMode.FULL); + ctx.currentSpan().name("ROOT"); + String rootSpanId = ctx.currentSpan().spanId(); + + try (PlainExecutor executor = new PlainExecutor()) { + try (Closeable ignored = TraceContextHolder::detach) { + TraceContextHolder.attach(ctx); + + HttpRequest request = HttpRequest.newBuilder() + .uri(URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/trace")) + .GET() + .build(); + + CompletableFuture> responseFuture = HttpClient.newBuilder() + .executor(executor) + .build() + .sendAsync(request, HttpResponse.BodyHandlers.ofString()); + + Assertions.assertEquals(rootSpanId, TraceContextHolder.current().currentSpan().spanId()); + + HttpResponse response = responseFuture.get(); + + Assertions.assertEquals(200, response.statusCode()); + } + ctx.currentSpan().finish(); + ctx.finish(); + } finally { + server.stop(0); + } + + Assertions.assertNotNull(traceParent.get()); + ITraceSpan httpClientSpan = awaitSpan(span -> "http-client".equals(span.name()) && "sendAsync".equals(span.method())); + + Assertions.assertEquals(httpClientSpan.spanId(), parentSpanId(traceParent.get())); + } + + private static HttpServer newTraceServer(AtomicReference traceParent) throws IOException { + HttpServer server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + server.createContext("/trace", exchange -> { + traceParent.set(exchange.getRequestHeaders().getFirst(W3CTraceContextHeader.TRACE_HEADER_PARENT)); + byte[] response = "OK".getBytes(StandardCharsets.UTF_8); + exchange.sendResponseHeaders(200, response.length); + exchange.getResponseBody().write(response); + exchange.close(); + }); + server.start(); + return server; + } + private static ITraceSpan awaitSpan(Predicate predicate) throws InterruptedException { long deadline = System.currentTimeMillis() + 5_000; while (System.currentTimeMillis() < deadline) { @@ -168,6 +217,20 @@ private static String parentSpanId(String traceParent) { return traceParent.split("-")[2]; } + private static class PlainExecutor implements Executor, Closeable { + private final ExecutorService delegate = Executors.newFixedThreadPool(1, runnable -> new Thread(runnable, "HttpClient-plain-Worker")); + + @Override + public void execute(Runnable command) { + delegate.execute(command); + } + + @Override + public void close() { + delegate.shutdownNow(); + } + } + private static class TracingExecutor implements Executor, Closeable { private final ExecutorService delegate = Executors.newFixedThreadPool(1, runnable -> new Thread(runnable, "HttpClient-test-Worker")); diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java index d65dca947a..ed1cd3b830 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java @@ -59,7 +59,12 @@ public List getInterceptors() { // Intercept the JDK's send-time user headers access to propagate tracing headers forClass("jdk.internal.net.http.HttpRequestImpl") + .onConstructor() + .andRawArgs("java.net.http.HttpRequest", "java.net.ProxySelector") + .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequestImpl$Ctor") + .onMethod("getUserHeaders") + .andArgsSize(0) .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequestImpl$GetUserHeaders") .build() ); diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java index fda84a66c0..95c29a7de6 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java @@ -23,6 +23,7 @@ import org.bithon.agent.instrumentation.aop.interceptor.declaration.AroundInterceptor; import org.bithon.agent.observability.metric.domain.httpclient.HttpOutgoingMetricsRegistry; import org.bithon.agent.observability.tracing.config.TraceConfig; +import org.bithon.agent.observability.tracing.context.ITraceContext; import org.bithon.agent.observability.tracing.context.ITraceSpan; import org.bithon.agent.observability.tracing.context.TraceContextFactory; import org.bithon.agent.plugin.httpclient.javanethttp.utils.HttpClientUtils; @@ -73,6 +74,7 @@ public InterceptionDecision before(AopContext aopContext) { }); aopContext.setUserContext(span.start()); + savePropagationContext(request, span); } return InterceptionDecision.CONTINUE; @@ -100,10 +102,12 @@ public void after(AopContext aopContext) { if (span != null) { span.finish(); } + clearPropagationContext(request); return; } final ITraceSpan httpClientSpan = span == null ? null : span.detach(); + clearPropagationContext(request); // For async operations, we need to hook into the returned CompletableFuture CompletableFuture> returnFuture = aopContext.getReturningAs(); @@ -154,4 +158,14 @@ public void after(AopContext aopContext) { // Replace the original future with our wrapped one aopContext.setReturning(wrappedFuture); } + + private static void savePropagationContext(HttpRequest request, ITraceSpan span) { + ITraceContext context = span.context().copy(); + context.newSpan(span.parentSpanId(), span.spanId()); + HttpRequestPropagationContext.set(request, context); + } + + private static void clearPropagationContext(HttpRequest request) { + HttpRequestPropagationContext.remove(request); + } } diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java new file mode 100644 index 0000000000..8fbd6670bf --- /dev/null +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java @@ -0,0 +1,44 @@ +/* + * Copyright 2020 bithon.org + * + * Licensed 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.bithon.agent.plugin.httpclient.javanethttp.interceptor; + +import org.bithon.agent.instrumentation.aop.IBithonObject; +import org.bithon.agent.instrumentation.aop.context.AopContext; +import org.bithon.agent.instrumentation.aop.interceptor.declaration.AfterInterceptor; +import org.bithon.agent.observability.tracing.context.ITraceContext; + +/** + * {@link jdk.internal.net.http.HttpRequestImpl#HttpRequestImpl(java.net.http.HttpRequest, java.net.ProxySelector)} + * + * @author frank.chen021@outlook.com + * @date 2026/7/1 + */ +public class HttpRequestImpl$Ctor extends AfterInterceptor { + + @Override + public void after(AopContext aopContext) { + if (aopContext.hasException()) { + return; + } + + Object source = aopContext.getArgAs(0); + ITraceContext context = HttpRequestPropagationContext.get(source); + if (context != null) { + ((IBithonObject) aopContext.getTarget()).setInjectedObject(context); + } + } +} diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java index 3434e3ff64..3200bbddd4 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java @@ -41,8 +41,8 @@ public void after(AopContext aopContext) { if (aopContext.hasException()) { return; } - ITraceContext ctx = TraceContextHolder.current(); - if (ctx == null || ctx.traceMode() == TraceMode.LOGGING) { + ITraceContext ctx = getPropagationContext(aopContext); + if (ctx == null) { return; } HttpHeaders headers = aopContext.getReturningAs(); @@ -52,4 +52,18 @@ public void after(AopContext aopContext) { aopContext.setReturning(HttpHeaders.of(headerMap, (a, b) -> true)); } + + private static ITraceContext getPropagationContext(AopContext aopContext) { + ITraceContext ctx = TraceContextHolder.current(); + if (isPropagatable(ctx)) { + return ctx; + } + + ctx = HttpRequestPropagationContext.get(aopContext.getTarget()); + return isPropagatable(ctx) ? ctx : null; + } + + private static boolean isPropagatable(ITraceContext ctx) { + return ctx != null && ctx.traceMode() != TraceMode.LOGGING && ctx.currentSpan() != null; + } } diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java new file mode 100644 index 0000000000..8f78f4eb72 --- /dev/null +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java @@ -0,0 +1,59 @@ +/* + * Copyright 2020 bithon.org + * + * Licensed 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.bithon.agent.plugin.httpclient.javanethttp.interceptor; + +import org.bithon.agent.instrumentation.aop.IBithonObject; +import org.bithon.agent.observability.tracing.context.ITraceContext; + +import java.net.http.HttpRequest; +import java.util.Collections; +import java.util.Map; +import java.util.WeakHashMap; + +final class HttpRequestPropagationContext { + private static final Map CONTEXTS = Collections.synchronizedMap(new WeakHashMap<>()); + + private HttpRequestPropagationContext() { + } + + static void set(HttpRequest request, ITraceContext context) { + if (request instanceof IBithonObject) { + ((IBithonObject) request).setInjectedObject(context); + } + CONTEXTS.put(request, context); + } + + static ITraceContext get(Object request) { + if (request instanceof IBithonObject) { + Object injected = ((IBithonObject) request).getInjectedObject(); + if (injected instanceof ITraceContext) { + return (ITraceContext) injected; + } + } + if (request instanceof HttpRequest) { + return CONTEXTS.get(request); + } + return null; + } + + static void remove(HttpRequest request) { + if (request instanceof IBithonObject) { + ((IBithonObject) request).setInjectedObject(null); + } + CONTEXTS.remove(request); + } +} diff --git a/pom.xml b/pom.xml index af4df77743..c83aca20a9 100644 --- a/pom.xml +++ b/pom.xml @@ -56,7 +56,6 @@ component agent - server @@ -491,6 +490,16 @@ + + server + + [21,) + + + server + + + shaded From e4a9ede503a96eae042ecbc46d47c454bece5ed3 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 2 Jul 2026 10:07:26 +0800 Subject: [PATCH 4/4] Remove java.net.http async fallback propagation --- ...lientThreadPropagationInterceptorTest.java | 57 ------------------ .../javanethttp/JavaNetHttpClientPlugin.java | 4 -- .../interceptor/HttpClientImpl$SendAsync.java | 14 ----- .../interceptor/HttpRequestImpl$Ctor.java | 44 -------------- .../HttpRequestImpl$GetUserHeaders.java | 18 +----- .../HttpRequestPropagationContext.java | 59 ------------------- 6 files changed, 2 insertions(+), 194 deletions(-) delete mode 100644 agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java delete mode 100644 agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java diff --git a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java index d181a3ac2d..ea29f51089 100644 --- a/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java +++ b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java @@ -140,49 +140,6 @@ public void testSendAsyncPropagatesWorkerAsyncTaskSpanContext() throws Exception Assertions.assertNotEquals(httpClientSpan.spanId(), propagatedParentSpanId); } - @Test - public void testSendAsyncFallsBackToClientSpanContextWithoutWorkerTraceContext() throws Exception { - installInterceptor(); - - AtomicReference traceParent = new AtomicReference<>(); - HttpServer server = newTraceServer(traceParent); - - ITraceContext ctx = TraceContextFactory.newContext(SamplingMode.FULL); - ctx.currentSpan().name("ROOT"); - String rootSpanId = ctx.currentSpan().spanId(); - - try (PlainExecutor executor = new PlainExecutor()) { - try (Closeable ignored = TraceContextHolder::detach) { - TraceContextHolder.attach(ctx); - - HttpRequest request = HttpRequest.newBuilder() - .uri(URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/trace")) - .GET() - .build(); - - CompletableFuture> responseFuture = HttpClient.newBuilder() - .executor(executor) - .build() - .sendAsync(request, HttpResponse.BodyHandlers.ofString()); - - Assertions.assertEquals(rootSpanId, TraceContextHolder.current().currentSpan().spanId()); - - HttpResponse response = responseFuture.get(); - - Assertions.assertEquals(200, response.statusCode()); - } - ctx.currentSpan().finish(); - ctx.finish(); - } finally { - server.stop(0); - } - - Assertions.assertNotNull(traceParent.get()); - ITraceSpan httpClientSpan = awaitSpan(span -> "http-client".equals(span.name()) && "sendAsync".equals(span.method())); - - Assertions.assertEquals(httpClientSpan.spanId(), parentSpanId(traceParent.get())); - } - private static HttpServer newTraceServer(AtomicReference traceParent) throws IOException { HttpServer server = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); server.createContext("/trace", exchange -> { @@ -217,20 +174,6 @@ private static String parentSpanId(String traceParent) { return traceParent.split("-")[2]; } - private static class PlainExecutor implements Executor, Closeable { - private final ExecutorService delegate = Executors.newFixedThreadPool(1, runnable -> new Thread(runnable, "HttpClient-plain-Worker")); - - @Override - public void execute(Runnable command) { - delegate.execute(command); - } - - @Override - public void close() { - delegate.shutdownNow(); - } - } - private static class TracingExecutor implements Executor, Closeable { private final ExecutorService delegate = Executors.newFixedThreadPool(1, runnable -> new Thread(runnable, "HttpClient-test-Worker")); diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java index ed1cd3b830..13d969e9b2 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/JavaNetHttpClientPlugin.java @@ -59,10 +59,6 @@ public List getInterceptors() { // Intercept the JDK's send-time user headers access to propagate tracing headers forClass("jdk.internal.net.http.HttpRequestImpl") - .onConstructor() - .andRawArgs("java.net.http.HttpRequest", "java.net.ProxySelector") - .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequestImpl$Ctor") - .onMethod("getUserHeaders") .andArgsSize(0) .interceptedBy("org.bithon.agent.plugin.httpclient.javanethttp.interceptor.HttpRequestImpl$GetUserHeaders") diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java index 95c29a7de6..fda84a66c0 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpClientImpl$SendAsync.java @@ -23,7 +23,6 @@ import org.bithon.agent.instrumentation.aop.interceptor.declaration.AroundInterceptor; import org.bithon.agent.observability.metric.domain.httpclient.HttpOutgoingMetricsRegistry; import org.bithon.agent.observability.tracing.config.TraceConfig; -import org.bithon.agent.observability.tracing.context.ITraceContext; import org.bithon.agent.observability.tracing.context.ITraceSpan; import org.bithon.agent.observability.tracing.context.TraceContextFactory; import org.bithon.agent.plugin.httpclient.javanethttp.utils.HttpClientUtils; @@ -74,7 +73,6 @@ public InterceptionDecision before(AopContext aopContext) { }); aopContext.setUserContext(span.start()); - savePropagationContext(request, span); } return InterceptionDecision.CONTINUE; @@ -102,12 +100,10 @@ public void after(AopContext aopContext) { if (span != null) { span.finish(); } - clearPropagationContext(request); return; } final ITraceSpan httpClientSpan = span == null ? null : span.detach(); - clearPropagationContext(request); // For async operations, we need to hook into the returned CompletableFuture CompletableFuture> returnFuture = aopContext.getReturningAs(); @@ -158,14 +154,4 @@ public void after(AopContext aopContext) { // Replace the original future with our wrapped one aopContext.setReturning(wrappedFuture); } - - private static void savePropagationContext(HttpRequest request, ITraceSpan span) { - ITraceContext context = span.context().copy(); - context.newSpan(span.parentSpanId(), span.spanId()); - HttpRequestPropagationContext.set(request, context); - } - - private static void clearPropagationContext(HttpRequest request) { - HttpRequestPropagationContext.remove(request); - } } diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java deleted file mode 100644 index 8fbd6670bf..0000000000 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$Ctor.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * Copyright 2020 bithon.org - * - * Licensed 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.bithon.agent.plugin.httpclient.javanethttp.interceptor; - -import org.bithon.agent.instrumentation.aop.IBithonObject; -import org.bithon.agent.instrumentation.aop.context.AopContext; -import org.bithon.agent.instrumentation.aop.interceptor.declaration.AfterInterceptor; -import org.bithon.agent.observability.tracing.context.ITraceContext; - -/** - * {@link jdk.internal.net.http.HttpRequestImpl#HttpRequestImpl(java.net.http.HttpRequest, java.net.ProxySelector)} - * - * @author frank.chen021@outlook.com - * @date 2026/7/1 - */ -public class HttpRequestImpl$Ctor extends AfterInterceptor { - - @Override - public void after(AopContext aopContext) { - if (aopContext.hasException()) { - return; - } - - Object source = aopContext.getArgAs(0); - ITraceContext context = HttpRequestPropagationContext.get(source); - if (context != null) { - ((IBithonObject) aopContext.getTarget()).setInjectedObject(context); - } - } -} diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java index 3200bbddd4..3434e3ff64 100644 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java +++ b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestImpl$GetUserHeaders.java @@ -41,8 +41,8 @@ public void after(AopContext aopContext) { if (aopContext.hasException()) { return; } - ITraceContext ctx = getPropagationContext(aopContext); - if (ctx == null) { + ITraceContext ctx = TraceContextHolder.current(); + if (ctx == null || ctx.traceMode() == TraceMode.LOGGING) { return; } HttpHeaders headers = aopContext.getReturningAs(); @@ -52,18 +52,4 @@ public void after(AopContext aopContext) { aopContext.setReturning(HttpHeaders.of(headerMap, (a, b) -> true)); } - - private static ITraceContext getPropagationContext(AopContext aopContext) { - ITraceContext ctx = TraceContextHolder.current(); - if (isPropagatable(ctx)) { - return ctx; - } - - ctx = HttpRequestPropagationContext.get(aopContext.getTarget()); - return isPropagatable(ctx) ? ctx : null; - } - - private static boolean isPropagatable(ITraceContext ctx) { - return ctx != null && ctx.traceMode() != TraceMode.LOGGING && ctx.currentSpan() != null; - } } diff --git a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java b/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java deleted file mode 100644 index 8f78f4eb72..0000000000 --- a/agent/agent-plugins/httpclient/java-net-http/src/main/java/org/bithon/agent/plugin/httpclient/javanethttp/interceptor/HttpRequestPropagationContext.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright 2020 bithon.org - * - * Licensed 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.bithon.agent.plugin.httpclient.javanethttp.interceptor; - -import org.bithon.agent.instrumentation.aop.IBithonObject; -import org.bithon.agent.observability.tracing.context.ITraceContext; - -import java.net.http.HttpRequest; -import java.util.Collections; -import java.util.Map; -import java.util.WeakHashMap; - -final class HttpRequestPropagationContext { - private static final Map CONTEXTS = Collections.synchronizedMap(new WeakHashMap<>()); - - private HttpRequestPropagationContext() { - } - - static void set(HttpRequest request, ITraceContext context) { - if (request instanceof IBithonObject) { - ((IBithonObject) request).setInjectedObject(context); - } - CONTEXTS.put(request, context); - } - - static ITraceContext get(Object request) { - if (request instanceof IBithonObject) { - Object injected = ((IBithonObject) request).getInjectedObject(); - if (injected instanceof ITraceContext) { - return (ITraceContext) injected; - } - } - if (request instanceof HttpRequest) { - return CONTEXTS.get(request); - } - return null; - } - - static void remove(HttpRequest request) { - if (request instanceof IBithonObject) { - ((IBithonObject) request).setInjectedObject(null); - } - CONTEXTS.remove(request); - } -}