Fix problems with lambdas on recursive path

Author: JonasKunz

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

@@ -52,10 +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.AssignReturned.ToReturned
+    public static Class<?> onExit(
         @Advice.Enter DefineClassContext context, @Advice.Return Class<?> returned) {
       DefineClassHelper.afterDefineClass(context);
+      return returned;
     }
   }
 
@@ -69,10 +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.AssignReturned.ToReturned
+    public static Class<?> onExit(
         @Advice.Enter DefineClassContext context, @Advice.Return Class<?> returned) {
       DefineClassHelper.afterDefineClass(context);
+      return returned;
     }
   }
 }

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

@@ -8,6 +8,7 @@
 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;
 
@@ -20,13 +21,31 @@ class AdviceBootstrapState implements AutoCloseable {
   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;
   }
 
-  public static AdviceBootstrapState enter(
+  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,
@@ -39,8 +58,7 @@ public static AdviceBootstrapState enter(
             adviceClassName,
             adviceMethodName,
             adviceMethodDescriptor);
-    AdviceBootstrapState state =
-        stateForCurrentThread.get().computeIfAbsent(key, k -> new AdviceBootstrapState(key));
+    AdviceBootstrapState state = stateForCurrentThread.get().computeIfAbsent(key, CONSTRUCTOR);
     state.recursionDepth++;
     return state;
   }
@@ -56,6 +74,13 @@ public MutableCallSite getOrInitMutableCallSite(Supplier<MutableCallSite> initia
     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;
@@ -99,8 +124,12 @@ private Key(
 
     @Override
     public boolean equals(Object o) {
-      if (this == o) return true;
-      if (o == null || getClass() != o.getClass()) return false;
+      if (this == o) {
+        return true;
+      }
+      if (o == null || !(o instanceof Key)) {
+        return false;
+      }
 
       Key that = (Key) o;
       return instrumentedClass.equals(that.instrumentedClass)

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

@@ -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);
     }
@@ -176,20 +178,12 @@ private static CallSite bootstrapAdvice(
         // 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.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});
-        if (logger.isLoggable(Level.FINEST)) {
-          logger.log(
-              Level.FINEST,
-              "Stacktrace for nested invokedynamic instruction linkage:",
-              new Throwable());
+        MutableCallSite mutableCallSite = nestedState.getMutableCallSite();
+        if (mutableCallSite == null) {
+          mutableCallSite = new MutableCallSite(IndyBootstrapDispatcher.generateNoopMethodHandle(invokedynamicMethodType));
+          nestedState.initMutableCallSite(mutableCallSite);
         }
-        return nestedState.getOrInitMutableCallSite(
-            () ->
-                new MutableCallSite(
-                    IndyBootstrapDispatcher.generateNoopMethodHandle(invokedynamicMethodType)));
+        return mutableCallSite;
       }
 
       InstrumentationModuleClassLoader instrumentationClassloader =
@@ -212,6 +206,10 @@ private static CallSite bootstrapAdvice(
       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 instrumentaiton should be active now",
+            new Object[] {lookup.lookupClass().getName(), adviceClassName, adviceMethodName});
         nestedBootstrapCallSite.setTarget(methodHandle);
         MutableCallSite.syncAll(new MutableCallSite[] {nestedBootstrapCallSite});
         return nestedBootstrapCallSite;

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

@@ -25,7 +25,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -53,9 +52,6 @@
  */
 public class InstrumentationModuleClassLoader extends ClassLoader {
 
-  private static final Logger logger =
-      Logger.getLogger(InstrumentationModuleClassLoader.class.getName());
-
   static {
     ClassLoader.registerAsParallelCapable();
   }