Reduce memory usage for ClassLoaderHasClassesNamedMatcher (#7866)

See
#7698
This is an attempt to reduce memory usage for
`ClassLoaderHasClassesNamedMatcher`. Instead of having each matcher keep
a `Map<ClassLoader, Boolean>` we can have one `Map<ClassLoader, BitSet>`
where each matcher uses one bit in the `BitSet`. Alternatively
`Map<ClassLoader, Set<ClassLoaderHasClassesNamedMatcher>>` where set
contains matchers that match for given class loader would also work well
because these matchers usually don't match so we can expect to have only
a few elements in the set.

Author: laurit

gradle-plugins/src/main/kotlin/io/opentelemetry/javaagent/muzzle/matcher/MuzzleGradlePluginUtil.kt

@@ -57,9 +57,6 @@ class MuzzleGradlePluginUtil {
 
       val matcherClass = agentClassLoader.loadClass("io.opentelemetry.javaagent.tooling.muzzle.ClassLoaderMatcher")
 
-      // We cannot reference Mismatch class directly here, because we are loaded from a different
-      // class loader.
-
       // We cannot reference Mismatch class directly here, because we are loaded from a different
       // class loader.
       val allMismatches = matcherClass

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/internal/ClassLoaderMatcherCacheHolder.java

@@ -22,19 +22,19 @@
 public final class ClassLoaderMatcherCacheHolder {
 
   @GuardedBy("allCaches")
-  private static final List<Cache<ClassLoader, Boolean>> allCaches = new ArrayList<>();
+  private static final List<Cache<ClassLoader, ?>> allCaches = new ArrayList<>();
 
   private ClassLoaderMatcherCacheHolder() {}
 
-  public static void addCache(Cache<ClassLoader, Boolean> cache) {
+  public static void addCache(Cache<ClassLoader, ?> cache) {
     synchronized (allCaches) {
       allCaches.add(cache);
     }
   }
 
   public static void invalidateAllCachesForClassLoader(ClassLoader loader) {
     synchronized (allCaches) {
-      for (Cache<ClassLoader, Boolean> cache : allCaches) {
+      for (Cache<ClassLoader, ?> cache : allCaches) {
         cache.remove(loader);
       }
     }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java

@@ -8,20 +8,29 @@
 import io.opentelemetry.instrumentation.api.internal.cache.Cache;
 import io.opentelemetry.javaagent.bootstrap.internal.ClassLoaderMatcherCacheHolder;
 import io.opentelemetry.javaagent.bootstrap.internal.InClassLoaderMatcher;
+import java.util.BitSet;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicInteger;
 import net.bytebuddy.matcher.ElementMatcher;
 
 class ClassLoaderHasClassesNamedMatcher extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
-
-  private final Cache<ClassLoader, Boolean> cache = Cache.weak();
+  // caching is disabled for build time muzzle checks
+  // this field is set via reflection from ClassLoaderMatcher
+  static boolean useCache = true;
+  private static final AtomicInteger counter = new AtomicInteger();
 
   private final String[] resources;
+  private final int index = counter.getAndIncrement();
 
   ClassLoaderHasClassesNamedMatcher(String... classNames) {
     resources = classNames;
     for (int i = 0; i < resources.length; i++) {
       resources[i] = resources[i].replace(".", "/") + ".class";
     }
-    ClassLoaderMatcherCacheHolder.addCache(cache);
+    if (useCache) {
+      Manager.INSTANCE.add(this);
+    }
   }
 
   @Override
@@ -30,10 +39,14 @@ public boolean matches(ClassLoader cl) {
       // Can't match the bootstrap class loader.
       return false;
     }
-    return cache.computeIfAbsent(cl, this::hasResources);
+    if (useCache) {
+      return Manager.INSTANCE.match(this, cl);
+    } else {
+      return hasResources(cl, resources);
+    }
   }
 
-  private boolean hasResources(ClassLoader cl) {
+  private static boolean hasResources(ClassLoader cl, String... resources) {
     boolean priorValue = InClassLoaderMatcher.getAndSet(true);
     try {
       for (String resource : resources) {
@@ -46,4 +59,41 @@ private boolean hasResources(ClassLoader cl) {
     }
     return true;
   }
+
+  private static class Manager {
+    private static final BitSet EMPTY = new BitSet(0);
+    static final Manager INSTANCE = new Manager();
+    private final List<ClassLoaderHasClassesNamedMatcher> matchers = new CopyOnWriteArrayList<>();
+    private final Cache<ClassLoader, BitSet> enabled = Cache.weak();
+    private volatile boolean matchCalled = false;
+
+    Manager() {
+      ClassLoaderMatcherCacheHolder.addCache(enabled);
+    }
+
+    void add(ClassLoaderHasClassesNamedMatcher matcher) {
+      if (matchCalled) {
+        throw new IllegalStateException("All matchers should be create before match is called");
+      }
+      matchers.add(matcher);
+    }
+
+    boolean match(ClassLoaderHasClassesNamedMatcher matcher, ClassLoader cl) {
+      matchCalled = true;
+      BitSet set = enabled.get(cl);
+      if (set == null) {
+        set = new BitSet(counter.get());
+        for (ClassLoaderHasClassesNamedMatcher m : matchers) {
+          if (hasResources(cl, m.resources)) {
+            set.set(m.index);
+          }
+        }
+        enabled.put(cl, set.isEmpty() ? EMPTY : set);
+        enabled.put(cl, set);
+      } else if (set.isEmpty()) {
+        return false;
+      }
+      return set.get(matcher.index);
+    }
+  }
 }

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ClassLoaderMatcher.java

@@ -7,6 +7,7 @@
 
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.HelperInjector;
+import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -32,6 +33,8 @@ public class ClassLoaderMatcher {
    */
   public static Map<String, List<Mismatch>> matchesAll(
       ClassLoader classLoader, boolean injectHelpers, Set<String> excludedInstrumentationNames) {
+    disableMatcherCache();
+
     Map<String, List<Mismatch>> result = new HashMap<>();
     ServiceLoader.load(InstrumentationModule.class, ClassLoaderMatcher.class.getClassLoader())
         .forEach(
@@ -101,5 +104,19 @@ private static List<Mismatch> checkHelperInjection(
     return mismatches;
   }
 
+  private static void disableMatcherCache() {
+    try {
+      Class<?> matcherClass =
+          Class.forName(
+              "io.opentelemetry.javaagent.extension.matcher.ClassLoaderHasClassesNamedMatcher");
+      Field field = matcherClass.getDeclaredField("useCache");
+      field.setAccessible(true);
+      field.setBoolean(null, false);
+    } catch (Exception exception) {
+      throw new IllegalStateException(
+          "Failed to disable cache for ClassLoaderHasClassesNamedMatcher", exception);
+    }
+  }
+
   private ClassLoaderMatcher() {}
 }