Deep-merge Reference fields and methods

Plus misc cleanup.

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

@@ -132,7 +132,7 @@ public interface Instrumenter {
               log.debug(
                   "Instrumentation muzzled: {} -- {} on {}",
                   instrumentationPrimaryName,
-                  this.getClass().getName(),
+                  getClass().getName(),
                   classLoader);
               for (Reference.Mismatch mismatch : mismatches) {
                 log.debug("-- {}", mismatch);

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

@@ -84,8 +84,8 @@ public class Reference {
         className,
         superName,
         merge(interfaces, anotherReference.interfaces),
-        merge(fields, anotherReference.fields),
+        mergeFields(fields, anotherReference.fields),
-        merge(methods, anotherReference.methods));
+        mergeMethods(methods, anotherReference.methods));
   }
 
   private static <T> Set<T> merge(Set<T> set1, Set<T> set2) {
@@ -95,6 +95,32 @@ public class Reference {
     return set;
   }
 
+  private static Set<Method> mergeMethods(Set<Method> methods1, Set<Method> methods2) {
+    List<Method> merged = new ArrayList<>(methods1);
+    for (Method method : methods2) {
+      int i = merged.indexOf(method);
+      if (i == -1) {
+        merged.add(method);
+      } else {
+        merged.set(i, merged.get(i).merge(method));
+      }
+    }
+    return new HashSet<>(merged);
+  }
+
+  private static Set<Field> mergeFields(Set<Field> fields1, Set<Field> fields2) {
+    List<Field> merged = new ArrayList<>(fields1);
+    for (Field field : fields2) {
+      int i = merged.indexOf(field);
+      if (i == -1) {
+        merged.add(field);
+      } else {
+        merged.set(i, merged.get(i).merge(field));
+      }
+    }
+    return new HashSet<>(merged);
+  }
+
   private static Set<Flag> mergeFlags(Set<Flag> flags1, Set<Flag> flags2) {
     Set<Flag> merged = merge(flags1, flags2);
     // TODO: Assert flags are non-contradictory and resolve
@@ -144,8 +170,15 @@ public class Reference {
     }
   }
 
+  /**
+   * A mismatch between a Reference and a runtime class.
+   *
+   * <p>This class' toString returns a human-readable description of the mismatch along with
+   * source-code locations of the instrumentation which caused the mismatch.
+   */
   public abstract static class Mismatch {
-    final Source[] mismatchSources;
+    /** Instrumentation sources which caused the mismatch. */
+    private final Source[] mismatchSources;
 
     Mismatch(Source[] mismatchSources) {
       this.mismatchSources = mismatchSources;
@@ -160,10 +193,11 @@ public class Reference {
       }
     }
 
+    /** Human-readable string describing the mismatch. */
     abstract String getMismatchDetails();
 
     public static class MissingClass extends Mismatch {
-      final String className;
+      private final String className;
 
       public MissingClass(Source[] sources, String className) {
         super(sources);
@@ -177,9 +211,9 @@ public class Reference {
     }
 
     public static class MissingFlag extends Mismatch {
-      final Flag expectedFlag;
+      private final Flag expectedFlag;
-      final String classMethodOrFieldDesc;
+      private final String classMethodOrFieldDesc;
-      final int foundAccess;
+      private final int foundAccess;
 
       public MissingFlag(
           Source[] sources, String classMethodOrFieldDesc, Flag expectedFlag, int foundAccess) {
@@ -197,14 +231,14 @@ public class Reference {
 
     /** Fallback mismatch in case an unexpected exception occurs during reference checking. */
     public static class ReferenceCheckError extends Mismatch {
-      private final Exception referenceCheckExcetpion;
+      private final Exception referenceCheckException;
       private final Reference referenceBeingChecked;
       private final ClassLoader classLoaderBeingChecked;
 
       public ReferenceCheckError(
           Exception e, Reference referenceBeingChecked, ClassLoader classLoaderBeingChecked) {
         super(new Source[0]);
-        this.referenceCheckExcetpion = e;
+        this.referenceCheckException = e;
         this.referenceBeingChecked = referenceBeingChecked;
         this.classLoaderBeingChecked = classLoaderBeingChecked;
       }
@@ -219,7 +253,7 @@ public class Reference {
         sw.write("\n");
         // add exception message and stack trace
         final PrintWriter pw = new PrintWriter(sw);
-        referenceCheckExcetpion.printStackTrace(pw);
+        referenceCheckException.printStackTrace(pw);
         return sw.toString();
       }
     }
@@ -243,8 +277,8 @@ public class Reference {
     }
 
     public static class MissingMethod extends Mismatch {
-      final String className;
+      private final String className;
-      final String method;
+      private final String method;
 
       public MissingMethod(Source[] sources, String className, String method) {
         super(sources);
@@ -299,7 +333,7 @@ public class Reference {
 
       @Override
       public boolean matches(int asmFlags) {
-        return (Opcodes.ACC_PUBLIC & asmFlags) != 0 || (Opcodes.ACC_PROTECTED & asmFlags) != 0;
+        return PUBLIC.matches(asmFlags) || (Opcodes.ACC_PROTECTED & asmFlags) != 0;
       }
     },
     PRIVATE_OR_HIGHER {
@@ -464,8 +498,6 @@ public class Reference {
 
     @Override
     public String toString() {
-      // <init>()V
-      // toString()Ljava/lang/String;
       return name + getDescriptor();
     }
 
@@ -476,7 +508,8 @@ public class Reference {
     @Override
     public boolean equals(Object o) {
       if (o instanceof Method) {
-        return toString().equals(o.toString());
+        final Method m = (Method) o;
+        return name.equals(m.name) && getDescriptor().equals(m.getDescriptor());
       }
       return false;
     }

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

@@ -26,6 +26,14 @@ import net.bytebuddy.jar.asm.Type;
 // - inner class
 // - cast opcodes in method bodies
 public class ReferenceCreator extends ClassVisitor {
+  /**
+   * Classes in this namespace will be scanned and used to create references.
+   *
+   * <p>For now we're hardcoding this to the instrumentation package so we only create references
+   * from the method advice and helper classes.
+   */
+  private static String REFERENCE_CREATION_PACKAGE = "datadog.trace.instrumentation.";
+
   public static Map<String, Reference> createReferencesFrom(
       String entryPointClassName, ClassLoader loader) {
     return ReferenceCreator.createReferencesFrom(entryPointClassName, loader, true);
@@ -63,7 +71,7 @@ public class ReferenceCreator extends ClassVisitor {
           for (Map.Entry<String, Reference> entry : instrumentationReferences.entrySet()) {
             // Don't generate references created outside of the datadog instrumentation package.
             if (!visitedSources.contains(entry.getKey())
-                && entry.getKey().startsWith("datadog.trace.instrumentation.")) {
+                && entry.getKey().startsWith(REFERENCE_CREATION_PACKAGE)) {
               instrumentationQueue.add(entry.getKey());
             }
             if (references.containsKey(entry.getKey())) {
@@ -86,6 +94,11 @@ public class ReferenceCreator extends ClassVisitor {
     return references;
   }
 
+  /**
+   * Get the package of an internal class name.
+   *
+   * <p>foo/bar/Baz -> foo/bar/
+   */
   private static String internalPackageName(String internalName) {
     return internalName.replaceAll("/[^/]+$", "");
   }
@@ -139,6 +152,17 @@ public class ReferenceCreator extends ClassVisitor {
     }
   }
 
+  /**
+   * @return If TYPE is an array, return the underlying type. If TYPE is not an array simply return
+   *     the type.
+   */
+  private static Type underlyingType(Type type) {
+    while (type.getSort() == Type.ARRAY) {
+      type = type.getElementType();
+    }
+    return type;
+  }
+
   private Map<String, Reference> references = new HashMap<>();
   private String refSourceClassName;
   private Type refSourceType;
@@ -249,10 +273,7 @@ public class ReferenceCreator extends ClassVisitor {
                   fieldType)
               .build());
 
-      Type underlyingFieldType = fieldType;
+      final Type underlyingFieldType = underlyingType(fieldType);
-      while (underlyingFieldType.getSort() == Type.ARRAY) {
-        underlyingFieldType = underlyingFieldType.getElementType();
-      }
       if (underlyingFieldType.getSort() == Type.OBJECT) {
         addReference(
             new Reference.Builder(underlyingFieldType.getInternalName())
@@ -282,10 +303,7 @@ public class ReferenceCreator extends ClassVisitor {
       final Type methodType = Type.getMethodType(descriptor);
 
       { // ref for method return type
-        Type returnType = methodType.getReturnType();
+        final Type returnType = underlyingType(methodType.getReturnType());
-        while (returnType.getSort() == Type.ARRAY) {
-          returnType = returnType.getElementType();
-        }
         if (returnType.getSort() == Type.OBJECT) {
           addReference(
               new Reference.Builder(returnType.getInternalName())
@@ -296,9 +314,7 @@ public class ReferenceCreator extends ClassVisitor {
       }
       // refs for method param types
       for (Type paramType : methodType.getArgumentTypes()) {
-        while (paramType.getSort() == Type.ARRAY) {
+        paramType = underlyingType(paramType);
-          paramType = paramType.getElementType();
-        }
         if (paramType.getSort() == Type.OBJECT) {
           addReference(
               new Reference.Builder(paramType.getInternalName())