Skip to content

Commit 089bb3d

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

File tree

2 files changed

+214
-0
lines changed

2 files changed

+214
-0
lines changed

‎dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/ClassFileTransformerLister.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 ClassFileTransformerLister 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-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 abstract 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)