Skip to content

Commit 44c4979

Browse files
PerfectSlayerPerfectSlayer
authored
Improve virtual thread context tracking support (#10208)
* feat(test): Add instrumentation test for JUnit * feat(java-lang): Migrate tests to JUnit
1 parent 1f4c004 commit 44c4979

File tree

6 files changed

+390
-171
lines changed

6 files changed

+390
-171
lines changed

‎dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package datadog.trace.agent.test;
2+
3+
import static org.junit.jupiter.api.Assertions.assertNull;
4+
import static org.junit.jupiter.api.Assertions.assertTrue;
5+
6+
import datadog.instrument.classinject.ClassInjector;
7+
import datadog.trace.agent.tooling.AgentInstaller;
8+
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import datadog.trace.agent.tooling.TracerInstaller;
10+
import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
11+
import datadog.trace.api.Config;
12+
import datadog.trace.api.IdGenerationStrategy;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
15+
import datadog.trace.common.writer.ListWriter;
16+
import datadog.trace.core.CoreTracer;
17+
import datadog.trace.core.DDSpan;
18+
import datadog.trace.core.PendingTrace;
19+
import datadog.trace.core.TraceCollector;
20+
import java.lang.instrument.ClassFileTransformer;
21+
import java.lang.instrument.Instrumentation;
22+
import java.util.ServiceLoader;
23+
import java.util.concurrent.TimeUnit;
24+
import java.util.concurrent.TimeoutException;
25+
import net.bytebuddy.agent.ByteBuddyAgent;
26+
import org.junit.jupiter.api.AfterEach;
27+
import org.junit.jupiter.api.BeforeEach;
28+
import org.junit.jupiter.api.extension.ExtendWith;
29+
30+
@ExtendWith(TestClassShadowingExtension.class)
31+
public abstract class AbstractInstrumentationTest {
32+
static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
33+
34+
static final long TIMEOUT_MILLIS = TimeUnit.SECONDS.toMillis(20);
35+
36+
protected AgentTracer.TracerAPI tracer;
37+
38+
protected ListWriter writer;
39+
40+
protected ClassFileTransformer activeTransformer;
41+
protected ClassFileTransformerListener transformerLister;
42+
43+
@BeforeEach
44+
public void init() {
45+
// If this fails, it's likely the result of another test loading Config before it can be
46+
// injected into the bootstrap classpath.
47+
// If one test extends AgentTestRunner in a module, all tests must extend
48+
assertNull(Config.class.getClassLoader(), "Config must load on the bootstrap classpath.");
49+
50+
// Initialize test tracer
51+
this.writer = new ListWriter();
52+
// Initialize test tracer
53+
CoreTracer tracer =
54+
CoreTracer.builder()
55+
.writer(this.writer)
56+
.idGenerationStrategy(IdGenerationStrategy.fromName(idGenerationStrategyName()))
57+
.strictTraceWrites(useStrictTraceWrites())
58+
.build();
59+
TracerInstaller.forceInstallGlobalTracer(tracer);
60+
this.tracer = tracer;
61+
62+
ClassInjector.enableClassInjection(INSTRUMENTATION);
63+
64+
// if a test enables the instrumentation it verifies,
65+
// the cache needs to be recomputed taking into account that instrumentation's matchers
66+
ClassLoaderMatchers.resetState();
67+
68+
assertTrue(
69+
ServiceLoader.load(
70+
InstrumenterModule.class, AbstractInstrumentationTest.class.getClassLoader())
71+
.iterator()
72+
.hasNext(),
73+
"No instrumentation found");
74+
this.transformerLister = new ClassFileTransformerListener();
75+
this.activeTransformer =
76+
AgentInstaller.installBytebuddyAgent(
77+
INSTRUMENTATION, true, AgentInstaller.getEnabledSystems(), this.transformerLister);
78+
}
79+
80+
protected String idGenerationStrategyName() {
81+
return "SEQUENTIAL";
82+
}
83+
84+
private boolean useStrictTraceWrites() {
85+
return true;
86+
}
87+
88+
@AfterEach
89+
public void tearDown() {
90+
this.tracer.close();
91+
this.writer.close();
92+
if (this.activeTransformer != null) {
93+
INSTRUMENTATION.removeTransformer(this.activeTransformer);
94+
this.activeTransformer = null;
95+
}
96+
97+
// All cleanups should happen before these assertions.
98+
// If not, a failing assertion may prevent cleanup
99+
this.transformerLister.verify();
100+
this.transformerLister = null;
101+
}
102+
103+
protected void blockUntilChildSpansFinished(final int numberOfSpans) {
104+
blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
105+
}
106+
107+
static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
108+
if (span instanceof DDSpan) {
109+
TraceCollector traceCollector = ((DDSpan) span).context().getTraceCollector();
110+
if (!(traceCollector instanceof PendingTrace)) {
111+
throw new IllegalStateException(
112+
"Expected $PendingTrace.name trace collector, got $traceCollector.class.name");
113+
}
114+
115+
PendingTrace pendingTrace = (PendingTrace) traceCollector;
116+
long deadline = System.currentTimeMillis() + TIMEOUT_MILLIS;
117+
118+
while (pendingTrace.size() < numberOfSpans) {
119+
if (System.currentTimeMillis() > deadline) {
120+
throw new RuntimeException(
121+
new TimeoutException(
122+
"Timed out waiting for child spans. Received: " + pendingTrace.size()));
123+
}
124+
try {
125+
Thread.sleep(10);
126+
} catch (InterruptedException e) {
127+
Thread.currentThread().interrupt();
128+
}
129+
}
130+
}
131+
}
132+
}

‎dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/ClassFileTransformerListener.java‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package datadog.trace.agent.test;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertTrue;
5+
6+
import com.google.common.collect.Sets;
7+
import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores;
8+
import de.thetaphi.forbiddenapis.SuppressForbidden;
9+
import java.util.Set;
10+
import java.util.concurrent.atomic.AtomicInteger;
11+
import net.bytebuddy.agent.builder.AgentBuilder;
12+
import net.bytebuddy.description.type.TypeDescription;
13+
import net.bytebuddy.dynamic.DynamicType;
14+
import net.bytebuddy.utility.JavaModule;
15+
import net.bytebuddy.utility.nullability.MaybeNull;
16+
17+
public class ClassFileTransformerListener implements AgentBuilder.Listener {
18+
19+
final Set<String> transformedClassesNames = Sets.newConcurrentHashSet();
20+
final Set<TypeDescription> transformedClassesTypes = Sets.newConcurrentHashSet();
21+
final AtomicInteger instrumentationErrorCount = new AtomicInteger(0);
22+
23+
@Override
24+
public void onTransformation(
25+
TypeDescription typeDescription,
26+
@MaybeNull ClassLoader classLoader,
27+
@MaybeNull JavaModule module,
28+
boolean loaded,
29+
DynamicType dynamicType) {
30+
this.transformedClassesNames.add(typeDescription.getActualName());
31+
this.transformedClassesTypes.add(typeDescription);
32+
}
33+
34+
@SuppressForbidden // Allows System.out.println
35+
@Override
36+
public void onError(
37+
String typeName,
38+
ClassLoader classLoader,
39+
JavaModule module,
40+
boolean loaded,
41+
Throwable throwable) {
42+
// Incorrect* classes assert on incorrect api usage. Error expected.
43+
if (typeName.startsWith("context.FieldInjectionTestInstrumentation$Incorrect")
44+
&& throwable.getMessage().startsWith("Incorrect Context Api Usage detected.")) {
45+
return;
46+
}
47+
48+
System.out.println(
49+
"Unexpected instrumentation error when instrumenting " + typeName + " on " + classLoader);
50+
throwable.printStackTrace();
51+
instrumentationErrorCount.incrementAndGet();
52+
}
53+
54+
@Override
55+
public void onDiscovery(
56+
String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
57+
// Nothing special to do
58+
}
59+
60+
@Override
61+
public void onIgnored(
62+
TypeDescription typeDescription, ClassLoader classLoader, JavaModule module, boolean loaded) {
63+
// Nothing special to do
64+
}
65+
66+
@Override
67+
public void onComplete(
68+
String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {
69+
// Nothing special to do
70+
}
71+
72+
public void verify() {
73+
// Check instrumentation errors
74+
int errorCount = this.instrumentationErrorCount.get();
75+
assertEquals(0, errorCount, errorCount + " instrumentation errors during test");
76+
// Check effectively transformed classes that should have been ignored
77+
assertTrue(
78+
this.transformedClassesTypes.stream()
79+
.map(TypeDescription::getActualName)
80+
.noneMatch(GlobalIgnores::isAdditionallyIgnored),
81+
"Transformed classes match global libraries ignore matcher");
82+
}
83+
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@ idea {
2222
}
2323
}
2424

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)
25+
// Set test compile task to use JDK 21 to use the virtual threads API
26+
tasks.named("compileTestJava", JavaCompile) {
27+
configureCompiler(it, 21)
2828
}
2929

3030
dependencies {

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

Lines changed: 0 additions & 168 deletions
This file was deleted.

0 commit comments

Comments
 (0)