Remove usages of agent runtime specific code from reusable classes (#3492)

* Remove usages of agent runtime specific code from reusable classes

* Oops :)

* spotless

* Maybe a bit more clarity

* Doc over requireNonNull

Co-authored-by: Trask Stalnaker <trask.stalnaker@gmail.com>

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

@@ -127,7 +127,7 @@ public class AgentInstaller {
             .with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
             .with(AgentTooling.poolStrategy())
             .with(new ClassLoadListener())
-            .with(AgentTooling.locationStrategy());
+            .with(AgentTooling.locationStrategy(Utils.getBootstrapProxy()));
     // FIXME: we cannot enable it yet due to BB/JVM bug, see
     // https://github.com/raphw/byte-buddy/issues/558
     // .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)

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

@@ -15,11 +15,14 @@ import io.opentelemetry.javaagent.tooling.bytebuddy.AgentLocationStrategy;
  */
 public final class AgentTooling {
 
-  private static final AgentLocationStrategy LOCATION_STRATEGY = new AgentLocationStrategy();
   private static final AgentCachingPoolStrategy POOL_STRATEGY = new AgentCachingPoolStrategy();
 
   public static AgentLocationStrategy locationStrategy() {
-    return LOCATION_STRATEGY;
+    return locationStrategy(null);
+  }
+
+  public static AgentLocationStrategy locationStrategy(ClassLoader bootstrapProxy) {
+    return new AgentLocationStrategy(bootstrapProxy);
   }
 
   public static AgentCachingPoolStrategy poolStrategy() {

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bytebuddy/AgentLocationStrategy.java

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent.tooling.bytebuddy;
 
-import io.opentelemetry.javaagent.tooling.Utils;
 import java.util.ArrayList;
 import java.util.List;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -19,6 +18,12 @@ import net.bytebuddy.utility.JavaModule;
  */
 public class AgentLocationStrategy implements AgentBuilder.LocationStrategy {
 
+  private final ClassLoader bootstrapProxy;
+
+  public AgentLocationStrategy(ClassLoader bootstrapProxy) {
+    this.bootstrapProxy = bootstrapProxy;
+  }
+
   public ClassFileLocator classFileLocator(ClassLoader classLoader) {
     return classFileLocator(classLoader, null);
   }
@@ -26,7 +31,9 @@ public class AgentLocationStrategy implements AgentBuilder.LocationStrategy {
   @Override
   public ClassFileLocator classFileLocator(ClassLoader classLoader, JavaModule javaModule) {
     List<ClassFileLocator> locators = new ArrayList<>();
-    locators.add(ClassFileLocator.ForClassLoader.of(Utils.getBootstrapProxy()));
+    if (bootstrapProxy != null) {
+      locators.add(ClassFileLocator.ForClassLoader.of(bootstrapProxy));
+    }
     while (classLoader != null) {
       locators.add(ClassFileLocator.ForClassLoader.WeaklyReferenced.of(classLoader));
       classLoader = classLoader.getParent();

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.tooling.instrumentation;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.failSafe;
+import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
@@ -143,6 +144,9 @@ public final class InstrumentationModuleInstaller {
         Class<?> classBeingRedefined,
         ProtectionDomain protectionDomain) {
       ReferenceMatcher muzzle = getReferenceMatcher();
+      if (classLoader == BOOTSTRAP_LOADER) {
+        classLoader = Utils.getBootstrapProxy();
+      }
       boolean isMatch = muzzle.matches(classLoader);
 
       if (!isMatch) {

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

@@ -7,7 +7,6 @@ package io.opentelemetry.javaagent.tooling.muzzle.matcher;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
-import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
 import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.javaagent.extension.muzzle.ClassRef;
@@ -15,7 +14,6 @@ import io.opentelemetry.javaagent.extension.muzzle.FieldRef;
 import io.opentelemetry.javaagent.extension.muzzle.Flag;
 import io.opentelemetry.javaagent.extension.muzzle.MethodRef;
 import io.opentelemetry.javaagent.tooling.AgentTooling;
-import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationClassPredicate;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -53,16 +51,15 @@ public final class ReferenceMatcher {
   /**
    * Matcher used by ByteBuddy. Fails fast and only caches empty results, or complete results
    *
-   * @param userClassLoader Classloader to validate against (or null for bootstrap)
+   * @param userClassLoader Classloader to validate against (cannot be {@code null}, must pass
+   *     "bootstrap proxy" instead of bootstrap class loader)
    * @return true if all references match the classpath of loader
    */
   public boolean matches(ClassLoader userClassLoader) {
-    if (userClassLoader == BOOTSTRAP_LOADER) {
-      userClassLoader = Utils.getBootstrapProxy();
-    }
     return mismatchCache.computeIfAbsent(userClassLoader, this::doesMatch);
   }
 
+  // loader cannot be null, must pass "bootstrap proxy" instead of bootstrap class loader
   private boolean doesMatch(ClassLoader loader) {
     TypePool typePool = createTypePool(loader);
     for (ClassRef reference : references.values()) {
@@ -76,13 +73,11 @@ public final class ReferenceMatcher {
   /**
    * Loads the full list of mismatches. Used in debug contexts only
    *
-   * @param loader Classloader to validate against (or null for bootstrap)
+   * @param loader Classloader to validate against (cannot be {@code null}, must pass "bootstrap *
+   *     proxy" instead of bootstrap class loader)
    * @return A list of all mismatches between this ReferenceMatcher and loader's classpath.
    */
   public List<Mismatch> getMismatchedReferenceSources(ClassLoader loader) {
-    if (loader == BOOTSTRAP_LOADER) {
-      loader = Utils.getBootstrapProxy();
-    }
     TypePool typePool = createTypePool(loader);
 
     List<Mismatch> mismatches = emptyList();
@@ -94,7 +89,9 @@ public final class ReferenceMatcher {
     return mismatches;
   }
 
+  // loader cannot be null, must pass "bootstrap proxy" instead of bootstrap class loader
   private static TypePool createTypePool(ClassLoader loader) {
+    // ok to use locationStrategy() without fallback bootstrap proxy here since loader is non-null
     return AgentTooling.poolStrategy()
         .typePool(AgentTooling.locationStrategy().classFileLocator(loader), loader);
   }