Add DDCachingPoolStrategy to reuse type descriptions

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())

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java

@@ -0,0 +1,30 @@
+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;
+  }
+}

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()));

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 =

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)