Skip to content

Commit 49476f7

Browse files
PerfectSlayerPerfectSlayer
committed
feat(java-lang): Add support for VirtualThread
Add context tracking for VirtualThread API
1 parent 52803bf commit 49476f7

File tree

8 files changed

+421
-0
lines changed

8 files changed

+421
-0
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package datadog.trace.bootstrap.instrumentation.java.lang;
2+
3+
public final class VirtualThreadHelper {
4+
public static final String VIRTUAL_THREAD_CLASS_NAME = "java.lang.VirtualThread";
5+
/**
6+
* {@link datadog.trace.bootstrap.instrumentation.api.AgentScope} class name as string literal.
7+
* This is mandatory for {@link datadog.trace.bootstrap.ContextStore} API call.
8+
*/
9+
public static final String AGENT_SCOPE_CLASS_NAME =
10+
"datadog.trace.bootstrap.instrumentation.api.AgentScope";
11+
}

‎dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@
5050
0 java.lang.ProcessImpl
5151
# allow Runtime instrumentation for RASP
5252
0 java.lang.Runtime
53+
# allow context tracking for VirtualThread
54+
0 java.lang.VirtualThread
5355
0 java.net.http.*
5456
0 java.net.HttpURLConnection
5557
0 java.net.Socket

‎dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/build.gradle‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
plugins {
2+
id 'idea'
3+
}
4+
5+
apply from: "$rootDir/gradle/java.gradle"
6+
// Use slf4j-simple as default; logback has a high chance of getting stuck in a deadlock on CI.
7+
apply from: "$rootDir/gradle/slf4j-simple.gradle"
8+
9+
testJvmConstraints {
10+
minJavaVersion = JavaVersion.VERSION_21
11+
}
12+
13+
muzzle {
14+
pass {
15+
coreJdk('21')
16+
}
17+
}
18+
19+
idea {
20+
module {
21+
jdkName = '21'
22+
}
23+
}
24+
25+
// Set all compile tasks to use JDK21 but let instrumentation code targets 1.8 compatibility
26+
tasks.withType(AbstractCompile).configureEach {
27+
configureCompiler(it, 21, JavaVersion.VERSION_1_8)
28+
}
29+
30+
31+
// tasks.named("test", Test) {
32+
// jvmArgs = ['-Dnet.bytebuddy.dump=/Users/bruce.bujon/tmp/vt-dump']
33+
// }
34+
35+
dependencies {
36+
testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
37+
// testImplementation project(':dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-1.8')
38+
}
39+

‎dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/ContinuationInstrumentation.java‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package datadog.trace.instrumentation.java.lang.jdk21;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
5+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
6+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
7+
import static java.util.Collections.singletonMap;
8+
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
9+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
10+
11+
import com.google.auto.service.AutoService;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.bootstrap.InstrumentationContext;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
16+
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
17+
import java.util.Map;
18+
import net.bytebuddy.asm.Advice;
19+
import net.bytebuddy.asm.Advice.OnMethodEnter;
20+
import net.bytebuddy.asm.Advice.OnMethodExit;
21+
22+
@SuppressWarnings("unused")
23+
@AutoService(InstrumenterModule.class)
24+
public class ContinuationInstrumentation extends InstrumenterModule.Tracing
25+
implements Instrumenter.ForBootstrap,
26+
// Instrumenter.ForTypeHierarchy,
27+
Instrumenter.ForSingleType,
28+
Instrumenter.HasMethodAdvice {
29+
30+
private static final String VIRTUAL_THREAD_CLASS_NAME = "java.lang.VirtualThread";
31+
32+
public ContinuationInstrumentation() {
33+
super("java_lang", "continuation");
34+
}
35+
36+
@Override
37+
public String instrumentedType() {
38+
return VIRTUAL_THREAD_CLASS_NAME;
39+
}
40+
41+
// @Override
42+
// public String hierarchyMarkerType() {
43+
// return "jdk.internal.vm.Continuation";
44+
// }
45+
//
46+
// @Override
47+
// public ElementMatcher<TypeDescription> hierarchyMatcher() {
48+
// return extendsClass(named(hierarchyMarkerType()));
49+
// }
50+
51+
@Override
52+
public boolean isEnabled() {
53+
return false; // TODO REMOVE
54+
// return JavaVirtualMachine.isJavaVersionAtLeast(19) && super.isEnabled();
55+
}
56+
57+
@Override
58+
public Map<String, String> contextStore() {
59+
return singletonMap(VIRTUAL_THREAD_CLASS_NAME, State.class.getName());
60+
}
61+
62+
@Override
63+
public void methodAdvice(MethodTransformer transformer) {
64+
transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
65+
transformer.applyAdvice(
66+
isMethod().and(named("runContinuation")), getClass().getName() + "$RunContinuation");
67+
68+
// transformer.applyAdvice(isMethod().and(named("mount")), getClass().getName() + "$Activate");
69+
// transformer.applyAdvice(isMethod().and(named("unmount")), getClass().getName() + "Close");
70+
}
71+
72+
public static final class Construct {
73+
@OnMethodExit(suppress = Throwable.class)
74+
public static void captureScope(@Advice.This Object virtualThread) {
75+
capture(
76+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, State.class.getName()),
77+
virtualThread);
78+
}
79+
}
80+
81+
public static final class RunContinuation {
82+
@OnMethodEnter(suppress = Throwable.class)
83+
public static AgentScope activate(@Advice.This Object virtualThread) {
84+
// System.out.println(">>> INJECTED ENTER");
85+
return startTaskScope(
86+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, State.class.getName()),
87+
virtualThread);
88+
}
89+
90+
@OnMethodExit(suppress = Throwable.class)
91+
public static void close(@Advice.Enter AgentScope scope) {
92+
// System.out.println(">>> INJECTED AFTER");
93+
endTaskScope(scope);
94+
}
95+
}
96+
97+
// public static final class Capture {
98+
// @OnMethodExit(suppress = Throwable.class)
99+
// public static void captureScope(@Advice.This Object virtualThread) {
100+
// System.out.println(">>> INJECTED");
101+
// capture(InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, State.class.getName()),
102+
// virtualThread);
103+
// }
104+
// }
105+
//
106+
// public static final class Activate {
107+
// @OnMethodEnter(suppress = Throwable.class)
108+
// public static AgentScope activate(@Advice.This Object virtualThread) {
109+
// System.out.println(">>> INJECTED2");
110+
// return startTaskScope(InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME,
111+
// State.class.getName()), virtualThread);
112+
// }
113+
// }
114+
115+
// public static final class Close {
116+
// @OnMethodExit
117+
// public static void closeScope(@Advice.This Object continuation) {
118+
// endTaskScope(InstrumentationContext.get("jdk.internal.vm.Continuation", Stat));
119+
// }
120+
// }
121+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
package datadog.trace.instrumentation.java.lang.jdk21;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
5+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
6+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
7+
import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.AGENT_SCOPE_CLASS_NAME;
8+
import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.VIRTUAL_THREAD_CLASS_NAME;
9+
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
10+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
11+
12+
import com.google.auto.service.AutoService;
13+
import datadog.environment.JavaVirtualMachine;
14+
import datadog.trace.agent.tooling.Instrumenter;
15+
import datadog.trace.agent.tooling.InstrumenterModule;
16+
import datadog.trace.bootstrap.ContextStore;
17+
import datadog.trace.bootstrap.InstrumentationContext;
18+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
19+
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
import net.bytebuddy.asm.Advice;
23+
import net.bytebuddy.asm.Advice.OnMethodEnter;
24+
import net.bytebuddy.asm.Advice.OnMethodExit;
25+
26+
/**
27+
* Instruments {@code VirtualThread} to capture active state at creation, active it on continuation
28+
* mount, and close the scope from activation on continuation unmount.
29+
*
30+
* <p>The instrumentation reuses the context store from {@link Runnable} (as {@code VirtualThread}
31+
* inherits from {@link Runnable}) to store the {@link State} captured state to restore. It
32+
* additionally stores the {@link AgentScope} to be able to close it later as activation / close is
33+
* not done around the same method (so passing the scope from {@link OnMethodEnter} / {@link
34+
* OnMethodExit} using advice return value is not possible).
35+
*
36+
* <p>Instrumenting the internal `VirtualThread.runContinuation()` method does not work as the
37+
* current thread is still the carrier thread and not a virtual thread. Activating the state when on
38+
* the carrier thread (ie a platform thread) would store the active context into ThreadLocal using
39+
* the platform thread as key, making the tracer unable to retrieve the stored context from the
40+
* current virtual thread (ThreadLocal will not return the value associated to the underlying
41+
* platform thread as they are considered to be different).
42+
*/
43+
@SuppressWarnings("unused")
44+
@AutoService(InstrumenterModule.class)
45+
public final class VirtualThreadInstrumentation extends InstrumenterModule.Tracing
46+
implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
47+
48+
public VirtualThreadInstrumentation() {
49+
super("java-lang", "virtual-thread");
50+
}
51+
52+
@Override
53+
public String instrumentedType() {
54+
return VIRTUAL_THREAD_CLASS_NAME;
55+
}
56+
57+
@Override
58+
public boolean isEnabled() {
59+
return JavaVirtualMachine.isJavaVersionAtLeast(21) && super.isEnabled();
60+
}
61+
62+
@Override
63+
public Map<String, String> contextStore() {
64+
Map<String, String> contextStore = new HashMap<>();
65+
contextStore.put(Runnable.class.getName(), State.class.getName());
66+
contextStore.put(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
67+
return contextStore;
68+
}
69+
70+
@Override
71+
public void methodAdvice(MethodTransformer transformer) {
72+
transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
73+
transformer.applyAdvice(isMethod().and(named("mount")), getClass().getName() + "$Activate");
74+
transformer.applyAdvice(isMethod().and(named("unmount")), getClass().getName() + "$Close");
75+
}
76+
77+
public static final class Construct {
78+
@OnMethodExit(suppress = Throwable.class)
79+
public static void captureScope(@Advice.This Object virtualThread) {
80+
capture(InstrumentationContext.get(Runnable.class, State.class), (Runnable) virtualThread);
81+
}
82+
}
83+
84+
public static final class Activate {
85+
@OnMethodExit(suppress = Throwable.class)
86+
public static void activate(@Advice.This Object virtualThread) {
87+
ContextStore<Runnable, State> stateStore =
88+
InstrumentationContext.get(Runnable.class, State.class);
89+
ContextStore<Object, Object> scopeStore =
90+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
91+
AgentScope agentScope = startTaskScope(stateStore, (Runnable) virtualThread);
92+
scopeStore.put(virtualThread, agentScope);
93+
}
94+
}
95+
96+
public static final class Close {
97+
@OnMethodEnter(suppress = Throwable.class)
98+
public static void close(@Advice.This Object virtualThread) {
99+
ContextStore<Object, Object> scopeStore =
100+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
101+
Object agentScope = scopeStore.get(virtualThread);
102+
if (agentScope instanceof AgentScope) {
103+
endTaskScope((AgentScope) agentScope);
104+
}
105+
}
106+
}
107+
}

‎dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/groovy/JavaAsyncChild.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import datadog.trace.api.Trace;
2+
import java.util.concurrent.Callable;
3+
import java.util.concurrent.atomic.AtomicBoolean;
4+
5+
public class JavaAsyncChild implements Runnable, Callable<Void> {
6+
private final AtomicBoolean blockThread;
7+
private final boolean doTraceableWork;
8+
9+
public JavaAsyncChild() {
10+
this(true, false);
11+
}
12+
13+
public JavaAsyncChild(final boolean doTraceableWork, final boolean blockThread) {
14+
this.doTraceableWork = doTraceableWork;
15+
this.blockThread = new AtomicBoolean(blockThread);
16+
}
17+
18+
public void unblock() {
19+
blockThread.set(false);
20+
}
21+
22+
@Override
23+
public void run() {
24+
runImpl();
25+
}
26+
27+
@Override
28+
public Void call() {
29+
runImpl();
30+
return null;
31+
}
32+
33+
private void runImpl() {
34+
System.out.println(">>> Running on " + Thread.currentThread());
35+
while (blockThread.get()) {
36+
// busy-wait to block thread
37+
}
38+
if (doTraceableWork) {
39+
asyncChild();
40+
}
41+
System.out.println(">>> Completed on " + Thread.currentThread());
42+
}
43+
44+
@Trace(operationName = "asyncChild")
45+
private void asyncChild() {}
46+
}

0 commit comments

Comments
 (0)