Added support for nested invokedynamic bootstrapping using MutableCallSite

Author: JonasKunz

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

@@ -55,7 +55,7 @@ public static CallSite bootstrap(
   }
 
   // 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,123 @@
+/*
+ * 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.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;
+
+  private AdviceBootstrapState(Key key) {
+    this.key = key;
+    // enter will increment it by one, so 0 is the value for non-recursive calls
+    recursionDepth = -1;
+  }
+
+  public 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, k -> new AdviceBootstrapState(key));
+    state.recursionDepth++;
+    return state;
+  }
+
+  public boolean isNestedInvocation() {
+    return recursionDepth > 0;
+  }
+
+  public MutableCallSite getOrInitMutableCallSite(Supplier<MutableCallSite> initializer) {
+    if (nestedCallSite == null) {
+      nestedCallSite = initializer.get();
+    }
+    return nestedCallSite;
+  }
+
+  @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 || getClass() != o.getClass()) 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/IndyBootstrap.java

@@ -5,14 +5,14 @@
 
 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;
 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,20 +84,14 @@ public class IndyBootstrap {
 
       MethodType bootstrapMethodType =
           MethodType.methodType(
-              ConstantCallSite.class,
+              CallSite.class,
               MethodHandles.Lookup.class,
               String.class,
               MethodType.class,
               Object[].class);
 
       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);
     }
@@ -111,7 +105,7 @@ public static Method getIndyBootstrapMethod() {
 
   @Nullable
   @SuppressWarnings({"unused", "removal"}) // SecurityManager and AccessController are deprecated
-  private static ConstantCallSite bootstrap(
+  private static CallSite bootstrap(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
       MethodType adviceMethodType,
@@ -124,11 +118,11 @@ private static ConstantCallSite bootstrap(
     // 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,
@@ -163,29 +157,39 @@ private static ConstantCallSite internalBootstrap(
     }
   }
 
-  private static ConstantCallSite bootstrapAdvice(
+  private static CallSite bootstrapAdvice(
       MethodHandles.Lookup lookup,
       String adviceMethodName,
       MethodType invokedynamicMethodType,
       String moduleClassName,
       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.SEVERE,
-            "Nested instrumented invokedynamic instruction linkage detected for instrumented class {0} and advice {1}.{2}, the instrumentation might not work correctly",
+            Level.FINE,
+            "Nested instrumented invokedynamic instruction linkage detected for instrumented class {0} and advice {1}.{2}, this invocation won't be instrumented",
             new Object[] {lookup.lookupClass().getName(), adviceClassName, adviceMethodName});
-        logger.log(
-            Level.SEVERE,
-            "Stacktrace for nested invokedynamic instruction linkage:",
-            new Throwable());
-        return null;
+        if (logger.isLoggable(Level.FINEST)) {
+          logger.log(
+              Level.FINEST,
+              "Stacktrace for nested invokedynamic instruction linkage:",
+              new Throwable());
+        }
+        return nestedState.getOrInitMutableCallSite(
+            () ->
+                new MutableCallSite(
+                    IndyBootstrapDispatcher.generateNoopMethodHandle(invokedynamicMethodType)));
       }
 
       InstrumentationModuleClassLoader instrumentationClassloader =
@@ -203,9 +207,17 @@ private static ConstantCallSite bootstrapAdvice(
               .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
+        nestedBootstrapCallSite.setTarget(methodHandle);
+        MutableCallSite.syncAll(new MutableCallSite[] {nestedBootstrapCallSite});
+        return nestedBootstrapCallSite;
+      } else {
+        return new ConstantCallSite(methodHandle);
+      }
     }
   }