Add DDCachingPoolStrategy to reuse type descriptions
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -21,6 +21,7 @@ import net.bytebuddy.utility.JavaModule;
@Slf4j
public class AgentInstaller {
public static final DDLocationStrategy LOCATION_STRATEGY = new DDLocationStrategy();
+ public static final AgentBuilder.PoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy();
private static volatile Instrumentation INSTRUMENTATION;
public static Instrumentation getInstrumentation() {
@@ -47,6 +48,7 @@ public class AgentInstaller {
.disableClassFormatChanges()
.with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
.with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
+ .with(POOL_STRATEGY)
.with(new LoggingListener())
.with(LOCATION_STRATEGY)
.ignore(any(), skipClassLoader())

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java
+package datadog.trace.agent.tooling;
+
+import static net.bytebuddy.agent.builder.AgentBuilder.*;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
+
+public class DDCachingPoolStrategy implements PoolStrategy {
+ private static final Map<ClassLoader, TypePool> typePoolCache =
+ Collections.synchronizedMap(new WeakHashMap<ClassLoader, TypePool>());
+
+ @Override
+ public TypePool typePool(ClassFileLocator classFileLocator, ClassLoader classLoader) {
+ final ClassLoader key = null == classLoader ? Utils.getBootstrapProxy() : classLoader;
+ TypePool cachedPool = typePoolCache.get(key);
+ if (null == cachedPool) {
+ synchronized (key) {
+ cachedPool = typePoolCache.get(key);
+ if (null == cachedPool) {
+ cachedPool = Default.FAST.typePool(classFileLocator, classLoader);
+ typePoolCache.put(key, cachedPool);
+ }
+ }
+ }
+ return cachedPool;
+ }
+}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDLocationStrategy.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDLocationStrategy.java
@@ -20,7 +20,7 @@ public class DDLocationStrategy implements AgentBuilder.LocationStrategy {
public ClassFileLocator classFileLocator(ClassLoader classLoader, final JavaModule javaModule) {
final List<ClassFileLocator> locators = new ArrayList<>();
while (classLoader != null) {
- locators.add(ClassFileLocator.ForClassLoader.of(classLoader));
+ locators.add(ClassFileLocator.ForClassLoader.WeaklyReferenced.of(classLoader));
classLoader = classLoader.getParent();
}
locators.add(ClassFileLocator.ForClassLoader.of(Utils.getBootstrapProxy()));

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java
@@ -82,7 +82,8 @@ public class ReferenceMatcher {
*/
public static List<Reference.Mismatch> checkMatch(Reference reference, ClassLoader loader) {
final TypePool typePool =
- TypePool.Default.of(AgentInstaller.LOCATION_STRATEGY.classFileLocator(loader));
+ AgentInstaller.POOL_STRATEGY.typePool(
+ AgentInstaller.LOCATION_STRATEGY.classFileLocator(loader), loader);
final List<Mismatch> mismatches = new ArrayList<>(0);
try {
final TypePool.Resolution resolution =

MODIFY
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceMatcherTest.groovy
dd-java-agent/testing/src/test/groovy/muzzle/ReferenceMatcherTest.groovy
@@ -45,6 +45,40 @@ class ReferenceMatcherTest extends AgentTestRunner {
MuzzleWeakReferenceTest.classLoaderRefIsGarbageCollected()
}
+ private static class CountingClassLoader extends URLClassLoader{
+ int count = 0
+
+ CountingClassLoader(URL[] urls, ClassLoader parent) {
+ super(urls, (ClassLoader) parent)
+ }
+
+ @Override
+ URL getResource(String name) {
+ count++
+ return super.getResource(name)
+ }
+ }
+
+ def "muzzle type pool caches"() {
+ setup:
+ ClassLoader cl = new CountingClassLoader(
+ [TestUtils.createJarWithClasses(MethodBodyAdvice.A,
+ MethodBodyAdvice.B,
+ MethodBodyAdvice.SomeInterface,
+ MethodBodyAdvice.SomeImplementation)] as URL[],
+ (ClassLoader) null)
+ Reference[] refs = ReferenceCreator.createReferencesFrom(MethodBodyAdvice.getName(), this.getClass().getClassLoader()).values().toArray(new Reference[0])
+ ReferenceMatcher refMatcher1 = new ReferenceMatcher(refs)
+ ReferenceMatcher refMatcher2 = new ReferenceMatcher(refs)
+ assert getMismatchClassSet(refMatcher1.getMismatchedReferenceSources(cl)) == new HashSet<>()
+ int countAfterFirstMatch = cl.count
+ // the second matcher should be able to used cached type descriptions from the first
+ assert getMismatchClassSet(refMatcher2.getMismatchedReferenceSources(cl)) == new HashSet<>()
+
+ expect:
+ cl.count == countAfterFirstMatch
+ }
+
def "matching ref #referenceName #referenceFlags against #classToCheck produces #expectedMismatches"() {
setup:
Reference.Builder builder = new Reference.Builder(referenceName)

Add DDCachingPoolStrategy to reuse type descriptions