Skip to content

Commit 1f73c23

Browse files
mccullsmcculls
authored
Simplify Muzzle side-classes to make them easier to unload and remove the need to modify the Instrumenter class (DataDog#3854)
1 parent 4568db6 commit 1f73c23

10 files changed

Lines changed: 390 additions & 528 deletions

File tree

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

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -162,11 +162,8 @@ public final void instrument(TransformerBuilder transformerBuilder) {
162162
/** Matches classes for which instrumentation is not muzzled. */
163163
public final boolean muzzleMatches(
164164
final ClassLoader classLoader, final Class<?> classBeingRedefined) {
165-
/* Optimization: calling getInstrumentationMuzzle() inside this method
166-
* prevents unnecessary loading of muzzle references during agentBuilder
167-
* setup.
168-
*/
169-
final ReferenceMatcher muzzle = getInstrumentationMuzzle();
165+
// Optimization: we delay calling getInstrumentationMuzzle() until we need the references
166+
ReferenceMatcher muzzle = getInstrumentationMuzzle();
170167
if (null != muzzle) {
171168
final boolean isMatch = muzzle.matches(classLoader);
172169
if (!isMatch) {
@@ -176,13 +173,13 @@ public final boolean muzzleMatches(
176173
log.debug(
177174
"Muzzled - instrumentation.names=[{}] instrumentation.class={} instrumentation.target.classloader={}",
178175
Strings.join(",", instrumentationNames),
179-
Instrumenter.Default.this.getClass().getName(),
176+
getClass().getName(),
180177
classLoader);
181178
for (final Reference.Mismatch mismatch : mismatches) {
182179
log.debug(
183180
"Muzzled mismatch - instrumentation.names=[{}] instrumentation.class={} instrumentation.target.classloader={} muzzle.mismatch=\"{}\"",
184181
Strings.join(",", instrumentationNames),
185-
Instrumenter.Default.this.getClass().getName(),
182+
getClass().getName(),
186183
classLoader,
187184
mismatch);
188185
}
@@ -192,7 +189,7 @@ public final boolean muzzleMatches(
192189
log.debug(
193190
"Instrumentation applied - instrumentation.names=[{}] instrumentation.class={} instrumentation.target.classloader={} instrumentation.target.class={}",
194191
Strings.join(",", instrumentationNames),
195-
Instrumenter.Default.this.getClass().getName(),
192+
getClass().getName(),
196193
classLoader,
197194
classBeingRedefined == null ? "null" : classBeingRedefined.getName());
198195
}
@@ -202,21 +199,33 @@ public final boolean muzzleMatches(
202199
return true;
203200
}
204201

205-
/**
206-
* This method is implemented dynamically by compile-time bytecode transformations.
207-
*
208-
* <p>{@see datadog.trace.agent.tooling.muzzle.MuzzleGradlePlugin}
209-
*/
210-
protected ReferenceMatcher getInstrumentationMuzzle() {
211-
return null;
202+
public final ReferenceMatcher getInstrumentationMuzzle() {
203+
String muzzleClassName = getClass().getName() + "$Muzzle";
204+
try {
205+
// Muzzle class contains static references captured at build-time
206+
// see datadog.trace.agent.tooling.muzzle.MuzzleGenerator
207+
ReferenceMatcher muzzle =
208+
(ReferenceMatcher)
209+
getClass()
210+
.getClassLoader()
211+
.loadClass(muzzleClassName)
212+
.getConstructor()
213+
.newInstance();
214+
// mix in any additional references captured at runtime
215+
muzzle.withReferenceProvider(runtimeMuzzleReferences());
216+
return muzzle;
217+
} catch (Throwable e) {
218+
log.warn("Failed to load - muzzle.class={}", muzzleClassName, e);
219+
return null;
220+
}
212221
}
213222

214223
/** @return Class names of helpers to inject into the user's classloader */
215224
public String[] helperClassNames() {
216225
return new String[0];
217226
}
218227

219-
/* Classes that the muzzle plugin assumes will be injected */
228+
/** Classes that the muzzle plugin assumes will be injected */
220229
public String[] muzzleIgnoredClassNames() {
221230
return helperClassNames();
222231
}

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

Lines changed: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,281 @@
1+
package datadog.trace.agent.tooling.muzzle;
2+
3+
import datadog.trace.agent.tooling.Instrumenter;
4+
import java.io.File;
5+
import java.io.IOException;
6+
import java.nio.file.Files;
7+
import java.util.ArrayList;
8+
import java.util.Arrays;
9+
import java.util.Collections;
10+
import java.util.HashSet;
11+
import java.util.LinkedHashMap;
12+
import java.util.List;
13+
import java.util.Map;
14+
import java.util.Set;
15+
import net.bytebuddy.asm.AsmVisitorWrapper;
16+
import net.bytebuddy.description.field.FieldDescription;
17+
import net.bytebuddy.description.field.FieldList;
18+
import net.bytebuddy.description.method.MethodDescription;
19+
import net.bytebuddy.description.method.MethodList;
20+
import net.bytebuddy.description.type.TypeDescription;
21+
import net.bytebuddy.implementation.Implementation;
22+
import net.bytebuddy.jar.asm.ClassVisitor;
23+
import net.bytebuddy.jar.asm.ClassWriter;
24+
import net.bytebuddy.jar.asm.MethodVisitor;
25+
import net.bytebuddy.jar.asm.Opcodes;
26+
import net.bytebuddy.jar.asm.Type;
27+
import net.bytebuddy.matcher.ElementMatcher;
28+
import net.bytebuddy.pool.TypePool;
29+
30+
/** Generates a 'Muzzle' side-class for each {@link Instrumenter}. */
31+
public class MuzzleGenerator implements AsmVisitorWrapper {
32+
private final File targetDir;
33+
34+
public MuzzleGenerator(File targetDir) {
35+
this.targetDir = targetDir;
36+
}
37+
38+
@Override
39+
public int mergeWriter(int flags) {
40+
return flags | ClassWriter.COMPUTE_MAXS;
41+
}
42+
43+
@Override
44+
public int mergeReader(int flags) {
45+
return flags;
46+
}
47+
48+
@Override
49+
public ClassVisitor wrap(
50+
final TypeDescription instrumentedType,
51+
final ClassVisitor classVisitor,
52+
final Implementation.Context implementationContext,
53+
final TypePool typePool,
54+
final FieldList<FieldDescription.InDefinedShape> fields,
55+
final MethodList<?> methods,
56+
final int writerFlags,
57+
final int readerFlags) {
58+
59+
Instrumenter.Default instrumenter;
60+
try {
61+
instrumenter =
62+
(Instrumenter.Default)
63+
MuzzleGenerator.class
64+
.getClassLoader()
65+
.loadClass(instrumentedType.getName())
66+
.getConstructor()
67+
.newInstance();
68+
} catch (ReflectiveOperationException e) {
69+
throw new RuntimeException(e);
70+
}
71+
72+
File muzzleClass = new File(targetDir, instrumentedType.getInternalName() + "$Muzzle.class");
73+
try {
74+
muzzleClass.getParentFile().mkdirs();
75+
Files.write(muzzleClass.toPath(), generateMuzzleClass(instrumenter));
76+
} catch (IOException e) {
77+
throw new RuntimeException(e);
78+
}
79+
return classVisitor;
80+
}
81+
82+
private static Reference[] generateReferences(Instrumenter.Default instrumenter) {
83+
// track sources we've generated references from to avoid recursion
84+
final Set<String> referenceSources = new HashSet<>();
85+
final Map<String, Reference> references = new LinkedHashMap<>();
86+
final Set<String> adviceClasses = new HashSet<>();
87+
instrumenter.adviceTransformations(
88+
new Instrumenter.AdviceTransformation() {
89+
@Override
90+
public void applyAdvice(ElementMatcher<? super MethodDescription> matcher, String name) {
91+
adviceClasses.add(name);
92+
}
93+
});
94+
for (String adviceClass : adviceClasses) {
95+
if (referenceSources.add(adviceClass)) {
96+
for (Map.Entry<String, Reference> entry :
97+
ReferenceCreator.createReferencesFrom(
98+
adviceClass, ReferenceMatcher.class.getClassLoader())
99+
.entrySet()) {
100+
Reference toMerge = references.get(entry.getKey());
101+
if (null == toMerge) {
102+
references.put(entry.getKey(), entry.getValue());
103+
} else {
104+
references.put(entry.getKey(), toMerge.merge(entry.getValue()));
105+
}
106+
}
107+
}
108+
}
109+
return references.values().toArray(new Reference[0]);
110+
}
111+
112+
/** This code is generated in a separate side-class. */
113+
private static byte[] generateMuzzleClass(Instrumenter.Default instrumenter) {
114+
115+
Set<String> ignoredClassNames =
116+
new HashSet<>(Arrays.asList(instrumenter.muzzleIgnoredClassNames()));
117+
118+
List<Reference> references = new ArrayList<>();
119+
for (Reference reference : generateReferences(instrumenter)) {
120+
// ignore helper classes, they will be injected by the instrumentation's HelperInjector.
121+
if (!ignoredClassNames.contains(reference.className)) {
122+
references.add(reference);
123+
}
124+
}
125+
Reference[] additionalReferences = instrumenter.additionalMuzzleReferences();
126+
if (null != additionalReferences) {
127+
Collections.addAll(references, additionalReferences);
128+
}
129+
130+
ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
131+
cw.visit(
132+
Opcodes.V1_7,
133+
Opcodes.ACC_PUBLIC | Opcodes.ACC_FINAL,
134+
Type.getInternalName(instrumenter.getClass()) + "$Muzzle",
135+
null,
136+
"datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
137+
null);
138+
139+
MethodVisitor mv = cw.visitMethod(Opcodes.ACC_PUBLIC, "<init>", "()V", null, null);
140+
141+
mv.visitCode();
142+
143+
mv.visitIntInsn(Opcodes.ALOAD, 0);
144+
mv.visitLdcInsn(references.size());
145+
mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference");
146+
147+
int i = 0;
148+
for (Reference reference : references) {
149+
mv.visitInsn(Opcodes.DUP);
150+
mv.visitLdcInsn(i++);
151+
writeReference(mv, reference);
152+
mv.visitInsn(Opcodes.AASTORE);
153+
}
154+
155+
mv.visitMethodInsn(
156+
Opcodes.INVOKESPECIAL,
157+
"datadog/trace/agent/tooling/muzzle/ReferenceMatcher",
158+
"<init>",
159+
"([Ldatadog/trace/agent/tooling/muzzle/Reference;)V",
160+
false);
161+
162+
mv.visitInsn(Opcodes.RETURN);
163+
164+
mv.visitMaxs(0, 0);
165+
mv.visitEnd();
166+
167+
return cw.toByteArray();
168+
}
169+
170+
private static void writeReference(MethodVisitor mv, Reference reference) {
171+
if (reference instanceof OrReference) {
172+
mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/OrReference");
173+
mv.visitInsn(Opcodes.DUP);
174+
}
175+
176+
mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference");
177+
mv.visitInsn(Opcodes.DUP);
178+
179+
writeStrings(mv, reference.sources);
180+
mv.visitLdcInsn(reference.flags);
181+
mv.visitLdcInsn(reference.className);
182+
if (null != reference.superName) {
183+
mv.visitLdcInsn(reference.superName);
184+
} else {
185+
mv.visitInsn(Opcodes.ACONST_NULL);
186+
}
187+
writeStrings(mv, reference.interfaces);
188+
writeFields(mv, reference.fields);
189+
writeMethods(mv, reference.methods);
190+
191+
mv.visitMethodInsn(
192+
Opcodes.INVOKESPECIAL,
193+
"datadog/trace/agent/tooling/muzzle/Reference",
194+
"<init>",
195+
"([Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;[Ljava/lang/String;"
196+
+ "[Ldatadog/trace/agent/tooling/muzzle/Reference$Field;"
197+
+ "[Ldatadog/trace/agent/tooling/muzzle/Reference$Method;)V",
198+
false);
199+
200+
if (reference instanceof OrReference) {
201+
Reference[] ors = ((OrReference) reference).ors;
202+
203+
mv.visitLdcInsn(ors.length);
204+
mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference");
205+
206+
int i = 0;
207+
for (Reference or : ors) {
208+
mv.visitInsn(Opcodes.DUP);
209+
mv.visitLdcInsn(i++);
210+
writeReference(mv, or);
211+
mv.visitInsn(Opcodes.AASTORE);
212+
}
213+
214+
mv.visitMethodInsn(
215+
Opcodes.INVOKESPECIAL,
216+
"datadog/trace/agent/tooling/muzzle/OrReference",
217+
"<init>",
218+
"(Ldatadog/trace/agent/tooling/muzzle/Reference;"
219+
+ "[Ldatadog/trace/agent/tooling/muzzle/Reference;)V",
220+
false);
221+
}
222+
}
223+
224+
private static void writeStrings(MethodVisitor mv, String[] strings) {
225+
mv.visitLdcInsn(strings.length);
226+
mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/String");
227+
int i = 0;
228+
for (String string : strings) {
229+
mv.visitInsn(Opcodes.DUP);
230+
mv.visitLdcInsn(i++);
231+
mv.visitLdcInsn(string);
232+
mv.visitInsn(Opcodes.AASTORE);
233+
}
234+
}
235+
236+
private static void writeFields(MethodVisitor mv, Reference.Field[] fields) {
237+
mv.visitLdcInsn(fields.length);
238+
mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference$Field");
239+
int i = 0;
240+
for (Reference.Field field : fields) {
241+
mv.visitInsn(Opcodes.DUP);
242+
mv.visitLdcInsn(i++);
243+
mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference$Field");
244+
mv.visitInsn(Opcodes.DUP);
245+
writeStrings(mv, field.sources);
246+
mv.visitLdcInsn(field.flags);
247+
mv.visitLdcInsn(field.name);
248+
mv.visitLdcInsn(field.fieldType);
249+
mv.visitMethodInsn(
250+
Opcodes.INVOKESPECIAL,
251+
"datadog/trace/agent/tooling/muzzle/Reference$Field",
252+
"<init>",
253+
"([Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;)V",
254+
false);
255+
mv.visitInsn(Opcodes.AASTORE);
256+
}
257+
}
258+
259+
private static void writeMethods(MethodVisitor mv, Reference.Method[] methods) {
260+
mv.visitLdcInsn(methods.length);
261+
mv.visitTypeInsn(Opcodes.ANEWARRAY, "datadog/trace/agent/tooling/muzzle/Reference$Method");
262+
int i = 0;
263+
for (Reference.Method method : methods) {
264+
mv.visitInsn(Opcodes.DUP);
265+
mv.visitLdcInsn(i++);
266+
mv.visitTypeInsn(Opcodes.NEW, "datadog/trace/agent/tooling/muzzle/Reference$Method");
267+
mv.visitInsn(Opcodes.DUP);
268+
writeStrings(mv, method.sources);
269+
mv.visitLdcInsn(method.flags);
270+
mv.visitLdcInsn(method.name);
271+
mv.visitLdcInsn(method.methodType);
272+
mv.visitMethodInsn(
273+
Opcodes.INVOKESPECIAL,
274+
"datadog/trace/agent/tooling/muzzle/Reference$Method",
275+
"<init>",
276+
"([Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;)V",
277+
false);
278+
mv.visitInsn(Opcodes.AASTORE);
279+
}
280+
}
281+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public DynamicType.Builder<?> apply(
3838
final DynamicType.Builder<?> builder,
3939
final TypeDescription typeDescription,
4040
final ClassFileLocator classFileLocator) {
41-
return builder.visit(new MuzzleVisitor(targetDir));
41+
return builder.visit(new MuzzleGenerator(targetDir));
4242
}
4343

4444
@Override

0 commit comments

Comments
 (0)