diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java index 597595496ab..964e3d87274 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java @@ -5,6 +5,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; import datadog.trace.api.Config; +import datadog.trace.api.GenericClassValue; import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; @@ -31,6 +32,15 @@ * context listener, so we simply swap in on mount and out on unmount. */ public final class VirtualThreadState { + private static final ClassValue PROPAGATE_CONTEXT = + GenericClassValue.of( + type -> !type.getName().equals("jdk.internal.net.http.HttpClientImpl$SelectorManager")); + + /** The HTTP selector belongs to its client, not to the request creating the client. */ + public static boolean shouldPropagateContext(Runnable task) { + return PROPAGATE_CONTEXT.get(task.getClass()); + } + // CWS observes scope changes through listeners, so retain context swaps on mount and unmount. private static final boolean USE_PER_MOUNT_CONTEXT = isJavaVersion(21) diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index e719f4684f1..bf582679d58 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -107,6 +107,9 @@ 0 sun.rmi.server.* 0 sun.rmi.transport.* 0 sun.nio.ch.FileChannelImpl +0 sun.nio.ch.Poller$Pollers +0 sun.nio.ch.Poller$VThreadsPollerGroup +0 sun.nio.ch.Poller$PollerPerCarrierPollerGroup # USM needs it 0 sun.security.ssl.* 0 javax.net.ssl.SSLSocket diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle index 9fbe0263a3c..e9e62bfc021 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle @@ -29,3 +29,9 @@ tasks.named("compileTestJava", JavaCompile) { dependencies { testImplementation project(':dd-java-agent:instrumentation:datadog:tracing:trace-annotation') } + +tasks.named('forkedTest', Test) { + jvmArgs '--add-opens=java.base/sun.nio.ch=ALL-UNNAMED', + '--add-opens=java.base/java.util.concurrent=ALL-UNNAMED', + '--add-opens=java.net.http/jdk.internal.net.http=ALL-UNNAMED' +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java new file mode 100644 index 00000000000..2b4d878a8a6 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/PollerInstrumentation.java @@ -0,0 +1,64 @@ +package datadog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.rootContext; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.context.Context; +import datadog.environment.JavaVirtualMachine; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import net.bytebuddy.asm.Advice; + +/** Keeps JDK I/O pollers independent of the request that first needs them. */ +@AutoService(InstrumenterModule.class) +public final class PollerInstrumentation extends InstrumenterModule.ContextTracking + implements Instrumenter.ForBootstrap, Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice { + public PollerInstrumentation() { + super("java-lang", "java-lang-21", "virtual-thread"); + } + + @Override + public boolean isEnabled() { + return JavaVirtualMachine.isJavaVersionAtLeast(22) && super.isEnabled(); + } + + @Override + public String[] knownMatchingTypes() { + return new String[] { + "sun.nio.ch.Poller$Pollers", + "sun.nio.ch.Poller$VThreadsPollerGroup", + "sun.nio.ch.Poller$PollerPerCarrierPollerGroup" + }; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("start") + .and(takesArguments(0)) + .and(returns(void.class)) + .or( + named("startReadPoller") + .and(takesArguments(0)) + .and(returns(named("sun.nio.ch.Poller")))), + getClass().getName() + "$StartAdvice"); + } + + public static final class StartAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static Context enter() { + // Swap clears the raw context even when the caller has reached the scope-depth limit. + return rootContext().swap(); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void exit(@Advice.Enter Context previous) { + if (previous != null) { + previous.swap(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java index d1ad6e30ebc..ef60ec9ba90 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java @@ -26,6 +26,7 @@ import datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState; import java.util.Collection; import java.util.Map; +import java.util.concurrent.Executor; import net.bytebuddy.asm.Advice; import net.bytebuddy.asm.Advice.OnMethodEnter; import net.bytebuddy.asm.Advice.OnMethodExit; @@ -105,7 +106,10 @@ public Map contextStore() { @Override public void methodAdvice(MethodTransformer transformer) { - transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct"); + transformer.applyAdvice( + isConstructor() + .and(takesArguments(Executor.class, String.class, int.class, Runnable.class)), + getClass().getName() + "$Construct"); // JDK 22+ enters run(Runnable) after the first mount with the virtual thread current. JDK 21 // update releases differ, so they retain context swaps on every mount and unmount. if (JavaVirtualMachine.isJavaVersionAtLeast(22)) { @@ -125,11 +129,15 @@ public void methodAdvice(MethodTransformer transformer) { public static final class Construct { @OnMethodExit(suppress = Throwable.class) - public static void afterInit(@Advice.This Object virtualThread) { + public static void afterInit( + @Advice.This Object virtualThread, @Advice.Argument(3) Runnable task) { Context context = currentContext(); if (context == rootContext()) { return; // No active context to propagate, avoid creating state } + if (!VirtualThreadState.shouldPropagateContext(task)) { + return; + } VirtualThreadState state = new VirtualThreadState(context, context.capture()); ContextStore store = InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME); diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/HttpClientSelectorContextForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/HttpClientSelectorContextForkedTest.java new file mode 100644 index 00000000000..e00a17c5b79 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/HttpClientSelectorContextForkedTest.java @@ -0,0 +1,51 @@ +package testdog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; +import static testdog.trace.instrumentation.java.lang.jdk21.PollerContextForkedTest.field; +import static testdog.trace.instrumentation.java.lang.jdk21.PollerContextForkedTest.virtualThreadState; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.net.http.HttpClient; +import org.junit.jupiter.api.Test; + +class HttpClientSelectorContextForkedTest extends AbstractInstrumentationTest { + @Test + void selectorDoesNotRetainCreatingRequest() throws Exception { + assumeTrue(Runtime.version().feature() >= 26); + AgentSpan parent = startSpan("test", "client-owner"); + try (HttpClient client = newHttpClient(parent)) { + try (ContextScope ignored = activateSpan(parent)) { + Object impl = field(client.getClass(), "impl").get(client); + Thread selector = (Thread) field(impl.getClass(), "selmgrThread").get(impl); + assertTrue(selector.isVirtual()); + assertTrue(selector.isAlive()); + assertNull(virtualThreadState(selector), "live selector must not retain request context"); + assertSame(parent, activeSpan()); + } finally { + parent.finish(); + } + // Publish while the client is still alive: closing it must not be needed to release parent. + assertTraces(trace(span().root().operationName("client-owner"))); + } + } + + private static HttpClient newHttpClient(Object parentSpan) { + AgentSpan parent = (AgentSpan) parentSpan; + try (ContextScope ignored = activateSpan(parent)) { + return HttpClient.newHttpClient(); + } catch (Throwable error) { + parent.finish(); + throw error; + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PerCarrierPollerContextForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PerCarrierPollerContextForkedTest.java new file mode 100644 index 00000000000..fbad192e2b3 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PerCarrierPollerContextForkedTest.java @@ -0,0 +1,100 @@ +package testdog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; +import static testdog.trace.instrumentation.java.lang.jdk21.PollerContextForkedTest.field; +import static testdog.trace.instrumentation.java.lang.jdk21.PollerContextForkedTest.virtualThreadState; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.test.util.PollingConditions; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.time.Duration; +import java.util.Set; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.OS; + +@EnabledOnOs({OS.LINUX, OS.MAC}) +class PerCarrierPollerContextForkedTest extends AbstractInstrumentationTest { + @Test + void firstReadOnCarrierDoesNotRetainRequest() throws Exception { + assumeTrue(Runtime.version().feature() >= 27); + // A separate fork is needed: poller mode is fixed by Poller.. + System.setProperty("jdk.pollerMode", "3"); + System.setProperty("jdk.writePollers", "1"); + assertFalse(GlobalIgnores.isIgnored("sun.nio.ch.Poller$PollerPerCarrierPollerGroup", false)); + + AgentSpan parent = startSpan("test", "carrier-poller-owner"); + Object parentSpan = parent; + AtomicReference failure = new AtomicReference<>(); + Thread reader = null; + try (ContextScope ignored = activateSpan(parent)) { + Class poller = Class.forName("sun.nio.ch.Poller"); + Object group = field(poller, "POLLER_GROUP").get(null); + @SuppressWarnings("unchecked") + Set readPollers = (Set) field(group.getClass(), "readPollers").get(group); + assertTrue(readPollers.isEmpty(), "read pollers must be created lazily by I/O"); + + try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress()); + Socket socket = new Socket(InetAddress.getLoopbackAddress(), server.getLocalPort()); + Socket peer = server.accept()) { + socket.setSoTimeout(10_000); + reader = + Thread.ofVirtual() + .start( + () -> { + try { + assertSame(parentSpan, activeSpan()); + assertEquals(1, socket.getInputStream().read()); + assertSame( + parentSpan, + activeSpan(), + "poller creation must restore request context"); + } catch (Throwable error) { + failure.set(error); + } + }); + + new PollingConditions(10) + .eventually( + () -> { + assertFalse(readPollers.isEmpty()); + for (Object readPoller : readPollers) { + Thread owner = (Thread) field(poller, "owner").get(readPoller); + assertTrue(owner != null && owner.isAlive()); + assertTrue(owner.isVirtual()); + assertNull( + virtualThreadState(owner), + "live carrier poller must retain neither context nor continuation"); + } + }); + peer.getOutputStream().write(1); + assertTrue(reader.join(Duration.ofSeconds(10))); + assertNull(failure.get()); + assertSame(parent, activeSpan()); + } + } finally { + // Closing the sockets also releases the reader if an assertion above fails. + if (reader != null) { + reader.join(Duration.ofSeconds(10)); + } + parent.finish(); + } + assertTraces(trace(span().root().operationName("carrier-poller-owner"))); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PollerContextForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PollerContextForkedTest.java new file mode 100644 index 00000000000..8e690e7e845 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/PollerContextForkedTest.java @@ -0,0 +1,75 @@ +package testdog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores; +import datadog.trace.bootstrap.FieldBackedContextStores; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.lang.reflect.Field; +import java.util.Set; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledForJreRange; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.JRE; +import org.junit.jupiter.api.condition.OS; + +@EnabledOnOs({OS.LINUX, OS.MAC}) +class PollerContextForkedTest extends AbstractInstrumentationTest { + @Test + @EnabledForJreRange(min = JRE.JAVA_22) + void firstPollerInitializationDoesNotRetainRequest() throws Exception { + // Set these before Poller.; changing them afterward cannot reconfigure its pollers. + // The fork isolates the properties and poller initialization from other tests. + System.setProperty("jdk.pollerMode", "2"); + System.setProperty("jdk.readPollers", "1"); + System.setProperty("jdk.writePollers", "1"); + assertFalse(GlobalIgnores.isIgnored("sun.nio.ch.Poller$Pollers", false)); + assertFalse(GlobalIgnores.isIgnored("sun.nio.ch.Poller$VThreadsPollerGroup", false)); + assertTrue(GlobalIgnores.isIgnored("sun.nio.ch.Poller", false)); + AgentSpan parent = startSpan("test", "poller-owner"); + try (ContextScope ignored = activateSpan(parent)) { + Class poller = Class.forName("sun.nio.ch.Poller"); + Object pollers = + field(poller, Runtime.version().feature() >= 27 ? "POLLER_GROUP" : "POLLERS").get(null); + Object executor = field(pollers.getClass(), "executor").get(pollers); + @SuppressWarnings("unchecked") + Set threads = (Set) field(executor.getClass(), "threads").get(executor); + assertEquals(2, threads.size(), "one read poller and one write poller must be running"); + for (Thread thread : threads) { + assertTrue(thread.isVirtual()); + assertTrue(thread.isAlive()); + assertNull( + virtualThreadState(thread), "poller must retain neither context nor continuation"); + } + assertSame(parent, activeSpan(), "poller startup must restore the caller context"); + } finally { + parent.finish(); + } + assertTraces(trace(span().root().operationName("poller-owner"))); + } + + static Object virtualThreadState(Thread thread) { + int storeId = + FieldBackedContextStores.getContextStoreId( + "java.lang.VirtualThread", + "datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState"); + return FieldBackedContextStores.getContextStore(storeId).get(thread); + } + + static Field field(Class type, String name) throws NoSuchFieldException { + Field field = type.getDeclaredField(name); + field.setAccessible(true); + return field; + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java index 2c20740f4dc..0ff9c66cd72 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java @@ -4,13 +4,20 @@ import static datadog.trace.agent.test.assertions.SpanMatcher.span; import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import datadog.context.ContextScope; import datadog.trace.agent.test.AbstractInstrumentationTest; import datadog.trace.api.CorrelationIdentifier; import datadog.trace.api.GlobalTracer; import datadog.trace.api.Trace; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState; import java.time.Duration; import java.util.ArrayList; @@ -29,6 +36,26 @@ void testContextPathSelectedForRuntime() { assertEquals(isJavaVersion(21), VirtualThreadState.usePerMountContext()); } + @Test + void propagationDoesNotDependOnInheritableThreadLocals() throws InterruptedException { + AgentSpan parent = startSpan("test", "parent"); + try (ContextScope ignored = activateSpan(parent)) { + for (boolean inheritThreadLocals : new boolean[] {true, false}) { + AtomicReference inherited = new AtomicReference<>(); + Thread child = + Thread.ofVirtual() + .inheritInheritableThreadLocals(inheritThreadLocals) + .start(() -> inherited.set(activeSpan())); + assertTrue(child.join(TIMEOUT)); + assertSame(parent, inherited.get()); + assertSame(parent, activeSpan()); + } + } finally { + parent.finish(); + } + assertTraces(trace(span().root().operationName("parent"))); + } + @DisplayName("test context restored after virtual thread remounts") @Test void testContextRestoredAfterVirtualThreadRemount() {