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..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 @@ -146,6 +146,18 @@ 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. + *

+ * 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; + } + /** * 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..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 @@ -153,6 +153,21 @@ void onSpanStarted(TracingSpan span) { TraceContextListener.getInstance().onSpanStarted(span); } + boolean detach(TracingSpan span) { + if (!spanStack.isEmpty() && spanStack.peek().equals(span)) { + spanStack.pop(); + return true; + } + + 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; + } + 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..18ad3ef6f7 --- /dev/null +++ b/agent/agent-observability/src/test/java/org/bithon/agent/observability/tracing/context/TraceSpanDetachTest.java @@ -0,0 +1,91 @@ +/* + * 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 testDetachReturnsCopiedContextForCurrentSpan() { + ITraceContext context = TraceContextFactory.newContext(SamplingMode.FULL); + ITraceSpan rootSpan = context.currentSpan().name("root"); + ITraceSpan childSpan = rootSpan.newChildSpan("child"); + + 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(); + + 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..ea29f51089 --- /dev/null +++ b/agent/agent-plugins-test/src/test/java/org/bithon/agent/plugins/test/httpclient/JavaNetHttpClientThreadPropagationInterceptorTest.java @@ -0,0 +1,231 @@ +/* + * 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 org.junit.jupiter.api.condition.EnabledForJreRange; +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; +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) +@EnabledForJreRange(min = JRE.JAVA_11) +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 = newTraceServer(traceParent); + + 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 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) { + 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..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 @@ -57,14 +57,11 @@ 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") + .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 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)); 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