Merge pull request #441 from DataDog/ark/muzzle-integration-branch

muzzle integration branch

buildSrc/src/main/groovy/MuzzlePlugin.groovy

@@ -14,45 +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.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-ittests/src/test/groovy/datadog/trace/agent/integration/classloading/ClassLoadingTest.groovy

@@ -56,6 +56,7 @@ class ClassLoadingTest extends Specification {
     @Override
     URL getResource(String name) {
       count++
+      return super.getResource(name)
     }
   }
 
@@ -65,10 +66,13 @@ class ClassLoadingTest extends Specification {
 
     when:
     loader.loadClass(ClassToInstrument.getName())
+    int countAfterFirstLoad = loader.count
     loader.loadClass(ClassToInstrumentChild.getName())
 
     then:
-    loader.count == 2
+    // ClassToInstrumentChild won't cause an additional getResource() because its TypeDescription is created from transformation bytes.
+    loader.count > 0
+    loader.count == countAfterFirstLoad
   }
 
   def "make sure that ByteBuddy doesn't resue cached type descriptions between different classloaders"() {
@@ -81,8 +85,9 @@ class ClassLoadingTest extends Specification {
     loader2.loadClass(ClassToInstrument.getName())
 
     then:
-    loader1.count == 1
-    loader2.count == 1
+    loader1.count > 0
+    loader2.count > 0
+    loader1.count == loader2.count
   }
 
   def "can find bootstrap resources"() {

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/PatchLogger.java

@@ -99,6 +99,20 @@ public void logrb(
       String msg,
       Throwable thrown) {}
 
+  public void severe(String msg) {}
+
+  public void warning(String msg) {}
+
+  public void info(String msg) {}
+
+  public void config(String msg) {}
+
+  public void fine(String msg) {}
+
+  public void finer(String msg) {}
+
+  public void finest(String msg) {}
+
   public void throwing(String sourceClass, String sourceMethod, Throwable thrown) {}
 
   public void setLevel(Level newLevel) throws SecurityException {}

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

@@ -20,6 +20,8 @@
 
 @Slf4j
 public class AgentInstaller {
+  public static final DDLocationStrategy LOCATION_STRATEGY = new DDLocationStrategy();
+  public static final AgentBuilder.PoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy();
   private static volatile Instrumentation INSTRUMENTATION;
 
   public static Instrumentation getInstrumentation() {
@@ -46,8 +48,9 @@ public static ResettableClassFileTransformer installBytebuddyAgent(
             .disableClassFormatChanges()
             .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
             .with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
+            .with(POOL_STRATEGY)
             .with(new LoggingListener())
-            .with(new DDLocationStrategy())
+            .with(LOCATION_STRATEGY)
             .ignore(any(), skipClassLoader())
             .or(nameStartsWith("datadog.trace."))
             .or(nameStartsWith("datadog.opentracing."))

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

@@ -1,14 +1,14 @@
 package datadog.trace.agent.tooling;
 
-import static net.bytebuddy.matcher.ElementMatchers.erasure;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.build.HashCodeAndEqualsPlugin;
 import net.bytebuddy.description.type.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.description.type.TypeList;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 
@@ -30,7 +30,7 @@ public class ByteBuddyElementMatchers {
    */
   public static <T extends TypeDescription> ElementMatcher.Junction<T> safeHasSuperType(
       final ElementMatcher<? super TypeDescription> matcher) {
-    return safeHasGenericSuperType(erasure(matcher));
+    return safeHasGenericSuperType(new SafeErasureMatcher(matcher));
   }
 
   /**
@@ -48,6 +48,29 @@ public static <T extends TypeDescription> ElementMatcher.Junction<T> safeHasGene
     return new SafeHasSuperTypeMatcher<>(matcher);
   }
 
+  /**
+   * Wraps another matcher to assure that an element is not matched in case that the matching causes
+   * an {@link Exception}. Logs exception if it happens.
+   *
+   * @param matcher The element matcher that potentially throws an exception.
+   * @param <T> The type of the matched object.
+   * @return A matcher that returns {@code false} in case that the given matcher throws an
+   *     exception.
+   */
+  public static <T> ElementMatcher.Junction<T> failSafe(
+      final ElementMatcher<? super T> matcher, final String description) {
+    return new SafeMatcher<>(matcher, false, description);
+  }
+
+  private static TypeDescription safeAsErasure(final TypeDefinition target) {
+    try {
+      return target.asErasure();
+    } catch (final Exception e) {
+      log.debug("Exception trying to get interfaces:", e);
+      return null;
+    }
+  }
+
   /**
    * An element matcher that matches a super type. This is different from {@link
    * net.bytebuddy.matcher.HasSuperTypeMatcher} in the following way:
@@ -102,7 +125,7 @@ private TypeDefinition safeGetSuperClass(final TypeDefinition typeDefinition) {
       try {
         return typeDefinition.getSuperClass();
       } catch (final Exception e) {
-        log.info("Exception trying to get next type definition:", e);
+        log.debug("Exception trying to get next type definition:", e);
         return null;
       }
     }
@@ -117,27 +140,135 @@ private TypeDefinition safeGetSuperClass(final TypeDefinition typeDefinition) {
     private boolean hasInterface(
         final TypeDefinition typeDefinition, final Set<TypeDescription> checkedInterfaces) {
       for (final TypeDefinition interfaceType : safeGetInterfaces(typeDefinition)) {
-        if (checkedInterfaces.add(interfaceType.asErasure())
-            && (matcher.matches(interfaceType.asGenericType())
-                || hasInterface(interfaceType, checkedInterfaces))) {
-          return true;
+        final TypeDescription erasure = safeAsErasure(interfaceType);
+        if (erasure != null) {
+          if (checkedInterfaces.add(interfaceType.asErasure())
+              && (matcher.matches(interfaceType.asGenericType())
+                  || hasInterface(interfaceType, checkedInterfaces))) {
+            return true;
+          }
         }
       }
       return false;
     }
 
-    private TypeList.Generic safeGetInterfaces(final TypeDefinition typeDefinition) {
+    /**
+     * TypeDefinition#getInterfaces() produces an interator which may throw an exception during
+     * iteration if an interface is absent from the classpath.
+     *
+     * <p>This method exists to allow getting interfaces even if the lookup on one fails.
+     */
+    private List<TypeDefinition> safeGetInterfaces(final TypeDefinition typeDefinition) {
+      final List<TypeDefinition> interfaceTypes = new ArrayList<>();
       try {
-        return typeDefinition.getInterfaces();
+        final Iterator<TypeDescription.Generic> interfaceIter =
+            typeDefinition.getInterfaces().iterator();
+        while (interfaceIter.hasNext()) {
+          interfaceTypes.add(interfaceIter.next());
+        }
       } catch (final Exception e) {
-        log.info("Exception trying to get interfaces:", e);
-        return new TypeList.Generic.Empty();
+        log.debug("Exception trying to get interfaces:", e);
       }
+      return interfaceTypes;
     }
 
     @Override
     public String toString() {
       return "safeHasSuperType(" + matcher + ")";
     }
   }
+
+  /**
+   * An element matcher that matches its argument's {@link TypeDescription.Generic} raw type against
+   * the given matcher for a {@link TypeDescription}. As a wildcard does not define an erasure, a
+   * runtime exception is thrown when this matcher is applied to a wildcard.
+   *
+   * <p>Catches and logs exception if it was thrown when getting erasure, returning false.
+   *
+   * @param <T> The type of the matched entity.
+   * @see net.bytebuddy.matcher.ErasureMatcher
+   */
+  @HashCodeAndEqualsPlugin.Enhance
+  public static class SafeErasureMatcher<T extends TypeDefinition>
+      extends ElementMatcher.Junction.AbstractBase<T> {
+
+    /** The matcher to apply to the raw type of the matched element. */
+    private final ElementMatcher<? super TypeDescription> matcher;
+
+    /**
+     * Creates a new erasure matcher.
+     *
+     * @param matcher The matcher to apply to the raw type.
+     */
+    public SafeErasureMatcher(final ElementMatcher<? super TypeDescription> matcher) {
+      this.matcher = matcher;
+    }
+
+    @Override
+    public boolean matches(final T target) {
+      final TypeDescription erasure = safeAsErasure(target);
+      if (erasure == null) {
+        return false;
+      } else {
+        // We would like matcher exceptions to propagate
+        return matcher.matches(erasure);
+      }
+    }
+
+    @Override
+    public String toString() {
+      return "safeErasure(" + matcher + ")";
+    }
+  }
+
+  /**
+   * A fail-safe matcher catches exceptions that are thrown by a delegate matcher and returns an
+   * alternative value.
+   *
+   * <p>Logs exception if it was thrown.
+   *
+   * @param <T> The type of the matched entity.
+   * @see net.bytebuddy.matcher.FailSafeMatcher
+   */
+  @HashCodeAndEqualsPlugin.Enhance
+  public static class SafeMatcher<T> extends ElementMatcher.Junction.AbstractBase<T> {
+
+    /** The delegate matcher that might throw an exception. */
+    private final ElementMatcher<? super T> matcher;
+
+    /** The fallback value in case of an exception. */
+    private final boolean fallback;
+
+    /** The text description to log if exception happens. */
+    private final String description;
+
+    /**
+     * Creates a new fail-safe element matcher.
+     *
+     * @param matcher The delegate matcher that might throw an exception.
+     * @param fallback The fallback value in case of an exception.
+     * @param description Descriptive string to log along with exception.
+     */
+    public SafeMatcher(
+        final ElementMatcher<? super T> matcher, final boolean fallback, final String description) {
+      this.matcher = matcher;
+      this.fallback = fallback;
+      this.description = description;
+    }
+
+    @Override
+    public boolean matches(final T target) {
+      try {
+        return matcher.matches(target);
+      } catch (final Exception e) {
+        log.debug(description, e);
+        return fallback;
+      }
+    }
+
+    @Override
+    public String toString() {
+      return "safeMatcher(try(" + matcher + ") or " + fallback + ")";
+    }
+  }
 }