Make muzzle reference creation package(s) configurable (#2615)

* Make muzzle reference creation package(s) configurable

* Code review comments

instrumentation/couchbase/couchbase-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/couchbase/v3_1/CouchbaseInstrumentationModule.java

@@ -21,11 +21,8 @@ public CouchbaseInstrumentationModule() {
   }
 
   @Override
-  protected String[] additionalHelperClassNames() {
-    return new String[] {
-      "com.couchbase.client.tracing.opentelemetry.OpenTelemetryRequestSpan",
-      "com.couchbase.client.tracing.opentelemetry.OpenTelemetryRequestTracer"
-    };
+  public boolean isHelperClass(String className) {
+    return className.startsWith("com.couchbase.client.tracing.opentelemetry");
   }
 
   @Override

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/InstrumentationModule.java

@@ -28,6 +28,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Predicate;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationSource;
 import net.bytebuddy.description.method.MethodDescription;
@@ -251,6 +252,31 @@ private String mainInstrumentationName() {
     return instrumentationNames.iterator().next();
   }
 
+  /**
+   * This is an internal helper method for muzzle code generation: generating {@code invokedynamic}
+   * instructions in ASM is so painful that it's much simpler and readable to just have a plain old
+   * Java helper function here.
+   */
+  @SuppressWarnings("unused")
+  protected final Predicate<String> additionalLibraryInstrumentationPackage() {
+    return this::isHelperClass;
+  }
+
+  /**
+   * Instrumentation modules can override this method to specify additional packages (or classes)
+   * that should be treated as "library instrumentation" packages. Classes from those packages will
+   * be treated by muzzle as instrumentation helper classes: they will be scanned for references and
+   * automatically injected into the application classloader if they're used in any type
+   * instrumentation. The classes for which this predicate returns {@code true} will be treated as
+   * helper classes, in addition to the default ones defined in {@link
+   * InstrumentationClassPredicate}.
+   *
+   * @param className The name of the class that may or may not be a helper class.
+   */
+  public boolean isHelperClass(String className) {
+    return false;
+  }
+
   /**
    * The actual implementation of this method is generated automatically during compilation by the
    * {@link io.opentelemetry.javaagent.tooling.muzzle.collector.MuzzleCodeGenerationPlugin}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/InstrumentationClassPredicate.java

@@ -5,6 +5,8 @@
 
 package io.opentelemetry.javaagent.tooling.muzzle;
 
+import java.util.function.Predicate;
+
 public final class InstrumentationClassPredicate {
   // javaagent instrumentation packages
   private static final String JAVAAGENT_INSTRUMENTATION_PACKAGE =
@@ -16,16 +18,28 @@ public final class InstrumentationClassPredicate {
   private static final String LIBRARY_INSTRUMENTATION_PACKAGE = "io.opentelemetry.instrumentation.";
   private static final String INSTRUMENTATION_API_PACKAGE = "io.opentelemetry.instrumentation.api.";
 
+  private final Predicate<String> additionalLibraryInstrumentationPredicate;
+
+  public InstrumentationClassPredicate(
+      Predicate<String> additionalLibraryInstrumentationPredicate) {
+    this.additionalLibraryInstrumentationPredicate = additionalLibraryInstrumentationPredicate;
+  }
+
   /**
    * Defines which classes are treated by muzzle as "internal", "helper" instrumentation classes.
    *
    * <p>This set of classes is defined by a package naming convention: all javaagent and library
    * instrumentation classes are treated as "helper" classes and are subjected to the reference
    * collection process. All others (including {@code instrumentation-api} and {@code javaagent-api}
    * modules are not scanned for references (but references to them are collected).
+   *
+   * <p>Aside from "standard" instrumentation helper class packages, instrumentation modules can
+   * pass an additional predicate to include instrumentation helper classes from 3rd party packages.
    */
-  public static boolean isInstrumentationClass(String className) {
-    return isJavaagentInstrumentationClass(className) || isLibraryInstrumentationClass(className);
+  public boolean isInstrumentationClass(String className) {
+    return isJavaagentInstrumentationClass(className)
+        || isLibraryInstrumentationClass(className)
+        || additionalLibraryInstrumentationPredicate.test(className);
   }
 
   private static boolean isJavaagentInstrumentationClass(String className) {
@@ -37,6 +51,4 @@ private static boolean isLibraryInstrumentationClass(String className) {
     return className.startsWith(LIBRARY_INSTRUMENTATION_PACKAGE)
         && !className.startsWith(INSTRUMENTATION_API_PACKAGE);
   }
-
-  private InstrumentationClassPredicate() {}
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/MuzzleCodeGenerator.java

@@ -73,7 +73,7 @@ private static class GenerateMuzzleReferenceMatcherMethodAndField extends ClassV
     private final boolean frames;
 
     private String instrumentationClassName;
-    private InstrumentationModule instrumenter;
+    private InstrumentationModule instrumentationModule;
 
     public GenerateMuzzleReferenceMatcherMethodAndField(ClassVisitor classVisitor, boolean frames) {
       super(Opcodes.ASM7, classVisitor);
@@ -90,7 +90,7 @@ public void visit(
         String[] interfaces) {
       this.instrumentationClassName = name;
       try {
-        instrumenter =
+        instrumentationModule =
             (InstrumentationModule)
                 MuzzleCodeGenerator.class
                     .getClassLoader()
@@ -143,15 +143,15 @@ public void visitEnd() {
 
     private ReferenceCollector collectReferences() {
       Set<String> adviceClassNames =
-          instrumenter.typeInstrumentations().stream()
+          instrumentationModule.typeInstrumentations().stream()
               .flatMap(typeInstrumentation -> typeInstrumentation.transformers().values().stream())
               .collect(Collectors.toSet());
 
-      ReferenceCollector collector = new ReferenceCollector();
+      ReferenceCollector collector = new ReferenceCollector(instrumentationModule::isHelperClass);
       for (String adviceClass : adviceClassNames) {
         collector.collectReferencesFromAdvice(adviceClass);
       }
-      for (String resource : instrumenter.helperResourceNames()) {
+      for (String resource : instrumentationModule.helperResourceNames()) {
         collector.collectReferencesFromResource(resource);
       }
       return collector;
@@ -204,8 +204,9 @@ private void generateMuzzleReferenceMatcherMethod(ReferenceCollector collector)
        *   if (null == this.muzzleReferenceMatcher) {
        *     this.muzzleReferenceMatcher = new ReferenceMatcher(this.getAllHelperClassNames(),
        *                                                        new Reference[]{
-       *                                                                       //reference builders
-       *                                                                       });
+       *                                                          // reference builders
+       *                                                        },
+       *                                                        this.additionalLibraryInstrumentationPackage());
        *   }
        *   return this.muzzleReferenceMatcher;
        * }
@@ -467,11 +468,19 @@ private void generateMuzzleReferenceMatcherMethod(ReferenceCollector collector)
           mv.visitInsn(Opcodes.AASTORE);
         }
 
+        mv.visitVarInsn(Opcodes.ALOAD, 0);
+        mv.visitMethodInsn(
+            Opcodes.INVOKEVIRTUAL,
+            instrumentationClassName,
+            "additionalLibraryInstrumentationPackage",
+            "()Ljava/util/function/Predicate;",
+            false);
+
         mv.visitMethodInsn(
             Opcodes.INVOKESPECIAL,
             "io/opentelemetry/javaagent/tooling/muzzle/matcher/ReferenceMatcher",
             "<init>",
-            "(Ljava/util/List;[Lio/opentelemetry/javaagent/tooling/muzzle/Reference;)V",
+            "(Ljava/util/List;[Lio/opentelemetry/javaagent/tooling/muzzle/Reference;Ljava/util/function/Predicate;)V",
             false);
         mv.visitFieldInsn(
             Opcodes.PUTFIELD,

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/ReferenceCollectingClassVisitor.java

@@ -5,9 +5,8 @@
 
 package io.opentelemetry.javaagent.tooling.muzzle.collector;
 
-import static io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate.isInstrumentationClass;
-
 import io.opentelemetry.javaagent.tooling.Utils;
+import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference.Flag;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference.Flag.ManifestationFlag;
@@ -108,7 +107,9 @@ private static Type underlyingType(Type type) {
     return type;
   }
 
+  private final InstrumentationClassPredicate instrumentationClassPredicate;
   private final boolean isAdviceClass;
+
   private final Map<String, Reference> references = new LinkedHashMap<>();
   private final Set<String> helperClasses = new HashSet<>();
   // helper super classes which are themselves also helpers
@@ -117,8 +118,10 @@ private static Type underlyingType(Type type) {
   private String refSourceClassName;
   private Type refSourceType;
 
-  ReferenceCollectingClassVisitor(boolean isAdviceClass) {
+  ReferenceCollectingClassVisitor(
+      InstrumentationClassPredicate instrumentationClassPredicate, boolean isAdviceClass) {
     super(Opcodes.ASM7);
+    this.instrumentationClassPredicate = instrumentationClassPredicate;
     this.isAdviceClass = isAdviceClass;
   }
 
@@ -136,7 +139,7 @@ Set<String> getHelperSuperClasses() {
 
   private void addExtendsReference(Reference ref) {
     addReference(ref);
-    if (isInstrumentationClass(ref.getClassName())) {
+    if (instrumentationClassPredicate.isInstrumentationClass(ref.getClassName())) {
       helperSuperClasses.add(ref.getClassName());
     }
   }
@@ -150,7 +153,7 @@ private void addReference(Reference ref) {
         references.put(ref.getClassName(), reference.merge(ref));
       }
     }
-    if (isInstrumentationClass(ref.getClassName())) {
+    if (instrumentationClassPredicate.isInstrumentationClass(ref.getClassName())) {
       helperClasses.add(ref.getClassName());
     }
   }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/collector/ReferenceCollector.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.tooling.muzzle.collector;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate.isInstrumentationClass;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Collections.singleton;
 
@@ -16,6 +15,7 @@
 import com.google.common.graph.Graphs;
 import com.google.common.graph.MutableGraph;
 import io.opentelemetry.javaagent.tooling.Utils;
+import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -32,6 +32,7 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import net.bytebuddy.jar.asm.ClassReader;
 
@@ -46,6 +47,12 @@ public class ReferenceCollector {
   private final Map<String, Reference> references = new LinkedHashMap<>();
   private final MutableGraph<String> helperSuperClassGraph = GraphBuilder.directed().build();
   private final Set<String> visitedClasses = new HashSet<>();
+  private final InstrumentationClassPredicate instrumentationClassPredicate;
+
+  public ReferenceCollector(Predicate<String> libraryInstrumentationPredicate) {
+    this.instrumentationClassPredicate =
+        new InstrumentationClassPredicate(libraryInstrumentationPredicate);
+  }
 
   /**
    * If passed {@code resource} path points to an SPI file (either Java {@link
@@ -116,7 +123,8 @@ private void visitClassesAndCollectReferences(
 
       try (InputStream in = getClassFileStream(visitedClassName)) {
         // only start from method bodies for the advice class (skips class/method references)
-        ReferenceCollectingClassVisitor cv = new ReferenceCollectingClassVisitor(isAdviceClass);
+        ReferenceCollectingClassVisitor cv =
+            new ReferenceCollectingClassVisitor(instrumentationClassPredicate, isAdviceClass);
         ClassReader reader = new ClassReader(in);
         reader.accept(cv, ClassReader.SKIP_FRAMES);
 
@@ -125,7 +133,8 @@ private void visitClassesAndCollectReferences(
           Reference reference = entry.getValue();
 
           // Don't generate references created outside of the instrumentation package.
-          if (!visitedClasses.contains(refClassName) && isInstrumentationClass(refClassName)) {
+          if (!visitedClasses.contains(refClassName)
+              && instrumentationClassPredicate.isInstrumentationClass(refClassName)) {
             instrumentationQueue.add(refClassName);
           }
           addReference(refClassName, reference);

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/matcher/ReferenceMatcher.java

@@ -5,13 +5,13 @@
 
 package io.opentelemetry.javaagent.tooling.muzzle.matcher;
 
-import static io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate.isInstrumentationClass;
 import static java.util.Collections.emptyList;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
 import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.javaagent.tooling.AgentTooling;
 import io.opentelemetry.javaagent.tooling.Utils;
+import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference;
 import io.opentelemetry.javaagent.tooling.muzzle.Reference.Source;
 import io.opentelemetry.javaagent.tooling.muzzle.matcher.HelperReferenceWrapper.Factory;
@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Predicate;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -36,17 +37,19 @@ public final class ReferenceMatcher {
       Cache.newBuilder().setWeakKeys().build();
   private final Map<String, Reference> references;
   private final Set<String> helperClassNames;
+  private final InstrumentationClassPredicate instrumentationClassPredicate;
 
-  public ReferenceMatcher(Reference... references) {
-    this(emptyList(), references);
-  }
-
-  public ReferenceMatcher(List<String> helperClassNames, Reference[] references) {
+  public ReferenceMatcher(
+      List<String> helperClassNames,
+      Reference[] references,
+      Predicate<String> libraryInstrumentationPredicate) {
     this.references = new HashMap<>(references.length);
     for (Reference reference : references) {
       this.references.put(reference.getClassName(), reference);
     }
     this.helperClassNames = new HashSet<>(helperClassNames);
+    this.instrumentationClassPredicate =
+        new InstrumentationClassPredicate(libraryInstrumentationPredicate);
   }
 
   Collection<Reference> getReferences() {
@@ -106,7 +109,7 @@ private List<Mismatch> checkMatch(Reference reference, ClassLoader loader) {
         AgentTooling.poolStrategy()
             .typePool(AgentTooling.locationStrategy().classFileLocator(loader), loader);
     try {
-      if (isInstrumentationClass(reference.getClassName())) {
+      if (instrumentationClassPredicate.isInstrumentationClass(reference.getClassName())) {
         // make sure helper class is registered
         if (!helperClassNames.contains(reference.getClassName())) {
           return Collections.singletonList(

javaagent-tooling/src/test/groovy/io/opentelemetry/javaagent/tooling/muzzle/InstrumentationClassPredicateTest.groovy

@@ -11,19 +11,26 @@ import spock.lang.Unroll
 class InstrumentationClassPredicateTest extends Specification {
   @Unroll
   def "should collect references for #desc"() {
+    setup:
+    def predicate = new InstrumentationClassPredicate({ it.startsWith("com.example.instrumentation.library") })
+
     expect:
-    InstrumentationClassPredicate.isInstrumentationClass(className)
+    predicate.isInstrumentationClass(className)
 
     where:
-    desc                              | className
-    "javaagent instrumentation class" | "io.opentelemetry.javaagent.instrumentation.some_instrumentation.Advice"
-    "library instrumentation class"   | "io.opentelemetry.instrumentation.LibraryClass"
+    desc                                       | className
+    "javaagent instrumentation class"          | "io.opentelemetry.javaagent.instrumentation.some_instrumentation.Advice"
+    "library instrumentation class"            | "io.opentelemetry.instrumentation.LibraryClass"
+    "additional library instrumentation class" | "com.example.instrumentation.library.ThirdPartyExternalInstrumentation"
   }
 
   @Unroll
   def "should not collect references for #desc"() {
+    setup:
+    def predicate = new InstrumentationClassPredicate({ false })
+
     expect:
-    !InstrumentationClassPredicate.isInstrumentationClass(className)
+    !predicate.isInstrumentationClass(className)
 
     where:
     desc                        | className