Muzzle match only once in each class loader (#4543)

* Muzzle match only once in each class loader * Move muzzle matcher caching from ReferenceMatcher to InstrumentationModuleInstaller * Update comment as caching was moved to a different method * Fix comment
open-telemetry · Nov 1, 2021 · 8dce10f · 8dce10f
1 parent 8a33514
commit 8dce10f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/...va/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java b/...va/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java
@@ -10,6 +10,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
+import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.tooling.HelperInjector;
@@ -124,6 +125,7 @@ AgentBuilder install(
   private static class MuzzleMatcher implements AgentBuilder.RawMatcher {
     private final InstrumentationModule instrumentationModule;
     private final AtomicBoolean initialized = new AtomicBoolean(false);
+    private final Cache<ClassLoader, Boolean> matchCache = Cache.builder().setWeakKeys().build();
     private volatile ReferenceMatcher referenceMatcher;
 
     private MuzzleMatcher(InstrumentationModule instrumentationModule) {
@@ -137,10 +139,14 @@ public boolean matches(
         JavaModule module,
         Class<?> classBeingRedefined,
         ProtectionDomain protectionDomain) {
-      ReferenceMatcher muzzle = getReferenceMatcher();
       if (classLoader == BOOTSTRAP_LOADER) {
         classLoader = Utils.getBootstrapProxy();
       }
+      return matchCache.computeIfAbsent(classLoader, this::doesMatch);
+    }
+
+    private boolean doesMatch(ClassLoader classLoader) {
+      ReferenceMatcher muzzle = getReferenceMatcher();
       boolean isMatch = muzzle.matches(classLoader);
 
       if (!isMatch) {
@@ -168,10 +174,8 @@ public boolean matches(
       return isMatch;
     }
 
-    // ReferenceMatcher internally caches the muzzle check results per classloader, that's why we
-    // keep its instance in a field
-    // it is lazily created to avoid unnecessarily loading the muzzle references from the module
-    // during the agent setup
+    // ReferenceMatcher is lazily created to avoid unnecessarily loading the muzzle references from
+    // the module during the agent setup
     private ReferenceMatcher getReferenceMatcher() {
       if (initialized.compareAndSet(false, true)) {
         referenceMatcher = ReferenceMatcher.of(instrumentationModule);

diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ReferenceMatcher.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/ReferenceMatcher.java
@@ -8,7 +8,6 @@
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
-import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef;
 import io.opentelemetry.javaagent.tooling.muzzle.references.FieldRef;
@@ -31,7 +30,6 @@
 /** Matches a set of references against a classloader. */
 public final class ReferenceMatcher {
 
-  private final Cache<ClassLoader, Boolean> mismatchCache = Cache.builder().setWeakKeys().build();
   private final Map<String, ClassRef> references;
   private final Set<String> helperClassNames;
   private final HelperClassPredicate helperClassPredicate;
@@ -53,18 +51,14 @@ public static ReferenceMatcher of(InstrumentationModule instrumentationModule) {
   }
 
   /**
-   * Matcher used by ByteBuddy. Fails fast and only caches empty results, or complete results
+   * Matcher used by ByteBuddy. Caller is expected to cache the result if this method is called
+   * multiple times for given class loader.
    *
-   * @param userClassLoader Classloader to validate against (cannot be {@code null}, must pass
-   *     "bootstrap proxy" instead of bootstrap class loader)
+   * @param loader Classloader to validate against (cannot be {@code null}, must pass "bootstrap
+   *     proxy" instead of bootstrap class loader)
    * @return true if all references match the classpath of loader
    */
-  public boolean matches(ClassLoader userClassLoader) {
-    return mismatchCache.computeIfAbsent(userClassLoader, this::doesMatch);
-  }
-
-  // loader cannot be null, must pass "bootstrap proxy" instead of bootstrap class loader
-  private boolean doesMatch(ClassLoader loader) {
+  public boolean matches(ClassLoader loader) {
     TypePool typePool = createTypePool(loader);
     for (ClassRef reference : references.values()) {
       if (!checkMatch(reference, typePool, loader).isEmpty()) {
@@ -75,7 +69,7 @@ private boolean doesMatch(ClassLoader loader) {
   }
 
   /**
-   * Loads the full list of mismatches. Used in debug contexts only
+   * Loads the full list of mismatches. Used in debug contexts only.
    *
    * @param loader Classloader to validate against (cannot be {@code null}, must pass "bootstrap *
    *     proxy" instead of bootstrap class loader)