Add reference pretty-printer to muzzle plugin

buildSrc/src/main/groovy/MuzzlePlugin.groovy

@@ -14,48 +14,76 @@ class MuzzlePlugin implements Plugin<Project> {
     def bootstrapProject = project.rootProject.getChildProjects().get('dd-java-agent').getChildProjects().get('agent-bootstrap')
     def toolingProject = project.rootProject.getChildProjects().get('dd-java-agent').getChildProjects().get('agent-tooling')
     project.extensions.create("muzzle", MuzzleExtension)
+    def compileMuzzle = project.task('compileMuzzle') {
+      // not adding user and group to hide this from `gradle tasks`
+    }
     def muzzle = project.task('muzzle') {
       group = 'Muzzle'
       description = "Run instrumentation muzzle on compile time dependencies"
       doLast {
-        List<URL> userUrls = new ArrayList<>()
-        project.getLogger().info("Creating user classpath for: " + project.getName())
-        for (File f : project.configurations.compileOnly.getFiles()) {
-          project.getLogger().info('--' + f)
-          userUrls.add(f.toURI().toURL())
-        }
-        for (File f : bootstrapProject.sourceSets.main.runtimeClasspath.getFiles()) {
-          project.getLogger().info('--' + f)
-          userUrls.add(f.toURI().toURL())
-        }
-        final ClassLoader userCL = new URLClassLoader(userUrls.toArray(new URL[0]), (ClassLoader) null)
-
-        project.getLogger().info("Creating dd classpath for: " + project.getName())
-        Set<URL> ddUrls = new HashSet<>()
-        for (File f : toolingProject.sourceSets.main.runtimeClasspath.getFiles()) {
-          project.getLogger().info('--' + f)
-          ddUrls.add(f.toURI().toURL())
-        }
-        for (File f : project.sourceSets.main.runtimeClasspath.getFiles()) {
-          project.getLogger().info('--' + f)
-          ddUrls.add(f.toURI().toURL())
-        }
-
-        final ClassLoader agentCL = new URLClassLoader(ddUrls.toArray(new URL[0]), (ClassLoader) null)
+        final ClassLoader userCL = createUserClassLoader(project, bootstrapProject)
+        final ClassLoader agentCL = createDDClassloader(project, toolingProject)
         // find all instrumenters, get muzzle, and assert
         Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
           .getMethod('assertInstrumentationNotMuzzled', ClassLoader.class)
         assertionMethod.invoke(null, userCL)
       }
     }
-    // project.tasks.muzzle.dependsOn(bootstrapProject.tasks.shadowJar)
-    project.tasks.muzzle.dependsOn(bootstrapProject.tasks.compileJava)
-    project.tasks.muzzle.dependsOn(toolingProject.tasks.compileJava)
+    def printReferences = project.task('printReferences') {
+      group = 'Muzzle'
+      description = "Print references created by instrumentation muzzle"
+      doLast {
+        final ClassLoader agentCL = createDDClassloader(project, toolingProject)
+        Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
+          .getMethod('printMuzzleReferences')
+        assertionMethod.invoke(null)
+      }
+    }
+    project.tasks.compileMuzzle.dependsOn(bootstrapProject.tasks.compileJava)
+    project.tasks.compileMuzzle.dependsOn(toolingProject.tasks.compileJava)
     project.afterEvaluate {
-      project.tasks.muzzle.dependsOn(project.tasks.compileJava)
+      project.tasks.compileMuzzle.dependsOn(project.tasks.compileJava)
       if (project.tasks.getNames().contains("compileScala")) {
-        project.tasks.muzzle.dependsOn(project.tasks.compileScala)
+        project.tasks.compileMuzzle.dependsOn(project.tasks.compileScala)
       }
     }
+
+    project.tasks.muzzle.dependsOn(project.tasks.compileMuzzle)
+    project.tasks.printReferences.dependsOn(project.tasks.compileMuzzle)
+  }
+
+  /**
+   * Create a classloader with core agent classes and project instrumentation on the classpath.
+   */
+  private ClassLoader createDDClassloader(Project project, Project toolingProject) {
+    project.getLogger().info("Creating dd classpath for: " + project.getName())
+    Set<URL> ddUrls = new HashSet<>()
+    for (File f : toolingProject.sourceSets.main.runtimeClasspath.getFiles()) {
+      project.getLogger().info('--' + f)
+      ddUrls.add(f.toURI().toURL())
+    }
+    for (File f : project.sourceSets.main.runtimeClasspath.getFiles()) {
+      project.getLogger().info('--' + f)
+      ddUrls.add(f.toURI().toURL())
+    }
+
+    return new URLClassLoader(ddUrls.toArray(new URL[0]), (ClassLoader) null)
+  }
+
+  /**
+   * Create a classloader with user/library classes on the classpath.
+   */
+  private ClassLoader createUserClassLoader(Project project, Project bootstrapProject) {
+    List<URL> userUrls = new ArrayList<>()
+    project.getLogger().info("Creating user classpath for: " + project.getName())
+    for (File f : project.configurations.compileOnly.getFiles()) {
+      project.getLogger().info('--' + f)
+      userUrls.add(f.toURI().toURL())
+    }
+    for (File f : bootstrapProject.sourceSets.main.runtimeClasspath.getFiles()) {
+      project.getLogger().info('--' + f)
+      userUrls.add(f.toURI().toURL())
+    }
+    return new URLClassLoader(userUrls.toArray(new URL[0]), (ClassLoader) null)
   }
 }

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

@@ -82,5 +82,71 @@ public static void assertInstrumentationNotMuzzled(ClassLoader cl) throws Except
     }
   }
 
+  public static void printMuzzleReferences() {
+    for (Instrumenter instrumenter :
+        ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+      if (instrumenter instanceof Instrumenter.Default) {
+        try {
+          final Method getMuzzleMethod =
+              instrumenter.getClass().getDeclaredMethod("getInstrumentationMuzzle");
+          final ReferenceMatcher muzzle;
+          try {
+            getMuzzleMethod.setAccessible(true);
+            muzzle = (ReferenceMatcher) getMuzzleMethod.invoke(instrumenter);
+          } finally {
+            getMuzzleMethod.setAccessible(false);
+          }
+          System.out.println(instrumenter.getClass().getName());
+          for (Reference ref : muzzle.getReferences()) {
+            System.out.println(prettyPrint("  ", ref));
+          }
+        } catch (Exception e) {
+          System.out.println(
+              "Unexpected exception printing references for " + instrumenter.getClass().getName());
+          throw new RuntimeException(e);
+        }
+      } else {
+        throw new RuntimeException(
+            "class "
+                + instrumenter.getClass().getName()
+                + " is not a default instrumenter. No refs to print.");
+      }
+    }
+  }
+
+  private static String prettyPrint(String prefix, Reference ref) {
+    final StringBuilder builder = new StringBuilder(prefix).append(ref.getClassName());
+    if (ref.getSuperName() != null) {
+      builder.append(" extends<").append(ref.getSuperName()).append(">");
+    }
+    if (ref.getInterfaces().size() > 0) {
+      builder.append(" implements ");
+      for (String iface : ref.getInterfaces()) {
+        builder.append(" <").append(iface).append(">");
+      }
+    }
+    for (final Reference.Source source : ref.getSources()) {
+      builder.append("\n").append(prefix).append(prefix);
+      builder.append("Source: ").append(source.toString());
+    }
+    for (final Reference.Field field : ref.getFields()) {
+      builder.append("\n").append(prefix).append(prefix);
+      builder.append("Field: ");
+      for (Reference.Flag flag : field.getFlags()) {
+        builder.append(flag).append(" ");
+      }
+      builder.append(field.toString());
+    }
+    for (final Reference.Method method : ref.getMethods()) {
+      builder.append("\n").append(prefix).append(prefix);
+      builder.append("Method: ");
+      for (Reference.Flag flag : method.getFlags()) {
+        builder.append(flag).append(" ");
+      }
+      builder.append(method.toString());
+    }
+    return builder.toString();
+  }
+
   private MuzzleVersionScanPlugin() {}
 }

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

@@ -1,6 +1,5 @@
 package datadog.trace.agent.tooling.muzzle;
 
-import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
 import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
@@ -9,13 +8,12 @@
 import datadog.trace.agent.tooling.muzzle.Reference.Mismatch;
 import datadog.trace.agent.tooling.muzzle.Reference.Source;
 import datadog.trace.bootstrap.WeakMap;
-import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.WeakHashMap;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.method.MethodDescription;
@@ -38,6 +36,10 @@ public ReferenceMatcher(final String[] helperClassNames, final Reference[] refer
     this.helperClassNames = new HashSet<>(Arrays.asList(helperClassNames));
   }
 
+  public Reference[] getReferences() {
+    return references;
+  }
+
   /**
    * @param loader Classloader to validate against (or null for bootstrap)
    * @return true if all references match the classpath of loader