From 6f4cf22741ba6ed4f23d66bf0f5da3cd3db74af5 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Tue, 21 Feb 2023 15:48:15 +0200
Subject: [PATCH] Reduce memory usage for ClassLoaderHasClassesNamedMatcher

---
 .../ClassLoaderMatcherCacheHolder.java        |  6 +--
 .../ClassLoaderHasClassesNamedMatcher.java    | 51 +++++++++++++++++--
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/internal/ClassLoaderMatcherCacheHolder.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/internal/ClassLoaderMatcherCacheHolder.java
index 1c23650bbc7f..583c40d180e8 100644
--- a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/internal/ClassLoaderMatcherCacheHolder.java
+++ b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/internal/ClassLoaderMatcherCacheHolder.java
@@ -22,11 +22,11 @@
 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);
     }
@@ -34,7 +34,7 @@ public static void addCache(Cache<ClassLoader, Boolean> cache) {
 
   public static void invalidateAllCachesForClassLoader(ClassLoader loader) {
     synchronized (allCaches) {
-      for (Cache<ClassLoader, Boolean> cache : allCaches) {
+      for (Cache<ClassLoader, ?> cache : allCaches) {
         cache.remove(loader);
       }
     }
diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java
index 368effcd3447..1cf8cd68d9ca 100644
--- a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java
+++ b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java
@@ -8,20 +8,24 @@
 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();
+  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);
+    Manager.INSTANCE.add(this);
   }
 
   @Override
@@ -30,10 +34,10 @@ public boolean matches(ClassLoader cl) {
       // Can't match the bootstrap class loader.
       return false;
     }
-    return cache.computeIfAbsent(cl, this::hasResources);
+    return Manager.INSTANCE.match(this, cl);
   }
 
-  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 +50,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);
+    }
+  }
 }