Make internal classloader instrumentation indy compatible

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/BootDelegationInstrumentation.java

@@ -47,10 +47,7 @@ public class BootDelegationInstrumentation implements TypeInstrumentation {
   public ElementMatcher<TypeDescription> typeMatcher() {
     // just an optimization to exclude common class loaders that are known to delegate to the
     // bootstrap loader (or happen to _be_ the bootstrap loader)
-    return not(namedOneOf(
-            "java.lang.ClassLoader",
-            "com.ibm.oti.vm.BootstrapClassLoader",
-            "io.opentelemetry.javaagent.bootstrap.AgentClassLoader"))
+    return not(namedOneOf("java.lang.ClassLoader", "com.ibm.oti.vm.BootstrapClassLoader"))
         .and(extendsClass(named("java.lang.ClassLoader")));
   }
 
@@ -136,13 +133,14 @@ public static Class<?> onEnter(@Advice.Argument(0) String name) {
       return null;
     }
 
-    @Advice.OnMethodExit(onThrowable = Throwable.class)
-    public static void onExit(
-        @Advice.Return(readOnly = false) Class<?> result,
-        @Advice.Enter Class<?> resultFromBootstrapLoader) {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    @Advice.AssignReturned.ToReturned
+    public static Class<?> onExit(
+        @Advice.Return Class<?> result, @Advice.Enter Class<?> resultFromBootstrapLoader) {
       if (resultFromBootstrapLoader != null) {
-        result = resultFromBootstrapLoader;
+        return resultFromBootstrapLoader;
       }
+      return result;
     }
   }
 }

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ClassLoaderInstrumentationModule.java

@@ -10,11 +10,13 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class ClassLoaderInstrumentationModule extends InstrumentationModule {
+public class ClassLoaderInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public ClassLoaderInstrumentationModule() {
     super("internal-class-loader");
   }
@@ -26,13 +28,24 @@ public boolean defaultEnabled(ConfigProperties config) {
   }
 
   @Override
-  public boolean isIndyModule() {
-    return false;
+  public boolean isHelperClass(String className) {
+    // TODO: this can be removed when we drop inlined-advice support
+    // The advices can directly access this class in the AgentClassLoader with invokedynamic Advice
+    return className.equals("io.opentelemetry.javaagent.tooling.Constants");
   }
 
   @Override
-  public boolean isHelperClass(String className) {
-    return className.equals("io.opentelemetry.javaagent.tooling.Constants");
+  public boolean loadAdviceClassesEagerly() {
+    // This is required due to DefineClassInstrumentation
+    // Without this, we would get an infinite recursion on bootstrapping of that instrumentation:
+    // 1. ClassLoader.defineClass is invoked somewhere
+    // 2. The inserted invokedynamic for the instrumentation is reached
+    // 3. To bootstrap the advice, IndyBootstrap is invoked
+    // 4. IndyBootstrap tries to load the DefineClassInstrumentation Advice class
+    // 5. The loading calls ClassLoader.defineClass -> recursion, BOOM!
+    // We avoid this recursion by ensuring that the DefineClassInstrumentation Advice class
+    // is loaded eagerly before the corresponding invokedynamic is reached
+    return true;
   }
 
   @Override

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/DefineClassInstrumentation.java

@@ -52,9 +52,17 @@ public static DefineClassContext onEnter(
           classLoader, className, classBytes, offset, length);
     }
 
+    // TODO: the ToReturned does nothing except for signaling the AdviceTransformer that it must
+    // not touch this advice
+    // this is done to ensure that ClassLoaderInstrumentationModule.loadAdviceClassesEagerly works
+    // correctly
+    // we can therefore remove it, as soon as the AdviceTransformer is not applied anymore
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(@Advice.Enter DefineClassContext context) {
+    @Advice.AssignReturned.ToReturned
+    public static Class<?> onExit(
+        @Advice.Enter DefineClassContext context, @Advice.Return Class<?> returned) {
       DefineClassHelper.afterDefineClass(context);
+      return returned;
     }
   }
 
@@ -68,9 +76,17 @@ public static DefineClassContext onEnter(
       return DefineClassHelper.beforeDefineClass(classLoader, className, classBytes);
     }
 
+    // TODO: the ToReturned does nothing except for signaling the AdviceTransformer that it must
+    // not touch this advice
+    // this is done to ensure that ClassLoaderInstrumentationModule.loadAdviceClassesEagerly works
+    // correctly
+    // we can therefore remove it, as soon as the AdviceTransformer is not applied anymore
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(@Advice.Enter DefineClassContext context) {
+    @Advice.AssignReturned.ToReturned
+    public static Class<?> onExit(
+        @Advice.Enter DefineClassContext context, @Advice.Return Class<?> returned) {
       DefineClassHelper.afterDefineClass(context);
+      return returned;
     }
   }
 }

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/LoadInjectedClassInstrumentation.java

@@ -65,11 +65,13 @@ public static Class<?> onEnter(
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class)
-    public static void onExit(
-        @Advice.Return(readOnly = false) Class<?> result, @Advice.Enter Class<?> loadedClass) {
+    @Advice.AssignReturned.ToReturned
+    public static Class<?> onExit(
+        @Advice.Return Class<?> result, @Advice.Enter Class<?> loadedClass) {
       if (loadedClass != null) {
-        result = loadedClass;
+        return loadedClass;
       }
+      return result;
     }
   }
 }

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ResourceInjectionInstrumentation.java

@@ -8,6 +8,8 @@
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
@@ -32,7 +34,12 @@ public class ResourceInjectionInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return extendsClass(named("java.lang.ClassLoader"));
+    return extendsClass(named("java.lang.ClassLoader"))
+        .and(
+            not(
+                namedOneOf(
+                    "io.opentelemetry.javaagent.bootstrap.AgentClassLoader",
+                    "io.opentelemetry.javaagent.tooling.instrumentation.indy.InstrumentationModuleClassLoader")));
   }
 
   @Override
@@ -61,37 +68,37 @@ public void transform(TypeTransformer transformer) {
   public static class GetResourceAdvice {
 
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(
+    @Advice.AssignReturned.ToReturned
+    public static URL onExit(
         @Advice.This ClassLoader classLoader,
         @Advice.Argument(0) String name,
-        @Advice.Return(readOnly = false) URL resource) {
-      if (resource != null) {
-        return;
-      }
-
-      URL helper = HelperResources.loadOne(classLoader, name);
-      if (helper != null) {
-        resource = helper;
+        @Advice.Return URL resource) {
+      if (resource == null) {
+        URL helper = HelperResources.loadOne(classLoader, name);
+        if (helper != null) {
+          return helper;
+        }
       }
+      return resource;
     }
   }
 
   @SuppressWarnings("unused")
   public static class GetResourcesAdvice {
 
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(
+    @Advice.AssignReturned.ToReturned
+    public static Enumeration<URL> onExit(
         @Advice.This ClassLoader classLoader,
         @Advice.Argument(0) String name,
-        @Advice.Return(readOnly = false) Enumeration<URL> resources) {
+        @Advice.Return Enumeration<URL> resources) {
       List<URL> helpers = HelperResources.loadAll(classLoader, name);
       if (helpers.isEmpty()) {
-        return;
+        return resources;
       }
 
       if (!resources.hasMoreElements()) {
-        resources = Collections.enumeration(helpers);
-        return;
+        return Collections.enumeration(helpers);
       }
 
       List<URL> result = Collections.list(resources);
@@ -108,31 +115,30 @@ public static void onExit(
           result.add(helperUrl);
         }
       }
-
-      resources = Collections.enumeration(result);
+      return Collections.enumeration(result);
     }
   }
 
   @SuppressWarnings("unused")
   public static class GetResourceAsStreamAdvice {
 
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void onExit(
+    @Advice.AssignReturned.ToReturned
+    public static InputStream onExit(
         @Advice.This ClassLoader classLoader,
         @Advice.Argument(0) String name,
-        @Advice.Return(readOnly = false) InputStream inputStream) {
-      if (inputStream != null) {
-        return;
-      }
-
-      URL helper = HelperResources.loadOne(classLoader, name);
-      if (helper != null) {
-        try {
-          inputStream = helper.openStream();
-        } catch (IOException ignored) {
-          // ClassLoader.getResourceAsStream also ignores io exceptions from opening the stream
+        @Advice.Return InputStream inputStream) {
+      if (inputStream == null) {
+        URL helper = HelperResources.loadOne(classLoader, name);
+        if (helper != null) {
+          try {
+            return helper.openStream();
+          } catch (IOException ignored) {
+            // ClassLoader.getResourceAsStream also ignores io exceptions from opening the stream
+          }
         }
       }
+      return inputStream;
     }
   }
 }

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/AgentClassLoader.java

@@ -142,6 +142,15 @@ private static int getJavaVersion() {
     return Integer.parseInt(javaSpecVersion);
   }
 
+  @Override
+  public Class<?> loadClass(String name) throws ClassNotFoundException {
+    // We explicitly override loadClass from ClassLoader to ensure
+    // that loadClass is properly excluded from our internal ClassLoader Instrumentations
+    // (e.g. LoadInjectedClassInstrumentation, BooDelegationInstrumentation)
+    // Otherwise this will cause recursion in invokedynamic linkage
+    return loadClass(name, false);
+  }
+
   @Override
   public Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
     // ContextStorageOverride is meant for library instrumentation we don't want it to apply to our

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java

@@ -61,4 +61,19 @@ default String getModuleGroup() {
   default List<String> agentPackagesToHide() {
     return Collections.emptyList();
   }
+
+  /**
+   * By default, Advice classes are loaded lazily the first time the corresponding instrumented
+   * method is invoked. This function allows to change the behaviour so that all Advice classes are
+   * loaded and initialized as soon as a matching type is instrumented.
+   *
+   * <p>Note: this functionality currently does not work together with the AdviceTransformer.
+   * Therefore you should make your Advices indy-compatible (use @Advice.AssignReturned) before
+   * using this feature.
+   *
+   * @return true, if Advice classes should be loaded on instrumentation instead of first execution
+   */
+  default boolean loadAdviceClassesEagerly() {
+    return false;
+  }
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/AdviceTransformer.java

@@ -67,6 +67,9 @@ static byte[] transform(byte[] bytes) {
             }));
 
     TransformationContext context = new TransformationContext();
+    if (justDelegateAdvice) {
+      context.disableReturnTypeChange();
+    }
     ClassVisitor cv =
         new ClassVisitor(AsmApi.VERSION, cw) {
 

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyBootstrap.java

@@ -92,6 +92,12 @@ public class IndyBootstrap {
 
       IndyBootstrapDispatcher.init(
           MethodHandles.lookup().findStatic(IndyBootstrap.class, "bootstrap", bootstrapMethodType));
+
+      // Ensure that CallDepth is already loaded in case of bootstrapAdvice recursions with
+      // ClassLoader.loadClass
+      // This is required because CallDepth is a bootstrap class and therefore triggers our
+      // ClassLoader.loadClass instrumentations
+      Class.forName(CallDepth.class.getName());
     } catch (Exception e) {
       throw new IllegalStateException(e);
     }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java

@@ -25,6 +25,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -52,6 +54,9 @@
  */
 public class InstrumentationModuleClassLoader extends ClassLoader {
 
+  private static final Logger logger =
+      Logger.getLogger(InstrumentationModuleClassLoader.class.getName());
+
   static {
     ClassLoader.registerAsParallelCapable();
   }
@@ -150,8 +155,24 @@ public synchronized void installModule(InstrumentationModule module) {
                     className -> BytecodeWithUrl.create(className, agentOrExtensionCl)));
     installInjectedClasses(classesToInject);
     if (module instanceof ExperimentalInstrumentationModule) {
-      hiddenAgentPackages.addAll(
-          ((ExperimentalInstrumentationModule) module).agentPackagesToHide());
+      ExperimentalInstrumentationModule experimentalModule =
+          (ExperimentalInstrumentationModule) module;
+      hiddenAgentPackages.addAll(experimentalModule.agentPackagesToHide());
+      if (experimentalModule.loadAdviceClassesEagerly()) {
+        for (String adviceClass : getModuleAdviceNames(module)) {
+          try {
+            this.loadClass(adviceClass, true);
+          } catch (ClassNotFoundException e) {
+            logger.log(
+                Level.SEVERE,
+                "Failed to eagerly load advice class {0}",
+                new Object[] {adviceClass, e});
+          }
+        }
+        // We also eagerly load the LookupExposer, because that is also required for invokedynamic
+        // bootstrapping
+        getLookup();
+      }
     }
   }
 
@@ -196,6 +217,15 @@ public void applyTransformer(AgentBuilder.Transformer transformer) {}
 
   public static final Map<String, byte[]> bytecodeOverride = new ConcurrentHashMap<>();
 
+  @Override
+  public Class<?> loadClass(String name) throws ClassNotFoundException {
+    // We explicitly override loadClass from ClassLoader to ensure
+    // that loadClass is properly excluded from our internal ClassLoader Instrumentations
+    // (e.g. LoadInjectedClassInstrumentation, BooDelegationInstrumentation)
+    // Otherwise this will cause recursion in invokedynamic linkage
+    return loadClass(name, false);
+  }
+
   @Override
   @SuppressWarnings("removal") // AccessController is deprecated for removal
   protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {

testing-common/integration-tests/build.gradle.kts

@@ -8,6 +8,7 @@ dependencies {
 
   testCompileOnly(project(":instrumentation-api"))
   testCompileOnly(project(":javaagent-tooling"))
+  testCompileOnly(project(":javaagent-bootstrap"))
   testCompileOnly(project(":javaagent-extension-api"))
   testCompileOnly(project(":muzzle"))