Make internal classloader instrumentation indy compatible (#12242)

Co-authored-by: SylvainJuge <763082+SylvainJuge@users.noreply.github.com>
Co-authored-by: Trask Stalnaker <trask.stalnaker@gmail.com>

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 class BootDelegationInstrumentation implements TypeInstrumentation {
       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

@@ -25,13 +25,10 @@ public class ClassLoaderInstrumentationModule extends InstrumentationModule {
     return true;
   }
 
-  @Override
-  public boolean isIndyModule() {
-    return false;
-  }
-
   @Override
   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");
   }
 

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

@@ -52,9 +52,15 @@ public class DefineClassInstrumentation implements TypeInstrumentation {
           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 because we do not want the return values to be wrapped in array types
     @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 +74,15 @@ public class DefineClassInstrumentation implements TypeInstrumentation {
       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 because we do not want the return values to be wrapped in array types
     @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 class LoadInjectedClassInstrumentation implements TypeInstrumentation {
     }
 
     @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

@@ -61,18 +61,18 @@ public class ResourceInjectionInstrumentation implements TypeInstrumentation {
   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;
     }
   }
 
@@ -80,18 +80,18 @@ public class ResourceInjectionInstrumentation implements TypeInstrumentation {
   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,8 +108,7 @@ public class ResourceInjectionInstrumentation implements TypeInstrumentation {
           result.add(helperUrl);
         }
       }
-
-      resources = Collections.enumeration(result);
+      return Collections.enumeration(result);
     }
   }
 
@@ -117,22 +116,22 @@ public class ResourceInjectionInstrumentation implements TypeInstrumentation {
   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/IndyBootstrapDispatcher.java

@@ -54,8 +54,7 @@ public class IndyBootstrapDispatcher {
     return callSite;
   }
 
-  // package visibility for testing
-  static MethodHandle generateNoopMethodHandle(MethodType methodType) {
+  public static MethodHandle generateNoopMethodHandle(MethodType methodType) {
     Class<?> returnType = methodType.returnType();
     MethodHandle noopNoArg;
     if (returnType == void.class) {

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

@@ -0,0 +1,152 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.instrumentation.indy;
+
+import java.lang.invoke.MutableCallSite;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+class AdviceBootstrapState implements AutoCloseable {
+
+  private static final ThreadLocal<Map<Key, AdviceBootstrapState>> stateForCurrentThread =
+      ThreadLocal.withInitial(HashMap::new);
+
+  private final Key key;
+  private int recursionDepth;
+  @Nullable private MutableCallSite nestedCallSite;
+
+  /**
+   * We have to eagerly initialize to not cause a lambda construction during {@link #enter(Class,
+   * String, String, String, String)}.
+   */
+  private static final Function<Key, AdviceBootstrapState> CONSTRUCTOR = AdviceBootstrapState::new;
+
+  private AdviceBootstrapState(Key key) {
+    this.key = key;
+    // enter will increment it by one, so 0 is the value for non-recursive calls
+    recursionDepth = -1;
+  }
+
+  static void initialize() {
+    // Eager initialize everything because we could run into recursions doing this during advice
+    // bootstrapping
+    stateForCurrentThread.get();
+    stateForCurrentThread.remove();
+    try {
+      Class.forName(Key.class.getName());
+    } catch (ClassNotFoundException e) {
+      throw new IllegalStateException(e);
+    }
+  }
+
+  static AdviceBootstrapState enter(
+      Class<?> instrumentedClass,
+      String moduleClassName,
+      String adviceClassName,
+      String adviceMethodName,
+      String adviceMethodDescriptor) {
+    Key key =
+        new Key(
+            instrumentedClass,
+            moduleClassName,
+            adviceClassName,
+            adviceMethodName,
+            adviceMethodDescriptor);
+    AdviceBootstrapState state = stateForCurrentThread.get().computeIfAbsent(key, CONSTRUCTOR);
+    state.recursionDepth++;
+    return state;
+  }
+
+  public boolean isNestedInvocation() {
+    return recursionDepth > 0;
+  }
+
+  public MutableCallSite getOrInitMutableCallSite(Supplier<MutableCallSite> initializer) {
+    if (nestedCallSite == null) {
+      nestedCallSite = initializer.get();
+    }
+    return nestedCallSite;
+  }
+
+  public void initMutableCallSite(MutableCallSite mutableCallSite) {
+    if (nestedCallSite != null) {
+      throw new IllegalStateException("callsite has already been initialized");
+    }
+    nestedCallSite = mutableCallSite;
+  }
+
+  @Nullable
+  public MutableCallSite getMutableCallSite() {
+    return nestedCallSite;
+  }
+
+  @Override
+  public void close() {
+    if (recursionDepth == 0) {
+      Map<Key, AdviceBootstrapState> stateMap = stateForCurrentThread.get();
+      stateMap.remove(key);
+      if (stateMap.isEmpty()) {
+        // Do not leave an empty map dangling as thread local
+        stateForCurrentThread.remove();
+      }
+    } else {
+      recursionDepth--;
+    }
+  }
+
+  /** Key uniquely identifying a single invokedynamic instruction inserted for an advice */
+  private static class Key {
+
+    private final Class<?> instrumentedClass;
+    private final String moduleClassName;
+    private final String adviceClassName;
+    private final String adviceMethodName;
+    private final String adviceMethodDescriptor;
+
+    private Key(
+        Class<?> instrumentedClass,
+        String moduleClassName,
+        String adviceClassName,
+        String adviceMethodName,
+        String adviceMethodDescriptor) {
+      this.instrumentedClass = instrumentedClass;
+      this.moduleClassName = moduleClassName;
+      this.adviceClassName = adviceClassName;
+      this.adviceMethodName = adviceMethodName;
+      this.adviceMethodDescriptor = adviceMethodDescriptor;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || !(o instanceof Key)) {
+        return false;
+      }
+
+      Key that = (Key) o;
+      return instrumentedClass.equals(that.instrumentedClass)
+          && moduleClassName.equals(that.moduleClassName)
+          && adviceClassName.equals(that.adviceClassName)
+          && adviceMethodName.equals(that.adviceMethodName)
+          && adviceMethodDescriptor.equals(that.adviceMethodDescriptor);
+    }
+
+    @Override
+    public int hashCode() {
+      int result = instrumentedClass.hashCode();
+      result = 31 * result + moduleClassName.hashCode();
+      result = 31 * result + adviceClassName.hashCode();
+      result = 31 * result + adviceMethodName.hashCode();
+      result = 31 * result + adviceMethodDescriptor.hashCode();
+      return result;
+    }
+  }
+}

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

@@ -67,6 +67,9 @@ class AdviceTransformer {
             }));
 
     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

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
-import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.IndyBootstrapDispatcher;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import java.lang.invoke.CallSite;
@@ -13,6 +12,7 @@ import java.lang.invoke.ConstantCallSite;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.lang.invoke.MutableCallSite;
 import java.lang.reflect.Method;
 import java.security.PrivilegedAction;
 import java.util.Arrays;
@@ -84,7 +84,7 @@ public class IndyBootstrap {
 
       MethodType bootstrapMethodType =
           MethodType.methodType(
-              ConstantCallSite.class,
+              CallSite.class,
               MethodHandles.Lookup.class,
               String.class,
               MethodType.class,
@@ -92,6 +92,8 @@ public class IndyBootstrap {
 
       IndyBootstrapDispatcher.init(
           MethodHandles.lookup().findStatic(IndyBootstrap.class, "bootstrap", bootstrapMethodType));
+
+      AdviceBootstrapState.initialize();
     } catch (Exception e) {
       throw new IllegalStateException(e);
     }
@@ -105,7 +107,7 @@ public class IndyBootstrap {
 
   @Nullable
   @SuppressWarnings({"unused", "removal"}) // SecurityManager and AccessController are deprecated
-  private static ConstantCallSite bootstrap(
+  private static CallSite bootstrap(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
       MethodType adviceMethodType,
@@ -118,11 +120,11 @@ public class IndyBootstrap {
     // callsite resolution needs privileged access to call Class#getClassLoader() and
     // MethodHandles$Lookup#findStatic
     return java.security.AccessController.doPrivileged(
-        (PrivilegedAction<ConstantCallSite>)
+        (PrivilegedAction<CallSite>)
             () -> internalBootstrap(lookup, adviceMethodName, adviceMethodType, args));
   }
 
-  private static ConstantCallSite internalBootstrap(
+  private static CallSite internalBootstrap(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
       MethodType adviceMethodType,
@@ -157,7 +159,7 @@ public class IndyBootstrap {
     }
   }
 
-  private static ConstantCallSite bootstrapAdvice(
+  private static CallSite bootstrapAdvice(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
       MethodType invokedynamicMethodType,
@@ -165,17 +167,25 @@ public class IndyBootstrap {
       String adviceMethodDescriptor,
       String adviceClassName)
       throws NoSuchMethodException, IllegalAccessException, ClassNotFoundException {
-    CallDepth callDepth = CallDepth.forClass(IndyBootstrap.class);
-    try {
-      if (callDepth.getAndIncrement() > 0) {
+    try (AdviceBootstrapState nestedState =
+        AdviceBootstrapState.enter(
+            lookup.lookupClass(),
+            moduleClassName,
+            adviceClassName,
+            adviceMethodName,
+            adviceMethodDescriptor)) {
+      if (nestedState.isNestedInvocation()) {
         // avoid re-entrancy and stack overflow errors, which may happen when bootstrapping an
         // instrumentation that also gets triggered during the bootstrap
         // for example, adding correlation ids to the thread context when executing logger.debug.
-        logger.log(
-            Level.WARNING,
-            "Nested instrumented invokedynamic instruction linkage detected",
-            new Throwable());
-        return null;
+        MutableCallSite mutableCallSite = nestedState.getMutableCallSite();
+        if (mutableCallSite == null) {
+          mutableCallSite =
+              new MutableCallSite(
+                  IndyBootstrapDispatcher.generateNoopMethodHandle(invokedynamicMethodType));
+          nestedState.initMutableCallSite(mutableCallSite);
+        }
+        return mutableCallSite;
       }
 
       InstrumentationModuleClassLoader instrumentationClassloader =
@@ -193,9 +203,21 @@ public class IndyBootstrap {
               .getLookup()
               .findStatic(adviceClass, adviceMethodName, actualAdviceMethodType)
               .asType(invokedynamicMethodType);
-      return new ConstantCallSite(methodHandle);
-    } finally {
-      callDepth.decrementAndGet();
+
+      MutableCallSite nestedBootstrapCallSite = nestedState.getMutableCallSite();
+      if (nestedBootstrapCallSite != null) {
+        // There have been nested bootstrapping attempts
+        // Update the callsite of those to run the actual instrumentation
+        logger.log(
+            Level.FINE,
+            "Fixing nested instrumentation invokedynamic instruction bootstrapping for instrumented class {0} and advice {1}.{2}, the instrumentation should be active now",
+            new Object[] {lookup.lookupClass().getName(), adviceClassName, adviceMethodName});
+        nestedBootstrapCallSite.setTarget(methodHandle);
+        MutableCallSite.syncAll(new MutableCallSite[] {nestedBootstrapCallSite});
+        return nestedBootstrapCallSite;
+      } else {
+        return new ConstantCallSite(methodHandle);
+      }
     }
   }
 

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

@@ -196,6 +196,15 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
 
   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 {
@@ -317,12 +326,23 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
   }
 
   private Class<?> defineClassWithPackage(String name, byte[] bytecode) {
-    int lastDotIndex = name.lastIndexOf('.');
-    if (lastDotIndex != -1) {
-      String packageName = name.substring(0, lastDotIndex);
-      safeDefinePackage(packageName);
+    try {
+      int lastDotIndex = name.lastIndexOf('.');
+      if (lastDotIndex != -1) {
+        String packageName = name.substring(0, lastDotIndex);
+        safeDefinePackage(packageName);
+      }
+      return defineClass(name, bytecode, 0, bytecode.length, PROTECTION_DOMAIN);
+    } catch (LinkageError error) {
+      // Precaution against linkage error due to nested instrumentations happening
+      // it might be possible that e.g. an advice class has already been defined
+      // during an instrumentation of defineClass
+      Class<?> clazz = findLoadedClass(name);
+      if (clazz != null) {
+        return clazz;
+      }
+      throw error;
     }
-    return defineClass(name, bytecode, 0, bytecode.length, PROTECTION_DOMAIN);
   }
 
   private void safeDefinePackage(String packageName) {

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"))
 

testing-common/integration-tests/src/test/groovy/AgentInstrumentationSpecificationTest.groovy

@@ -6,7 +6,7 @@
 import com.google.common.reflect.ClassPath
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.instrumentation.test.utils.ClasspathUtils
-import io.opentelemetry.javaagent.tooling.Constants
+import io.opentelemetry.javaagent.bootstrap.BootstrapPackagePrefixesHolder
 import org.slf4j.LoggerFactory
 
 import java.util.concurrent.TimeoutException
@@ -17,12 +17,14 @@ import java.util.concurrent.TimeoutException
 class AgentInstrumentationSpecificationTest extends AgentInstrumentationSpecification {
   private static final ClassLoader BOOTSTRAP_CLASSLOADER = null
 
+  public static final List<String> BOOTSTRAP_PACKAGE_PREFIXES = BootstrapPackagePrefixesHolder.getBoostrapPackagePrefixes()
+
   def "classpath setup"() {
     setup:
     final List<String> bootstrapClassesIncorrectlyLoaded = []
     for (ClassPath.ClassInfo info : getTestClasspath().getAllClasses()) {
-      for (int i = 0; i < Constants.BOOTSTRAP_PACKAGE_PREFIXES.size(); ++i) {
-        if (info.getName().startsWith(Constants.BOOTSTRAP_PACKAGE_PREFIXES[i])) {
+      for (int i = 0; i < BOOTSTRAP_PACKAGE_PREFIXES.size(); ++i) {
+        if (info.getName().startsWith(BOOTSTRAP_PACKAGE_PREFIXES[i])) {
           Class<?> bootstrapClass = Class.forName(info.getName())
           def loader
           try {