Skip to content

Commit 7da065b

Browse files
PerfectSlayerPerfectSlayer
committed
feat(java-lang): Add support for VirtualThread
Add context tracking for VirtualThread API
1 parent 3b02385 commit 7da065b

File tree

8 files changed

+440
-0
lines changed

8 files changed

+440
-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: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
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.InstrumentationContext;
17+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
18+
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
import net.bytebuddy.asm.Advice;
22+
import net.bytebuddy.asm.Advice.OnMethodEnter;
23+
import net.bytebuddy.asm.Advice.OnMethodExit;
24+
25+
/**
26+
* Instruments {@code VirtualThread} to capture active state at creation, and active it on
27+
* continuation run.
28+
*
29+
* <p>The instrumentation reuses the context store from {@link Runnable} (as {@code VirtualThread}
30+
* inherits from {@link Runnable}) to store the {@link State} captured state to restore.
31+
*/
32+
@SuppressWarnings("unused")
33+
@AutoService(InstrumenterModule.class)
34+
public final class VirtualThreadInstrumentation extends InstrumenterModule.Tracing
35+
implements Instrumenter.ForBootstrap, Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
36+
37+
public VirtualThreadInstrumentation() {
38+
super("java_lang", "virtualthread");
39+
}
40+
41+
@Override
42+
public String instrumentedType() {
43+
return VIRTUAL_THREAD_CLASS_NAME;
44+
}
45+
46+
@Override
47+
public boolean isEnabled() {
48+
// return false;
49+
return JavaVirtualMachine.isJavaVersionAtLeast(19) && super.isEnabled();
50+
}
51+
52+
@Override
53+
public Map<String, String> contextStore() {
54+
Map<String, String> contextStore = new HashMap<>();
55+
contextStore.put(Runnable.class.getName(), State.class.getName());
56+
contextStore.put(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
57+
// contextStore.put(Integer.class.getName(), AgentScope.class.getName());
58+
return contextStore;
59+
// return singletonMap(Runnable.class.getName(), State.class.getName());
60+
}
61+
62+
@Override
63+
public void methodAdvice(MethodTransformer transformer) {
64+
transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
65+
// transformer.applyAdvice(isMethod().and(named("runContinuation")), getClass().getName() +
66+
// "$Run");
67+
transformer.applyAdvice(isMethod().and(named("mount")), getClass().getName() + "$Activate");
68+
transformer.applyAdvice(isMethod().and(named("unmount")), getClass().getName() + "$Close");
69+
}
70+
71+
public static final class Construct {
72+
@OnMethodExit(suppress = Throwable.class)
73+
public static void captureScope(@Advice.This Object virtualThread) {
74+
capture(InstrumentationContext.get(Runnable.class, State.class), (Runnable) virtualThread);
75+
}
76+
}
77+
78+
/*
79+
Decorating runContinuation is too early as thread haven't change and context will be restored on
80+
platform thread rather than on the virtual thread.
81+
*/
82+
//
83+
// public static final class Run {
84+
// @OnMethodEnter(suppress = Throwable.class)
85+
// public static AgentScope activate(@Advice.This Object virtualThread) {
86+
// return startTaskScope(InstrumentationContext.get(Runnable.class, State.class), (Runnable)
87+
// virtualThread);
88+
// }
89+
//
90+
// @OnMethodExit(onThrowable = Throwable.class)
91+
// public static void close(@Advice.Enter AgentScope scope) {
92+
// endTaskScope(scope);
93+
// }
94+
95+
public static final class Activate {
96+
@OnMethodExit(suppress = Throwable.class)
97+
public static void activate(@Advice.This Object virtualThread) {
98+
AgentScope agentScope =
99+
startTaskScope(
100+
InstrumentationContext.get(Runnable.class, State.class), (Runnable) virtualThread);
101+
102+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME)
103+
.put(virtualThread, agentScope);
104+
105+
// ContextStore<Integer, AgentScope> store = InstrumentationContext.get(Integer.class,
106+
// AgentScope.class);
107+
// store.put(virtualThread.hashCode(), agentScope);
108+
}
109+
}
110+
111+
public static final class Close {
112+
@OnMethodEnter(suppress = Throwable.class)
113+
public static void close(@Advice.This Object virtualThread) {
114+
Object agentScope =
115+
InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME)
116+
.get(virtualThread);
117+
118+
// ContextStore<Integer, AgentScope> store = InstrumentationContext.get(Integer.class,
119+
// AgentScope.class);
120+
// Object agentScope = store.get(virtualThread.hashCode());
121+
if (agentScope instanceof AgentScope) {
122+
endTaskScope((AgentScope) agentScope);
123+
}
124+
}
125+
}
126+
}

‎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)