diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java
new file mode 100644
index 00000000000..f51ee7d978c
--- /dev/null
+++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java
@@ -0,0 +1,12 @@
+package datadog.trace.bootstrap.instrumentation.java.lang;
+
+public final class VirtualThreadHelper {
+ public static final String VIRTUAL_THREAD_CLASS_NAME = "java.lang.VirtualThread";
+
+ /**
+ * {@link datadog.trace.bootstrap.instrumentation.api.AgentScope} class name as string literal.
+ * This is mandatory for {@link datadog.trace.bootstrap.ContextStore} API call.
+ */
+ public static final String AGENT_SCOPE_CLASS_NAME =
+ "datadog.trace.bootstrap.instrumentation.api.AgentScope";
+}
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 df1d38cd392..5430117b39a 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
@@ -50,6 +50,8 @@
0 java.lang.ProcessImpl
# allow Runtime instrumentation for RASP
0 java.lang.Runtime
+# allow context tracking for VirtualThread
+0 java.lang.VirtualThread
0 java.net.http.*
0 java.net.HttpURLConnection
0 java.net.Socket
diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/main/java/datadog/trace/instrumentation/java/concurrent/virtualthread/TaskRunnerInstrumentation.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/main/java/datadog/trace/instrumentation/java/concurrent/virtualthread/TaskRunnerInstrumentation.java
index e4e43fa9835..a53e7657282 100644
--- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/main/java/datadog/trace/instrumentation/java/concurrent/virtualthread/TaskRunnerInstrumentation.java
+++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/main/java/datadog/trace/instrumentation/java/concurrent/virtualthread/TaskRunnerInstrumentation.java
@@ -17,11 +17,14 @@
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
import java.util.Map;
import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.OnMethodEnter;
+import net.bytebuddy.asm.Advice.OnMethodExit;
/**
* Instruments {@code TaskRunner}, internal runnable for {@code ThreadPerTaskExecutor} (JDK 19+ as
* preview, 21+ as stable), the executor with default virtual thread factory.
*/
+@SuppressWarnings("unused")
@AutoService(InstrumenterModule.class)
public final class TaskRunnerInstrumentation extends InstrumenterModule.Tracing
implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
@@ -51,19 +54,19 @@ public void methodAdvice(MethodTransformer transformer) {
}
public static final class Construct {
- @Advice.OnMethodExit
+ @OnMethodExit(suppress = Throwable.class)
public static void captureScope(@Advice.This Runnable task) {
capture(InstrumentationContext.get(Runnable.class, State.class), task);
}
}
public static final class Run {
- @Advice.OnMethodEnter
+ @OnMethodEnter(suppress = Throwable.class)
public static AgentScope activate(@Advice.This Runnable task) {
return startTaskScope(InstrumentationContext.get(Runnable.class, State.class), task);
}
- @Advice.OnMethodExit(onThrowable = Throwable.class)
+ @OnMethodExit(suppress = Throwable.class)
public static void close(@Advice.Enter AgentScope scope) {
endTaskScope(scope);
}
diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadTest.groovy b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadPerTaskExecutorTest.groovy
similarity index 97%
rename from dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadTest.groovy
rename to dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadPerTaskExecutorTest.groovy
index 7bb36c7f024..cbdc42a4993 100644
--- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadTest.groovy
+++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-21.0/src/test/groovy/VirtualThreadPerTaskExecutorTest.groovy
@@ -1,14 +1,13 @@
import datadog.trace.agent.test.InstrumentationSpecification
import datadog.trace.api.Trace
import datadog.trace.core.DDSpan
-import spock.lang.Shared
-
import java.util.concurrent.Callable
import java.util.concurrent.ExecutorCompletionService
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit
+import spock.lang.Shared
-class VirtualThreadTest extends InstrumentationSpecification {
+class VirtualThreadPerTaskExecutorTest extends InstrumentationSpecification {
@Shared
def executeRunnable = { e, c -> e.execute((Runnable) c) }
@Shared
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
new file mode 100644
index 00000000000..00757ef832a
--- /dev/null
+++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle
@@ -0,0 +1,32 @@
+plugins {
+ id 'idea'
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+// Use slf4j-simple as default; logback has a high chance of getting stuck in a deadlock on CI.
+apply from: "$rootDir/gradle/slf4j-simple.gradle"
+
+testJvmConstraints {
+ minJavaVersion = JavaVersion.VERSION_21
+}
+
+muzzle {
+ pass {
+ coreJdk('21')
+ }
+}
+
+idea {
+ module {
+ jdkName = '21'
+ }
+}
+
+// Set all compile tasks to use JDK21 but let instrumentation code targets 1.8 compatibility
+tasks.withType(AbstractCompile).configureEach {
+ configureCompiler(it, 21, JavaVersion.VERSION_1_8)
+}
+
+dependencies {
+ testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
+}
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
new file mode 100644
index 00000000000..f590299e30e
--- /dev/null
+++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java
@@ -0,0 +1,107 @@
+package datadog.trace.instrumentation.java.lang.jdk21;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
+import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.AGENT_SCOPE_CLASS_NAME;
+import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.VIRTUAL_THREAD_CLASS_NAME;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.environment.JavaVirtualMachine;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.OnMethodEnter;
+import net.bytebuddy.asm.Advice.OnMethodExit;
+
+/**
+ * Instruments {@code VirtualThread} to capture active state at creation, activate it on
+ * continuation mount, and close the scope from activation on continuation unmount.
+ *
+ *
The instrumentation uses two context stores. The first from {@link Runnable} (as {@code
+ * VirtualThread} inherits from {@link Runnable}) to store the captured {@link State} to restore
+ * later. It additionally stores the {@link AgentScope} to be able to close it later as activation /
+ * close is not done around the same method (so passing the scope from {@link OnMethodEnter} /
+ * {@link OnMethodExit} using advice return value is not possible).
+ *
+ *
Instrumenting the internal {@code VirtualThread.runContinuation()} method does not work as the
+ * current thread is still the carrier thread and not a virtual thread. Activating the state when on
+ * the carrier thread (ie a platform thread) would store the active context into ThreadLocal using
+ * the platform thread as key, making the tracer unable to retrieve the stored context from the
+ * current virtual thread (ThreadLocal will not return the value associated to the underlying
+ * platform thread as they are considered to be different).
+ */
+@SuppressWarnings("unused")
+@AutoService(InstrumenterModule.class)
+public final class VirtualThreadInstrumentation extends InstrumenterModule.Tracing
+ implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+ public VirtualThreadInstrumentation() {
+ super("java-lang", "java-lang-21", "virtual-thread");
+ }
+
+ @Override
+ public String instrumentedType() {
+ return VIRTUAL_THREAD_CLASS_NAME;
+ }
+
+ @Override
+ public boolean isEnabled() {
+ return JavaVirtualMachine.isJavaVersionAtLeast(21) && super.isEnabled();
+ }
+
+ @Override
+ public Map contextStore() {
+ Map contextStore = new HashMap<>();
+ contextStore.put(Runnable.class.getName(), State.class.getName());
+ contextStore.put(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
+ return contextStore;
+ }
+
+ @Override
+ public void methodAdvice(MethodTransformer transformer) {
+ transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
+ transformer.applyAdvice(isMethod().and(named("mount")), getClass().getName() + "$Activate");
+ transformer.applyAdvice(isMethod().and(named("unmount")), getClass().getName() + "$Close");
+ }
+
+ public static final class Construct {
+ @OnMethodExit(suppress = Throwable.class)
+ public static void captureScope(@Advice.This Object virtualThread) {
+ capture(InstrumentationContext.get(Runnable.class, State.class), (Runnable) virtualThread);
+ }
+ }
+
+ public static final class Activate {
+ @OnMethodExit(suppress = Throwable.class)
+ public static void activate(@Advice.This Object virtualThread) {
+ ContextStore stateStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ ContextStore