Merge pull request #1251 from DataDog/devinsba/fast-reference-matcher

Fail fast in the matcher, let the debug outputs use the cache

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

@@ -139,10 +139,11 @@ public interface Instrumenter {
          */
         final ReferenceMatcher muzzle = getInstrumentationMuzzle();
         if (null != muzzle) {
-          final List<Reference.Mismatch> mismatches =
+          final boolean isMatch = muzzle.matches(classLoader);
-              muzzle.getMismatchedReferenceSources(classLoader);
+          if (!isMatch) {
-          if (mismatches.size() > 0) {
             if (log.isDebugEnabled()) {
+              final List<Reference.Mismatch> mismatches =
+                  muzzle.getMismatchedReferenceSources(classLoader);
               log.debug(
                   "Instrumentation muzzled: {} -- {} on {}",
                   instrumentationNames,
@@ -159,7 +160,7 @@ public interface Instrumenter {
                 Instrumenter.Default.this.getClass().getName(),
                 classLoader);
           }
-          return mismatches.size() == 0;
+          return isMatch;
         }
         return true;
       }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java

@@ -22,9 +22,8 @@ import net.bytebuddy.pool.TypePool;
 
 /** Matches a set of references against a classloader. */
 @Slf4j
-public class ReferenceMatcher
+public final class ReferenceMatcher implements WeakMap.ValueSupplier<ClassLoader, Boolean> {
-    implements WeakMap.ValueSupplier<ClassLoader, List<Reference.Mismatch>> {
+  private final WeakMap<ClassLoader, Boolean> mismatchCache = newWeakMap();
-  private final WeakMap<ClassLoader, List<Reference.Mismatch>> mismatchCache = newWeakMap();
   private final Reference[] references;
   private final Set<String> helperClassNames;
 
@@ -42,18 +41,12 @@ public class ReferenceMatcher
   }
 
   /**
+   * Matcher used by ByteBuddy. Fails fast and only caches empty results, or complete results
+   *
    * @param loader Classloader to validate against (or null for bootstrap)
    * @return true if all references match the classpath of loader
    */
-  public boolean matches(final ClassLoader loader) {
+  public boolean matches(ClassLoader loader) {
-    return getMismatchedReferenceSources(loader).isEmpty();
-  }
-
-  /**
-   * @param loader Classloader to validate against (or null for bootstrap)
-   * @return A list of all mismatches between this ReferenceMatcher and loader's classpath.
-   */
-  public List<Reference.Mismatch> getMismatchedReferenceSources(ClassLoader loader) {
     if (loader == BOOTSTRAP_LOADER) {
       loader = Utils.getBootstrapProxy();
     }
@@ -62,8 +55,32 @@ public class ReferenceMatcher
   }
 
   @Override
-  public List<Mismatch> get(final ClassLoader loader) {
+  public Boolean get(final ClassLoader loader) {
-    final List<Mismatch> mismatches = new ArrayList<>(0);
+    for (final Reference reference : references) {
+      // Don't reference-check helper classes.
+      // They will be injected by the instrumentation's HelperInjector.
+      if (!helperClassNames.contains(reference.getClassName())) {
+        if (!checkMatch(reference, loader).isEmpty()) {
+          return false;
+        }
+      }
+    }
+
+    return true;
+  }
+
+  /**
+   * Loads the full list of mismatches. Used in debug contexts only
+   *
+   * @param loader Classloader to validate against (or null for bootstrap)
+   * @return A list of all mismatches between this ReferenceMatcher and loader's classpath.
+   */
+  public List<Reference.Mismatch> getMismatchedReferenceSources(ClassLoader loader) {
+    if (loader == BOOTSTRAP_LOADER) {
+      loader = Utils.getBootstrapProxy();
+    }
+
+    final List<Mismatch> mismatches = new ArrayList<>();
 
     for (final Reference reference : references) {
       // Don't reference-check helper classes.
@@ -82,11 +99,12 @@ public class ReferenceMatcher
    * @param loader
    * @return A list of mismatched sources. A list of size 0 means the reference matches the class.
    */
-  public static List<Reference.Mismatch> checkMatch(Reference reference, ClassLoader loader) {
+  private static List<Reference.Mismatch> checkMatch(
+      final Reference reference, final ClassLoader loader) {
     final TypePool typePool =
         AgentTooling.poolStrategy()
             .typePool(AgentTooling.locationStrategy().classFileLocator(loader), loader);
-    final List<Mismatch> mismatches = new ArrayList<>(0);
+    final List<Mismatch> mismatches = new ArrayList<>();
     try {
       final TypePool.Resolution resolution =
           typePool.describe(Utils.getClassName(reference.getClassName()));
@@ -96,7 +114,7 @@ public class ReferenceMatcher
                 reference.getSources().toArray(new Source[0]), reference.getClassName()));
       }
       return checkMatch(reference, resolution.resolve());
-    } catch (Exception e) {
+    } catch (final Exception e) {
       if (e.getMessage().startsWith("Cannot resolve type description for ")) {
         // bytebuddy throws an illegal state exception with this message if it cannot resolve types
         // TODO: handle missing type resolutions without catching bytebuddy's exceptions
@@ -112,10 +130,10 @@ public class ReferenceMatcher
   }
 
   public static List<Reference.Mismatch> checkMatch(
-      Reference reference, TypeDescription typeOnClasspath) {
+      final Reference reference, final TypeDescription typeOnClasspath) {
-    final List<Mismatch> mismatches = new ArrayList<>(0);
+    final List<Mismatch> mismatches = new ArrayList<>();
 
-    for (Reference.Flag flag : reference.getFlags()) {
+    for (final Reference.Flag flag : reference.getFlags()) {
       if (!flag.matches(typeOnClasspath.getModifiers())) {
         final String desc = reference.getClassName();
         mismatches.add(
@@ -127,8 +145,8 @@ public class ReferenceMatcher
       }
     }
 
-    for (Reference.Field fieldRef : reference.getFields()) {
+    for (final Reference.Field fieldRef : reference.getFields()) {
-      FieldDescription.InDefinedShape fieldDescription = findField(fieldRef, typeOnClasspath);
+      final FieldDescription.InDefinedShape fieldDescription = findField(fieldRef, typeOnClasspath);
       if (fieldDescription == null) {
         mismatches.add(
             new Reference.Mismatch.MissingField(
@@ -137,7 +155,7 @@ public class ReferenceMatcher
                 fieldRef.getName(),
                 fieldRef.getType().getInternalName()));
       } else {
-        for (Reference.Flag flag : fieldRef.getFlags()) {
+        for (final Reference.Flag flag : fieldRef.getFlags()) {
           if (!flag.matches(fieldDescription.getModifiers())) {
             final String desc =
                 reference.getClassName()
@@ -155,7 +173,7 @@ public class ReferenceMatcher
       }
     }
 
-    for (Reference.Method methodRef : reference.getMethods()) {
+    for (final Reference.Method methodRef : reference.getMethods()) {
       final MethodDescription.InDefinedShape methodDescription =
           findMethod(methodRef, typeOnClasspath);
       if (methodDescription == null) {
@@ -165,7 +183,7 @@ public class ReferenceMatcher
                 methodRef.getName(),
                 methodRef.getDescriptor()));
       } else {
-        for (Reference.Flag flag : methodRef.getFlags()) {
+        for (final Reference.Flag flag : methodRef.getFlags()) {
           if (!flag.matches(methodDescription.getModifiers())) {
             final String desc =
                 reference.getClassName() + "#" + methodRef.getName() + methodRef.getDescriptor();
@@ -184,8 +202,8 @@ public class ReferenceMatcher
   }
 
   private static FieldDescription.InDefinedShape findField(
-      Reference.Field fieldRef, TypeDescription typeOnClasspath) {
+      final Reference.Field fieldRef, final TypeDescription typeOnClasspath) {
-    for (FieldDescription.InDefinedShape fieldType : typeOnClasspath.getDeclaredFields()) {
+    for (final FieldDescription.InDefinedShape fieldType : typeOnClasspath.getDeclaredFields()) {
       if (fieldType.getName().equals(fieldRef.getName())
           && fieldType
               .getType()
@@ -196,14 +214,14 @@ public class ReferenceMatcher
       }
     }
     if (typeOnClasspath.getSuperClass() != null) {
-      FieldDescription.InDefinedShape fieldOnSupertype =
+      final FieldDescription.InDefinedShape fieldOnSupertype =
           findField(fieldRef, typeOnClasspath.getSuperClass().asErasure());
       if (fieldOnSupertype != null) {
         return fieldOnSupertype;
       }
     }
-    for (TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
+    for (final TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
-      FieldDescription.InDefinedShape fieldOnSupertype =
+      final FieldDescription.InDefinedShape fieldOnSupertype =
           findField(fieldRef, interfaceType.asErasure());
       if (fieldOnSupertype != null) {
         return fieldOnSupertype;
@@ -213,8 +231,8 @@ public class ReferenceMatcher
   }
 
   private static MethodDescription.InDefinedShape findMethod(
-      Reference.Method methodRef, TypeDescription typeOnClasspath) {
+      final Reference.Method methodRef, final TypeDescription typeOnClasspath) {
-    for (MethodDescription.InDefinedShape methodDescription :
+    for (final MethodDescription.InDefinedShape methodDescription :
         typeOnClasspath.getDeclaredMethods()) {
       if (methodDescription.getInternalName().equals(methodRef.getName())
           && methodDescription.getDescriptor().equals(methodRef.getDescriptor())) {
@@ -222,14 +240,14 @@ public class ReferenceMatcher
       }
     }
     if (typeOnClasspath.getSuperClass() != null) {
-      MethodDescription.InDefinedShape methodOnSupertype =
+      final MethodDescription.InDefinedShape methodOnSupertype =
           findMethod(methodRef, typeOnClasspath.getSuperClass().asErasure());
       if (methodOnSupertype != null) {
         return methodOnSupertype;
       }
     }
-    for (TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
+    for (final TypeDescription.Generic interfaceType : typeOnClasspath.getInterfaces()) {
-      MethodDescription.InDefinedShape methodOnSupertype =
+      final MethodDescription.InDefinedShape methodOnSupertype =
           findMethod(methodRef, interfaceType.asErasure());
       if (methodOnSupertype != null) {
         return methodOnSupertype;