Make muzzle reference creation package(s) configurable

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

@@ -12,6 +12,7 @@
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Predicate;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
@@ -21,11 +22,8 @@ public CouchbaseInstrumentationModule() {
  }
 
  @Override
- protected String[] additionalHelperClassNames() {
- return new String[] {
- "com.couchbase.client.tracing.opentelemetry.OpenTelemetryRequestSpan",
- "com.couchbase.client.tracing.opentelemetry.OpenTelemetryRequestTracer"
- };
+ public Predicate<String> additionalLibraryInstrumentationPackage() {
+ return className -> 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,18 @@ private String mainInstrumentationName() {
  return instrumentationNames.iterator().next();
  }
 
+ /**
+ * 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 class name predicate returned by this method is added to the default ones
+ * defined in {@link InstrumentationClassPredicate}.
+ */
+ public Predicate<String> additionalLibraryInstrumentationPackage() {
+ return className -> 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,16 @@ 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.additionalLibraryInstrumentationPackage());
  for (String adviceClass : adviceClassNames) {
  collector.collectReferencesFromAdvice(adviceClass);
  }
- for (String resource : instrumenter.helperResourceNames()) {
+ for (String resource : instrumentationModule.helperResourceNames()) {
  collector.collectReferencesFromResource(resource);
  }
  return collector;
@@ -204,8 +205,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 +469,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.javaagent.bootstrap.WeakCache;
 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;
@@ -35,17 +36,19 @@ public final class ReferenceMatcher {
  private final WeakCache<ClassLoader, Boolean> mismatchCache = AgentTooling.newWeakCache();
  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