Skip to content

Commit b98b760

Browse files
mccullsmcculls
authored
Simplify HelperInjector to make it easier to switch out the underlying approach (#9643)
* Test helper consistency in MuzzleVersionScanPlugin without requiring HelperInjector This helps decouple internals and provides the same feedback with simpler code * Remove duplicate test: we already test unloading of an injected class-loader in this class * Replace use of ByteBuddy's ClassInjector when redefining modules to add reads * Remove bogus check from MuzzleVersionScanPluginTest: if an instrumentation is muzzled, its helpers won't be injected
1 parent 6c26663 commit b98b760

File tree

5 files changed

+87
-84
lines changed

5 files changed

+87
-84
lines changed

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java‎

Lines changed: 33 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
11
package datadog.trace.agent.tooling;
22

33
import static datadog.trace.bootstrap.AgentClassLoading.INJECTING_HELPERS;
4+
import static java.util.Arrays.asList;
45

56
import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
7+
import datadog.trace.util.JDK9ModuleAccess;
68
import java.io.IOException;
79
import java.lang.ref.WeakReference;
10+
import java.lang.reflect.AnnotatedElement;
811
import java.security.CodeSource;
912
import java.security.ProtectionDomain;
10-
import java.util.Arrays;
1113
import java.util.Collections;
14+
import java.util.HashSet;
1215
import java.util.LinkedHashMap;
1316
import java.util.LinkedHashSet;
1417
import java.util.List;
@@ -20,7 +23,6 @@
2023
import net.bytebuddy.dynamic.ClassFileLocator;
2124
import net.bytebuddy.dynamic.DynamicType;
2225
import net.bytebuddy.dynamic.loading.ClassInjector;
23-
import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
2426
import net.bytebuddy.utility.JavaModule;
2527
import org.slf4j.Logger;
2628
import org.slf4j.LoggerFactory;
@@ -40,9 +42,9 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
4042
private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
4143

4244
private final Map<ClassLoader, Boolean> injectedClassLoaders =
43-
Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
45+
Collections.synchronizedMap(new WeakHashMap<>());
4446

45-
private final List<WeakReference<Object>> helperModules = new CopyOnWriteArrayList<>();
47+
private final List<WeakReference<AnnotatedElement>> helperModules = new CopyOnWriteArrayList<>();
4648

4749
/**
4850
* Construct HelperInjector.
@@ -71,7 +73,7 @@ public HelperInjector(
7173
this.requestingName = requestingName;
7274
this.adviceShader = adviceShader;
7375

74-
this.helperClassNames = new LinkedHashSet<>(Arrays.asList(helperClassNames));
76+
this.helperClassNames = new LinkedHashSet<>(asList(helperClassNames));
7577
}
7678

7779
public HelperInjector(
@@ -132,12 +134,10 @@ public DynamicType.Builder<?> transform(
132134
final Map<String, byte[]> classnameToBytes = getHelperMap();
133135
final Map<String, Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
134136

135-
// All datadog helper classes are in the unnamed module
136-
// And there's exactly one unnamed module per classloader
137-
// Use the module of the first class for convenience
137+
// all datadog helper classes are in the unnamed module
138+
// and there's exactly one unnamed module per classloader
138139
if (JavaModule.isSupported()) {
139-
final JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
140-
helperModules.add(new WeakReference<>(javaModule.unwrap()));
140+
helperModules.add(new WeakReference<>(JDK9ModuleAccess.getUnnamedModule(classLoader)));
141141
}
142142

143143
// forcibly initialize any eager helpers
@@ -177,43 +177,42 @@ private Map<String, Class<?>> injectClassLoader(
177177
final ClassLoader classLoader, final Map<String, byte[]> classnameToBytes) {
178178
INJECTING_HELPERS.begin();
179179
try {
180-
ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
181-
return new ClassInjector.UsingReflection(classLoader, protectionDomain)
182-
.injectRaw(classnameToBytes);
180+
if (useAgentCodeSource) {
181+
ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
182+
return new ClassInjector.UsingReflection(classLoader, protectionDomain)
183+
.injectRaw(classnameToBytes);
184+
} else {
185+
return new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
186+
}
183187
} finally {
184188
INJECTING_HELPERS.end();
185189
}
186190
}
187191

188192
private ProtectionDomain createProtectionDomain(final ClassLoader classLoader) {
189-
if (useAgentCodeSource) {
190-
CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
191-
return new ProtectionDomain(codeSource, null, classLoader, null);
192-
} else {
193-
return ClassLoadingStrategy.NO_PROTECTION_DOMAIN;
194-
}
193+
CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
194+
return new ProtectionDomain(codeSource, null, classLoader, null);
195195
}
196196

197197
private void ensureModuleCanReadHelperModules(final JavaModule target) {
198198
if (JavaModule.isSupported() && target != JavaModule.UNSUPPORTED && target.isNamed()) {
199-
for (final WeakReference<Object> helperModuleReference : helperModules) {
200-
final Object realModule = helperModuleReference.get();
201-
if (realModule != null) {
202-
final JavaModule helperModule = JavaModule.of(realModule);
203-
204-
if (!target.canRead(helperModule)) {
205-
log.debug("Adding module read from {} to {}", target, helperModule);
206-
ClassInjector.UsingInstrumentation.redefineModule(
207-
Utils.getInstrumentation(),
208-
target,
209-
Collections.singleton(helperModule),
210-
Collections.<String, Set<JavaModule>>emptyMap(),
211-
Collections.<String, Set<JavaModule>>emptyMap(),
212-
Collections.<Class<?>>emptySet(),
213-
Collections.<Class<?>, List<Class<?>>>emptyMap());
199+
AnnotatedElement targetModule = (AnnotatedElement) target.unwrap();
200+
Set<AnnotatedElement> extraReads = null;
201+
for (final WeakReference<AnnotatedElement> helperModuleReference : helperModules) {
202+
final AnnotatedElement helperModule = helperModuleReference.get();
203+
if (helperModule != null) {
204+
if (!JDK9ModuleAccess.canRead(targetModule, helperModule)) {
205+
if (extraReads == null) {
206+
extraReads = new HashSet<>();
207+
}
208+
extraReads.add(helperModule);
214209
}
215210
}
216211
}
212+
if (extraReads != null) {
213+
log.debug("Adding module reads from {} to {}", targetModule, extraReads);
214+
JDK9ModuleAccess.addModuleReads(Utils.getInstrumentation(), targetModule, extraReads);
215+
}
217216
}
218217
}
219218
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package datadog.trace.agent.tooling.muzzle;
22

33
import datadog.trace.agent.tooling.AdviceShader;
4-
import datadog.trace.agent.tooling.HelperInjector;
54
import datadog.trace.agent.tooling.InstrumenterModule;
65
import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
76
import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
@@ -72,16 +71,15 @@ public static void assertInstrumentationMuzzled(
7271
if (assertPass) {
7372
for (InstrumenterModule module : toBeTested) {
7473
try {
75-
// verify helper injector works
74+
// verify helper consistency
7675
final String[] helperClassNames = module.helperClassNames();
7776
if (helperClassNames.length > 0) {
78-
new HelperInjector(
79-
module.useAgentCodeSource(),
80-
MuzzleVersionScanPlugin.class.getSimpleName(),
81-
createHelperMap(module))
82-
.transform(null, null, testApplicationLoader, null, null);
77+
HelperClassLoader helperClassLoader = new HelperClassLoader(testApplicationLoader);
78+
for (Map.Entry<String, byte[]> helper : createHelperMap(module).entrySet()) {
79+
helperClassLoader.injectClass(helper.getKey(), helper.getValue());
80+
}
8381
}
84-
} catch (final Exception e) {
82+
} catch (final Throwable e) {
8583
System.err.println(
8684
"FAILED HELPER INJECTION. Are Helpers being injected in the correct order?");
8785
System.err.println(e.getMessage());
@@ -107,6 +105,18 @@ private static synchronized List<InstrumenterModule> toBeTested(
107105
return toBeTested;
108106
}
109107

108+
// Exposes ClassLoader.defineClass() to test helper consistency
109+
// without requiring java.lang.instrument.Instrumentation agent
110+
static final class HelperClassLoader extends ClassLoader {
111+
HelperClassLoader(ClassLoader parent) {
112+
super(parent);
113+
}
114+
115+
public void injectClass(String name, byte[] bytecode) {
116+
defineClass(name, bytecode, 0, bytecode.length);
117+
}
118+
}
119+
110120
private static Map<String, byte[]> createHelperMap(final InstrumenterModule module)
111121
throws IOException {
112122
String[] helperClasses = module.helperClassNames();

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy‎

Lines changed: 0 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@ package datadog.trace.agent.test
33
import datadog.trace.agent.tooling.HelperInjector
44
import datadog.trace.agent.tooling.Utils
55
import datadog.trace.test.util.DDSpecification
6-
import net.bytebuddy.description.type.TypeDescription
7-
import net.bytebuddy.dynamic.ClassFileLocator
8-
import net.bytebuddy.dynamic.loading.ClassInjector
96

107
import java.lang.ref.WeakReference
118
import java.util.concurrent.atomic.AtomicReference
@@ -16,7 +13,6 @@ import static datadog.trace.test.util.GCUtils.awaitGC
1613
class HelperInjectionTest extends DDSpecification {
1714
static final String HELPER_CLASS_NAME = 'datadog.trace.agent.test.HelperClass'
1815

19-
//@Flaky("awaitGC usage is flaky")
2016
def "helpers injected to non-delegating classloader"() {
2117
setup:
2218
HelperInjector injector = new HelperInjector(false, "test", HELPER_CLASS_NAME)
@@ -45,39 +41,4 @@ class HelperInjectionTest extends DDSpecification {
4541
then: "HelperInjector doesn't prevent it from being collected"
4642
null == ref.get()
4743
}
48-
49-
//@Flaky("awaitGC usage is flaky")
50-
def "check hard references on class injection"() {
51-
setup:
52-
53-
// Copied from HelperInjector:
54-
final ClassFileLocator locator =
55-
ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader())
56-
final byte[] classBytes = locator.locate(HELPER_CLASS_NAME).resolve()
57-
final TypeDescription typeDesc =
58-
new TypeDescription.Latent(
59-
HELPER_CLASS_NAME, 0, null, Collections.<TypeDescription.Generic> emptyList())
60-
61-
AtomicReference<URLClassLoader> emptyLoader = new AtomicReference<>(new URLClassLoader(new URL[0], (ClassLoader) null))
62-
AtomicReference<ClassInjector> injector = new AtomicReference<>(new ClassInjector.UsingReflection(emptyLoader.get()))
63-
injector.get().inject([(typeDesc): classBytes])
64-
65-
when:
66-
def injectorRef = new WeakReference(injector.get())
67-
injector.set(null)
68-
69-
awaitGC(injectorRef)
70-
71-
then:
72-
null == injectorRef.get()
73-
74-
when:
75-
def loaderRef = new WeakReference(emptyLoader.get())
76-
emptyLoader.set(null)
77-
78-
awaitGC(loaderRef)
79-
80-
then:
81-
null == loaderRef.get()
82-
}
8344
}

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,6 @@ class MuzzleVersionScanPluginTest extends DDSpecification {
7878

7979
expect:
8080
MuzzleVersionScanPlugin.assertInstrumentationMuzzled(instrumentationLoader, testApplicationLoader, true, null)
81-
!helpers.findAll {
82-
testApplicationLoader.loadClass(it.name) != null
83-
}.isEmpty()
8481

8582
where:
8683
// spotless:off

‎internal-api/internal-api-9/src/main/java/datadog/trace/util/JDK9ModuleAccess.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package datadog.trace.util;
2+
3+
import static java.util.Collections.emptyMap;
4+
import static java.util.Collections.emptySet;
5+
6+
import java.lang.instrument.Instrumentation;
7+
import java.lang.reflect.AnnotatedElement;
8+
import java.util.Set;
9+
10+
/** Use standard API to work with JPMS modules on Java9+. */
11+
@SuppressWarnings("Since15")
12+
public final class JDK9ModuleAccess {
13+
14+
/** Retrieves a class-loader's unnamed module. */
15+
public static AnnotatedElement getUnnamedModule(ClassLoader cl) {
16+
return cl.getUnnamedModule();
17+
}
18+
19+
/** Returns {@code true} if the first module can read the second module. */
20+
public static boolean canRead(AnnotatedElement module, AnnotatedElement anotherModule) {
21+
return ((java.lang.Module) module).canRead((java.lang.Module) anotherModule);
22+
}
23+
24+
/** Adds extra module reads to the given module. */
25+
@SuppressWarnings({"rawtypes", "unchecked"})
26+
public static void addModuleReads(
27+
Instrumentation inst, AnnotatedElement module, Set<AnnotatedElement> extraReads) {
28+
inst.redefineModule(
29+
(java.lang.Module) module,
30+
(Set) extraReads,
31+
emptyMap(),
32+
emptyMap(),
33+
emptySet(),
34+
emptyMap());
35+
}
36+
}

0 commit comments

Comments
 (0)