Skip to content

Commit 7420cd1

Browse files
PerfectSlayerPerfectSlayer
committed
feat(test): Add instrumentation test for JUnit
1 parent 552ae2b commit 7420cd1

File tree

2 files changed

+217
-0
lines changed

2 files changed

+217
-0
lines changed

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

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

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
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 class InstrumentationTest {
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 ClassFileTransformerLister 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(InstrumenterModule.class, InstrumentationTest.class.getClassLoader())
70+
.iterator()
71+
.hasNext(),
72+
"No instrumentation found");
73+
this.transformerLister = new ClassFileTransformerLister();
74+
this.activeTransformer =
75+
AgentInstaller.installBytebuddyAgent(
76+
INSTRUMENTATION, true, AgentInstaller.getEnabledSystems(), this.transformerLister);
77+
}
78+
79+
protected String idGenerationStrategyName() {
80+
return "SEQUENTIAL";
81+
}
82+
83+
private boolean useStrictTraceWrites() {
84+
return true;
85+
}
86+
87+
@AfterEach
88+
public void tearDown() {
89+
this.tracer.close();
90+
this.writer.close();
91+
if (this.activeTransformer != null) {
92+
INSTRUMENTATION.removeTransformer(this.activeTransformer);
93+
this.activeTransformer = null;
94+
}
95+
96+
// All cleanups should happen before these assertions.
97+
// If not, a failing assertion may prevent cleanup
98+
this.transformerLister.verify();
99+
this.transformerLister = null;
100+
}
101+
102+
protected void blockUntilChildSpansFinished(final int numberOfSpans) {
103+
blockUntilChildSpansFinished(this.tracer.activeSpan(), numberOfSpans);
104+
}
105+
106+
static void blockUntilChildSpansFinished(AgentSpan span, int numberOfSpans) {
107+
if (span instanceof DDSpan) {
108+
TraceCollector traceCollector = ((DDSpan) span).context().getTraceCollector();
109+
if (!(traceCollector instanceof PendingTrace)) {
110+
throw new IllegalStateException(
111+
"Expected $PendingTrace.name trace collector, got $traceCollector.class.name");
112+
}
113+
114+
PendingTrace pendingTrace = (PendingTrace) traceCollector;
115+
long deadline = System.currentTimeMillis() + TIMEOUT_MILLIS;
116+
117+
while (pendingTrace.size() < numberOfSpans) {
118+
if (System.currentTimeMillis() > deadline) {
119+
throw new RuntimeException(
120+
new TimeoutException(
121+
"Timed out waiting for child spans. Received: " + pendingTrace.size()));
122+
}
123+
try {
124+
Thread.sleep(10);
125+
} catch (InterruptedException e) {
126+
Thread.currentThread().interrupt();
127+
}
128+
}
129+
}
130+
}
131+
}

0 commit comments

Comments
 (0)