Define helper classes in loadClass (#5528)

* Define helper classes in loadClass similarly to regular classes

* fix test

* spotless

* address review comments

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

@@ -30,17 +30,17 @@ import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.slf4j.LoggerFactory;
 
-/*
- * Some class loaders do not delegate to their parent, so classes in those class loaders
- * will not be able to see classes in the bootstrap class loader.
+/**
+ * Some class loaders do not delegate to their parent, so classes in those class loaders will not be
+ * able to see classes in the bootstrap class loader.
  *
- * In particular, instrumentation on classes in those class loaders will not be able to see
- * the shaded OpenTelemetry API classes in the bootstrap class loader.
+ * <p>In particular, instrumentation on classes in those class loaders will not be able to see the
+ * shaded OpenTelemetry API classes in the bootstrap class loader.
  *
- * This instrumentation forces all class loaders to delegate to the bootstrap class loader
- * for the classes that we have put in the bootstrap class loader.
+ * <p>This instrumentation forces all class loaders to delegate to the bootstrap class loader for
+ * the classes that we have put in the bootstrap class loader.
  */
-public class ClassLoaderInstrumentation implements TypeInstrumentation {
+public class BootDelegationInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
@@ -67,7 +67,7 @@ public class ClassLoaderInstrumentation implements TypeInstrumentation {
                             .and(takesArgument(1, boolean.class))))
             .and(isPublic().or(isProtected()))
             .and(not(isStatic())),
-        ClassLoaderInstrumentation.class.getName() + "$LoadClassAdvice");
+        BootDelegationInstrumentation.class.getName() + "$LoadClassAdvice");
   }
 
   public static class Holder {

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

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.instrumentation.internal.classloader;
 
 import static java.util.Arrays.asList;
-import static java.util.Collections.singletonList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -30,17 +29,11 @@ public class ClassLoaderInstrumentationModule extends InstrumentationModule {
     return className.equals("io.opentelemetry.javaagent.tooling.Constants");
   }
 
-  @Override
-  public List<String> getAdditionalHelperClassNames() {
-    return singletonList(
-        "io.opentelemetry.javaagent.instrumentation.internal.classloader.DefineClassUtil");
-  }
-
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
-        new ClassLoaderInstrumentation(),
-        new ResourceInjectionInstrumentation(),
-        new DefineClassInstrumentation());
+        new BootDelegationInstrumentation(),
+        new LoadInjectedClassInstrumentation(),
+        new ResourceInjectionInstrumentation());
   }
 }

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

@@ -1,187 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.javaagent.instrumentation.internal.classloader;
-
-import static net.bytebuddy.matcher.ElementMatchers.named;
-
-import io.opentelemetry.javaagent.bootstrap.DefineClassContext;
-import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
-import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import net.bytebuddy.asm.AsmVisitorWrapper;
-import net.bytebuddy.description.field.FieldDescription;
-import net.bytebuddy.description.field.FieldList;
-import net.bytebuddy.description.method.MethodList;
-import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.implementation.Implementation;
-import net.bytebuddy.matcher.ElementMatcher;
-import net.bytebuddy.pool.TypePool;
-import org.objectweb.asm.ClassVisitor;
-import org.objectweb.asm.ClassWriter;
-import org.objectweb.asm.Label;
-import org.objectweb.asm.MethodVisitor;
-import org.objectweb.asm.Opcodes;
-import org.objectweb.asm.Type;
-
-public class DefineClassInstrumentation implements TypeInstrumentation {
-
-  @Override
-  public ElementMatcher<TypeDescription> typeMatcher() {
-    return named("java.lang.ClassLoader");
-  }
-
-  @Override
-  public void transform(TypeTransformer transformer) {
-    transformer.applyTransformer(
-        (builder, typeDescription, classLoader, module) ->
-            builder.visit(
-                new AsmVisitorWrapper() {
-                  @Override
-                  public int mergeWriter(int flags) {
-                    return flags | ClassWriter.COMPUTE_MAXS;
-                  }
-
-                  @Override
-                  public int mergeReader(int flags) {
-                    return flags;
-                  }
-
-                  @Override
-                  public ClassVisitor wrap(
-                      TypeDescription instrumentedType,
-                      ClassVisitor classVisitor,
-                      Implementation.Context implementationContext,
-                      TypePool typePool,
-                      FieldList<FieldDescription.InDefinedShape> fields,
-                      MethodList<?> methods,
-                      int writerFlags,
-                      int readerFlags) {
-                    return new ClassLoaderClassVisitor(classVisitor);
-                  }
-                }));
-  }
-
-  private static class ClassLoaderClassVisitor extends ClassVisitor {
-
-    ClassLoaderClassVisitor(ClassVisitor cv) {
-      super(Opcodes.ASM7, cv);
-    }
-
-    @Override
-    public MethodVisitor visitMethod(
-        int access, String name, String descriptor, String signature, String[] exceptions) {
-      MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
-      // apply the following transformation to defineClass method
-      /*
-      DefineClassContext.enter();
-      try {
-        // original method body
-        DefineClassContext.exit();
-        return result;
-      } catch (LinkageError error) {
-        boolean helpersInjected = DefineClassContext.exitAndGet();
-        Class<?> loaded = findLoadedClass(className);
-        return DefineClassUtil.handleLinkageError(error, helpersInjected, loaded);
-      } catch (Throwable throwable) {
-        DefineClassContext.exit();
-        throw throwable;
-      }
-       */
-      if ("defineClass".equals(name)
-          && ("(Ljava/lang/String;[BIILjava/security/ProtectionDomain;)Ljava/lang/Class;"
-                  .equals(descriptor)
-              || "(Ljava/lang/String;Ljava/nio/ByteBuffer;Ljava/security/ProtectionDomain;)Ljava/lang/Class;"
-                  .equals(descriptor))) {
-        mv =
-            new MethodVisitor(api, mv) {
-              Label start = new Label();
-              Label end = new Label();
-              Label handler = new Label();
-
-              @Override
-              public void visitCode() {
-                mv.visitMethodInsn(
-                    Opcodes.INVOKESTATIC,
-                    Type.getInternalName(DefineClassContext.class),
-                    "enter",
-                    "()V",
-                    false);
-                mv.visitTryCatchBlock(start, end, end, "java/lang/LinkageError");
-                // catch other exceptions
-                mv.visitTryCatchBlock(start, end, handler, null);
-                mv.visitLabel(start);
-
-                super.visitCode();
-              }
-
-              @Override
-              public void visitInsn(int opcode) {
-                if (opcode == Opcodes.ARETURN) {
-                  mv.visitMethodInsn(
-                      Opcodes.INVOKESTATIC,
-                      Type.getInternalName(DefineClassContext.class),
-                      "exit",
-                      "()V",
-                      false);
-                }
-                super.visitInsn(opcode);
-              }
-
-              @Override
-              public void visitMaxs(int maxStack, int maxLocals) {
-                // handle LinkageError
-                mv.visitLabel(end);
-                mv.visitFrame(
-                    Opcodes.F_FULL,
-                    2,
-                    new Object[] {"java/lang/ClassLoader", "java/lang/String"},
-                    1,
-                    new Object[] {"java/lang/LinkageError"});
-                mv.visitMethodInsn(
-                    Opcodes.INVOKESTATIC,
-                    Type.getInternalName(DefineClassContext.class),
-                    "exitAndGet",
-                    "()Z",
-                    false);
-                mv.visitVarInsn(Opcodes.ALOAD, 0);
-                mv.visitVarInsn(Opcodes.ALOAD, 1);
-                mv.visitMethodInsn(
-                    Opcodes.INVOKEVIRTUAL,
-                    "java/lang/ClassLoader",
-                    "findLoadedClass",
-                    "(Ljava/lang/String;)Ljava/lang/Class;",
-                    false);
-                mv.visitMethodInsn(
-                    Opcodes.INVOKESTATIC,
-                    Type.getInternalName(DefineClassUtil.class),
-                    "handleLinkageError",
-                    "(Ljava/lang/LinkageError;ZLjava/lang/Class;)Ljava/lang/Class;",
-                    false);
-                mv.visitInsn(Opcodes.ARETURN);
-
-                // handle Throwable
-                mv.visitLabel(handler);
-                mv.visitFrame(
-                    Opcodes.F_FULL,
-                    2,
-                    new Object[] {"java/lang/ClassLoader", "java/lang/String"},
-                    1,
-                    new Object[] {"java/lang/Throwable"});
-                mv.visitMethodInsn(
-                    Opcodes.INVOKESTATIC,
-                    Type.getInternalName(DefineClassContext.class),
-                    "exit",
-                    "()V",
-                    false);
-                mv.visitInsn(Opcodes.ATHROW);
-
-                super.visitMaxs(maxStack, maxLocals);
-              }
-            };
-      }
-      return mv;
-    }
-  }
-}

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

@@ -1,41 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.javaagent.instrumentation.internal.classloader;
-
-@SuppressWarnings("unused")
-public final class DefineClassUtil {
-  private DefineClassUtil() {}
-
-  /**
-   * Handle LinkageError in ClassLoader.defineClass. Call to this method is inserted into
-   * ClassLoader.defineClass by DefineClassInstrumentation.
-   *
-   * @param linkageError LinkageError that happened in defineClass
-   * @param helpersInjected whether helpers were injected during defineClass call
-   * @param clazz Class that is being defined if it is already loaded
-   * @return give Class if LinkageError was a duplicate class definition error
-   */
-  public static Class<?> handleLinkageError(
-      LinkageError linkageError, boolean helpersInjected, Class<?> clazz) {
-    // only attempt to recover from duplicate class definition if helpers were injected during
-    // the defineClass call
-    if (!helpersInjected
-        // if exception was duplicate class definition we'll have access to the loaded class
-        || clazz == null
-        // duplicate class definition throws LinkageError, we can ignore its subclasses
-        || linkageError.getClass() != LinkageError.class) {
-      throw linkageError;
-    }
-    // check that the exception is a duplicate class or interface definition
-    String message = linkageError.getMessage();
-    if (message == null
-        || !(message.contains("duplicate interface definition")
-            || message.contains("duplicate class definition"))) {
-      throw linkageError;
-    }
-    return clazz;
-  }
-}

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isProtected;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * This instrumentation inserts loading of our injected helper classes at the start of {@code
+ * ClassLoader.loadClass} method.
+ */
+public class LoadInjectedClassInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return extendsClass(named("java.lang.ClassLoader"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("loadClass"))
+            .and(
+                takesArguments(1)
+                    .and(takesArgument(0, String.class))
+                    .or(
+                        takesArguments(2)
+                            .and(takesArgument(0, String.class))
+                            .and(takesArgument(1, boolean.class))))
+            .and(isPublic().or(isProtected()))
+            .and(not(isStatic())),
+        LoadInjectedClassInstrumentation.class.getName() + "$LoadClassAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class LoadClassAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static Class<?> onEnter(
+        @Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) {
+      Class<?> helperClass = InjectedClassHelper.loadHelperClass(classLoader, name);
+      if (helperClass != null) {
+        return helperClass;
+      }
+
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Return(readOnly = false) Class<?> result, @Advice.Enter Class<?> loadedClass) {
+      if (loadedClass != null) {
+        result = loadedClass;
+      }
+    }
+  }
+}

instrumentation/internal/internal-lambda/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/lambda/LambdaTransformer.java

@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.internal.lambda;
 
 import io.opentelemetry.javaagent.bootstrap.ClassFileTransformerHolder;
-import io.opentelemetry.javaagent.bootstrap.InjectedHelperClassDetector;
+import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
 import java.lang.instrument.ClassFileTransformer;
 
 /** Helper class for transforming lambda class bytes. */
@@ -31,7 +31,7 @@ public final class LambdaTransformer {
   @SuppressWarnings("unused")
   public static byte[] transform(byte[] classBytes, String slashClassName, Class<?> targetClass) {
     // Skip transforming lambdas of injected helper classes.
-    if (InjectedHelperClassDetector.isHelperClass(targetClass)) {
+    if (InjectedClassHelper.isHelperClass(targetClass)) {
       return classBytes;
     }
 

instrumentation/spring/spring-web-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/springweb/WebApplicationContextInstrumentation.java

@@ -65,14 +65,16 @@ public class WebApplicationContextInstrumentation implements TypeInstrumentation
           // Firstly check whether DispatcherServlet is present. We need to load an instrumented
           // class from spring-webmvc to trigger injection that makes
           // OpenTelemetryHandlerMappingFilter available.
-          beanFactory
-              .getBeanClassLoader()
-              .loadClass("org.springframework.web.servlet.DispatcherServlet");
-
-          // Now attempt to load our injected instrumentation class.
-          Class<?> clazz =
+          Class<?> dispatcherServletClass =
               beanFactory
                   .getBeanClassLoader()
+                  .loadClass("org.springframework.web.servlet.DispatcherServlet");
+
+          // Now attempt to load our injected instrumentation class from the same class loader as
+          // DispatcherServlet
+          Class<?> clazz =
+              dispatcherServletClass
+                  .getClassLoader()
                   .loadClass("org.springframework.web.servlet.OpenTelemetryHandlerMappingFilter");
           GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
           beanDefinition.setScope(SCOPE_SINGLETON);

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

@@ -0,0 +1,57 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap;
+
+import java.util.function.BiFunction;
+import java.util.function.BiPredicate;
+import java.util.function.Function;
+
+/** Helper class for detecting and loading injected helper classes. */
+public final class InjectedClassHelper {
+
+  private InjectedClassHelper() {}
+
+  private static volatile BiPredicate<ClassLoader, String> helperClassDetector;
+
+  /** Sets the {@link Function} for detecting injected helper classes. */
+  public static void internalSetHelperClassDetector(
+      BiPredicate<ClassLoader, String> helperClassDetector) {
+    if (InjectedClassHelper.helperClassDetector != null) {
+      // Only possible by misuse of this API, just ignore.
+      return;
+    }
+    InjectedClassHelper.helperClassDetector = helperClassDetector;
+  }
+
+  public static boolean isHelperClass(Class<?> clazz) {
+    return isHelperClass(clazz.getClassLoader(), clazz.getName());
+  }
+
+  public static boolean isHelperClass(ClassLoader classLoader, String className) {
+    if (helperClassDetector == null) {
+      return false;
+    }
+    return helperClassDetector.test(classLoader, className);
+  }
+
+  private static volatile BiFunction<ClassLoader, String, Class<?>> helperClassLoader;
+
+  public static void internalSetHelperClassLoader(
+      BiFunction<ClassLoader, String, Class<?>> helperClassLoader) {
+    if (InjectedClassHelper.helperClassLoader != null) {
+      // Only possible by misuse of this API, just ignore.
+      return;
+    }
+    InjectedClassHelper.helperClassLoader = helperClassLoader;
+  }
+
+  public static Class<?> loadHelperClass(ClassLoader classLoader, String className) {
+    if (helperClassLoader == null) {
+      return null;
+    }
+    return helperClassLoader.apply(classLoader, className);
+  }
+}

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

@@ -1,33 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.javaagent.bootstrap;
-
-import java.util.function.Function;
-
-/** Helper class for detecting whether given class is an injected helper class. */
-public final class InjectedHelperClassDetector {
-
-  private InjectedHelperClassDetector() {}
-
-  private static volatile Function<Class<?>, Boolean> helperClassDetector;
-
-  /** Sets the {@link Function} for detecting injected helper classes. */
-  public static void internalSetHelperClassDetector(
-      Function<Class<?>, Boolean> helperClassDetector) {
-    if (InjectedHelperClassDetector.helperClassDetector != null) {
-      // Only possible by misuse of this API, just ignore.
-      return;
-    }
-    InjectedHelperClassDetector.helperClassDetector = helperClassDetector;
-  }
-
-  public static boolean isHelperClass(Class<?> clazz) {
-    if (helperClassDetector == null) {
-      return false;
-    }
-    return helperClassDetector.apply(clazz);
-  }
-}

javaagent-tooling/javaagent-tooling-java9/src/main/java/io/opentelemetry/javaagent/tooling/ExposeAgentBootstrapListener.java

@@ -16,7 +16,10 @@ import net.bytebuddy.utility.JavaModule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-/** Ensures that transformed classes can read agent classes in bootstrap class loader. */
+/**
+ * Ensures that transformed classes can read agent classes in bootstrap class loader and injected
+ * classes in unnamed module of their class loader.
+ */
 public class ExposeAgentBootstrapListener extends AgentBuilder.Listener.Adapter {
   private static final Logger logger = LoggerFactory.getLogger(ExposeAgentBootstrapListener.class);
   // unnamed module in bootstrap class loader
@@ -29,7 +32,6 @@ public class ExposeAgentBootstrapListener extends AgentBuilder.Listener.Adapter
     this.instrumentation = instrumentation;
   }
 
-  @SuppressWarnings("ReferenceEquality")
   @Override
   public void onTransformation(
       TypeDescription typeDescription,
@@ -37,14 +39,24 @@ public class ExposeAgentBootstrapListener extends AgentBuilder.Listener.Adapter
       JavaModule javaModule,
       boolean b,
       DynamicType dynamicType) {
-    if (javaModule != JavaModule.UNSUPPORTED
-        && javaModule.isNamed()
-        && !javaModule.canRead(agentBootstrapModule)) {
-      logger.debug("Adding module read from {} to {}", javaModule, agentBootstrapModule);
+    // expose agent classes in unnamed module of bootstrap class loader
+    exposeModule(javaModule, agentBootstrapModule);
+    if (classLoader != null) {
+      // expose classes in unnamed module of current class loader
+      // this is needed so that advice code can access injected helper classes
+      exposeModule(javaModule, JavaModule.of(classLoader.getUnnamedModule()));
+    }
+  }
+
+  private void exposeModule(JavaModule fromModule, JavaModule targetModule) {
+    if (fromModule != JavaModule.UNSUPPORTED
+        && fromModule.isNamed()
+        && !fromModule.canRead(targetModule)) {
+      logger.debug("Adding module read from {} to {}", fromModule, targetModule);
       ClassInjector.UsingInstrumentation.redefineModule(
           instrumentation,
-          javaModule,
-          Collections.singleton(agentBootstrapModule),
+          fromModule,
+          Collections.singleton(targetModule),
           Collections.emptyMap(),
           Collections.emptyMap(),
           Collections.emptySet(),

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

@@ -217,7 +217,11 @@ public class AgentInstaller {
 
     return agentBuilder
         .ignore(any(), new IgnoredClassLoadersMatcher(builder.buildIgnoredClassLoadersTrie()))
-        .or(new IgnoredTypesMatcher(builder.buildIgnoredTypesTrie()));
+        .or(new IgnoredTypesMatcher(builder.buildIgnoredTypesTrie()))
+        .or(
+            (typeDescription, classLoader, module, classBeingRedefined, protectionDomain) -> {
+              return HelperInjector.isInjectedClass(classLoader, typeDescription.getName());
+            });
   }
 
   private static void runAfterAgentListeners(

javaagent-tooling/src/test/groovy/io/opentelemetry/javaagent/test/HelperInjectionTest.groovy

@@ -39,6 +39,7 @@ class HelperInjectionTest extends Specification {
 
     when:
     injector.transform(null, null, emptyLoader.get(), null)
+    HelperInjector.loadHelperClass(emptyLoader.get(), helperClassName)
     emptyLoader.get().loadClass(helperClassName)
     then:
     isClassLoaded(helperClassName, emptyLoader.get())

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/HelperInjector.java

@@ -6,14 +6,12 @@
 package io.opentelemetry.javaagent.tooling;
 
 import io.opentelemetry.instrumentation.api.cache.Cache;
-import io.opentelemetry.javaagent.bootstrap.DefineClassContext;
 import io.opentelemetry.javaagent.bootstrap.HelperResources;
-import io.opentelemetry.javaagent.bootstrap.InjectedHelperClassDetector;
+import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
 import io.opentelemetry.javaagent.tooling.muzzle.HelperResource;
 import java.io.File;
 import java.io.IOException;
 import java.lang.instrument.Instrumentation;
-import java.lang.ref.WeakReference;
 import java.net.URL;
 import java.nio.file.Files;
 import java.security.SecureClassLoader;
@@ -25,7 +23,7 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder.Transformer;
 import net.bytebuddy.description.type.TypeDescription;
@@ -52,7 +50,8 @@ public class HelperInjector implements Transformer {
       TransformSafeLogger.getLogger(HelperInjector.class);
 
   static {
-    InjectedHelperClassDetector.internalSetHelperClassDetector(HelperInjector::isInjectedClass);
+    InjectedClassHelper.internalSetHelperClassDetector(HelperInjector::isInjectedClass);
+    InjectedClassHelper.internalSetHelperClassLoader(HelperInjector::loadHelperClass);
   }
 
   // Need this because we can't put null into the injectedClassLoaders map.
@@ -64,7 +63,16 @@ public class HelperInjector implements Transformer {
         }
       };
 
-  private static final Cache<Class<?>, Boolean> injectedClasses = Cache.weak();
+  private static final HelperClassInjector BOOT_CLASS_INJECTOR =
+      new HelperClassInjector(null) {
+        @Override
+        Class<?> inject(ClassLoader classLoader, String className) {
+          throw new UnsupportedOperationException("should not be called");
+        }
+      };
+
+  private static final Cache<ClassLoader, Map<String, HelperClassInjector>> helperInjectors =
+      Cache.weak();
 
   private final String requestingName;
 
@@ -77,8 +85,6 @@ public class HelperInjector implements Transformer {
   private final Cache<ClassLoader, Boolean> injectedClassLoaders = Cache.weak();
   private final Cache<ClassLoader, Boolean> resourcesInjectedClassLoaders = Cache.weak();
 
-  private final List<WeakReference<Object>> helperModules = new CopyOnWriteArrayList<>();
-
   /**
    * Construct HelperInjector.
    *
@@ -187,9 +193,7 @@ public class HelperInjector implements Transformer {
 
   private void injectHelperClasses(
       TypeDescription typeDescription, ClassLoader classLoader, JavaModule module) {
-    if (classLoader == null) {
-      classLoader = BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
-    }
+    classLoader = maskNullClassLoader(classLoader);
     if (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER && instrumentation == null) {
       logger.error(
           "Cannot inject helpers into bootstrap classloader without an instance of Instrumentation. Programmer error!");
@@ -201,24 +205,24 @@ public class HelperInjector implements Transformer {
         cl -> {
           try {
             logger.debug("Injecting classes onto classloader {} -> {}", cl, helperClassNames);
-            DefineClassContext.helpersInjected();
 
             Map<String, byte[]> classnameToBytes = getHelperMap();
-            Map<String, Class<?>> classes;
-            if (cl == BOOTSTRAP_CLASSLOADER_PLACEHOLDER) {
-              classes = injectBootstrapClassLoader(classnameToBytes);
-            } else {
-              classes = injectClassLoader(cl, classnameToBytes);
+            Map<String, HelperClassInjector> map =
+                helperInjectors.computeIfAbsent(cl, (unused) -> new ConcurrentHashMap<>());
+            for (Map.Entry<String, byte[]> entry : classnameToBytes.entrySet()) {
+              // for boot loader we use a placeholder injector, we only need these classes to be
+              // in the injected classes map to later tell which of the classes are injected
+              HelperClassInjector injector =
+                  isBootClassLoader(cl)
+                      ? BOOT_CLASS_INJECTOR
+                      : new HelperClassInjector(entry.getValue());
+              map.put(entry.getKey(), injector);
             }
 
-            classes.values().forEach(c -> injectedClasses.put(c, Boolean.TRUE));
-
-            // All agent helper classes are in the unnamed module
-            // And there's exactly one unnamed module per classloader
-            // Use the module of the first class for convenience
-            if (JavaModule.isSupported()) {
-              JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
-              helperModules.add(new WeakReference<>(javaModule.unwrap()));
+            // For boot loader we define the classes immediately. For other loaders we load them
+            // from the loadClass method of the class loader.
+            if (isBootClassLoader(cl)) {
+              injectBootstrapClassLoader(classnameToBytes);
             }
           } catch (Exception e) {
             logger.error(
@@ -231,8 +235,6 @@ public class HelperInjector implements Transformer {
           }
           return true;
         });
-
-    ensureModuleCanReadHelperModules(module);
   }
 
   private Map<String, Class<?>> injectBootstrapClassLoader(Map<String, byte[]> classnameToBytes)
@@ -257,37 +259,6 @@ public class HelperInjector implements Transformer {
     }
   }
 
-  private static Map<String, Class<?>> injectClassLoader(
-      ClassLoader classLoader, Map<String, byte[]> classnameToBytes) {
-    return new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
-  }
-
-  // JavaModule.equals doesn't work for some reason
-  @SuppressWarnings("ReferenceEquality")
-  private void ensureModuleCanReadHelperModules(JavaModule target) {
-    if (JavaModule.isSupported() && target != JavaModule.UNSUPPORTED && target.isNamed()) {
-      for (WeakReference<Object> helperModuleReference : helperModules) {
-        Object realModule = helperModuleReference.get();
-        if (realModule != null) {
-          JavaModule helperModule = JavaModule.of(realModule);
-
-          if (!target.canRead(helperModule)) {
-            logger.debug("Adding module read from {} to {}", target, helperModule);
-            ClassInjector.UsingInstrumentation.redefineModule(
-                // TODO can we guarantee that this is always present?
-                instrumentation,
-                target,
-                Collections.singleton(helperModule),
-                Collections.emptyMap(),
-                Collections.emptyMap(),
-                Collections.emptySet(),
-                Collections.emptyMap());
-          }
-        }
-      }
-    }
-  }
-
   private static File createTempDir() throws IOException {
     return Files.createTempDirectory("opentelemetry-temp-jars").toFile();
   }
@@ -302,7 +273,54 @@ public class HelperInjector implements Transformer {
     }
   }
 
-  public static boolean isInjectedClass(Class<?> c) {
-    return Boolean.TRUE.equals(injectedClasses.get(c));
+  private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
+    return classLoader != null ? classLoader : BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
+  }
+
+  private static boolean isBootClassLoader(ClassLoader classLoader) {
+    return classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
+  }
+
+  public static boolean isInjectedClass(Class<?> clazz) {
+    return isInjectedClass(clazz.getClassLoader(), clazz.getName());
+  }
+
+  public static boolean isInjectedClass(ClassLoader classLoader, String className) {
+    Map<String, HelperClassInjector> injectorMap =
+        helperInjectors.get(maskNullClassLoader(classLoader));
+    if (injectorMap == null) {
+      return false;
+    }
+    return injectorMap.containsKey(className);
+  }
+
+  public static Class<?> loadHelperClass(ClassLoader classLoader, String className) {
+    if (classLoader == null) {
+      throw new IllegalStateException("boot loader not supported");
+    }
+    Map<String, HelperClassInjector> injectorMap = helperInjectors.get(classLoader);
+    if (injectorMap == null) {
+      return null;
+    }
+    HelperClassInjector helperClassInjector = injectorMap.get(className);
+    if (helperClassInjector == null) {
+      return null;
+    }
+    return helperClassInjector.inject(classLoader, className);
+  }
+
+  private static class HelperClassInjector {
+    private final byte[] bytes;
+
+    HelperClassInjector(byte[] bytes) {
+      this.bytes = bytes;
+    }
+
+    Class<?> inject(ClassLoader classLoader, String className) {
+      Map<String, Class<?>> result =
+          new ClassInjector.UsingReflection(classLoader)
+              .injectRaw(Collections.singletonMap(className, bytes));
+      return result.get(className);
+    }
   }
 }

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/HelperReferenceWrapper.java

@@ -152,10 +152,15 @@ interface HelperReferenceWrapper {
   class Factory {
     private final TypePool classpathPool;
     private final Map<String, ClassRef> helperReferences;
+    private final HelperClassPredicate helperClassPredicate;
 
-    public Factory(TypePool classpathPool, Map<String, ClassRef> helperReferences) {
+    public Factory(
+        TypePool classpathPool,
+        Map<String, ClassRef> helperReferences,
+        HelperClassPredicate helperClassPredicate) {
       this.classpathPool = classpathPool;
       this.helperReferences = helperReferences;
+      this.helperClassPredicate = helperClassPredicate;
     }
 
     public HelperReferenceWrapper create(ClassRef reference) {
@@ -163,9 +168,15 @@ interface HelperReferenceWrapper {
     }
 
     private HelperReferenceWrapper create(String className) {
-      Resolution resolution = classpathPool.describe(className);
-      if (resolution.isResolved()) {
-        return new ClasspathType(resolution.resolve());
+      // Looking up an injected helper is recorded in AgentCachingPoolStrategy as a failed resolve
+      // because injected helper can't be found by the TypePool that is used here. If that helper is
+      // defined before it gets evicted there will be a stack trace about resource not found which
+      // is flagged as a transformation failure that makes tests fail.
+      if (!helperClassPredicate.isHelperClass(className)) {
+        Resolution resolution = classpathPool.describe(className);
+        if (resolution.isResolved()) {
+          return new ClasspathType(resolution.resolve());
+        }
       }
       // checking helper references is needed when one helper class A extends another helper class B
       // and the subclass A also implements a library interface C

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

@@ -134,7 +134,8 @@ public final class ReferenceMatcher {
     List<Mismatch> mismatches = emptyList();
 
     HelperReferenceWrapper helperWrapper =
-        new HelperReferenceWrapper.Factory(typePool, references).create(helperClass);
+        new HelperReferenceWrapper.Factory(typePool, references, helperClassPredicate)
+            .create(helperClass);
 
     Set<HelperReferenceWrapper.Field> undeclaredFields =
         helperClass.getFields().stream()

muzzle/src/test/groovy/io/opentelemetry/javaagent/tooling/muzzle/HelperReferenceWrapperTest.groovy

@@ -46,7 +46,8 @@ class HelperReferenceWrapperTest extends Specification {
     ]
 
     when:
-    def helperWrapper = new HelperReferenceWrapper.Factory(typePool, references).create(helperClass)
+    def helperClassPredicate = new HelperClassPredicate({ false })
+    def helperWrapper = new HelperReferenceWrapper.Factory(typePool, references, helperClassPredicate).create(helperClass)
 
     then:
     with(helperWrapper) { helper ->

muzzle/src/test/groovy/io/opentelemetry/javaagent/tooling/muzzle/ReferenceMatcherTest.groovy

@@ -11,6 +11,9 @@ import io.opentelemetry.instrumentation.test.utils.ClasspathUtils
 import io.opentelemetry.javaagent.tooling.muzzle.references.ClassRef
 import io.opentelemetry.javaagent.tooling.muzzle.references.Flag
 import io.opentelemetry.javaagent.tooling.muzzle.references.Source
+import io.opentelemetry.test.AnotherTestInterface
+import io.opentelemetry.test.TestAbstractSuperClass
+import io.opentelemetry.test.TestInterface
 import muzzle.TestClasses
 import muzzle.TestClasses.MethodBodyAdvice
 import org.objectweb.asm.Type
@@ -209,7 +212,7 @@ class ReferenceMatcherTest extends Specification {
   def "should fail #desc helper classes that does not implement all abstract methods"() {
     given:
     def reference = ClassRef.builder(className)
-      .setSuperClassName(TestHelperClasses.HelperSuperClass.name)
+      .setSuperClassName(TestAbstractSuperClass.name)
       .addMethod(new Source[0], [] as Flag[], "someMethod", Type.VOID_TYPE)
       .build()
 
@@ -228,10 +231,10 @@ class ReferenceMatcherTest extends Specification {
 
   def "should fail #desc helper classes that do not implement all abstract methods - even if empty abstract class reference exists"() {
     given:
-    def emptySuperClassRef = ClassRef.builder(TestHelperClasses.HelperSuperClass.name)
+    def emptySuperClassRef = ClassRef.builder(TestAbstractSuperClass.name)
       .build()
     def reference = ClassRef.builder(className)
-      .setSuperClassName(TestHelperClasses.HelperSuperClass.name)
+      .setSuperClassName(TestAbstractSuperClass.name)
       .addMethod(new Source[0], [] as Flag[], "someMethod", Type.VOID_TYPE)
       .build()
 
@@ -254,13 +257,13 @@ class ReferenceMatcherTest extends Specification {
     given:
     def baseHelper = ClassRef.builder("io.opentelemetry.instrumentation.BaseHelper")
       .setSuperClassName(Object.name)
-      .addInterfaceName(TestHelperClasses.HelperInterface.name)
+      .addInterfaceName(TestInterface.name)
       .addMethod(new Source[0], [] as Flag[], "foo", Type.VOID_TYPE)
       .build()
     // abstract HelperInterface#foo() is implemented by BaseHelper
     def helper = ClassRef.builder(className)
       .setSuperClassName(baseHelper.className)
-      .addInterfaceName(TestHelperClasses.AnotherHelperInterface.name)
+      .addInterfaceName(AnotherTestInterface.name)
       .addMethod(new Source[0], [] as Flag[], "bar", Type.VOID_TYPE)
       .build()
 

muzzle/src/test/java/io/opentelemetry/test/AnotherTestInterface.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.test;
+
+public interface AnotherTestInterface extends TestInterface {
+  void bar();
+
+  @Override
+  int hashCode();
+
+  @Override
+  boolean equals(Object other);
+
+  Object clone();
+
+  @SuppressWarnings("checkstyle:NoFinalizer")
+  void finalize();
+}

muzzle/src/test/java/io/opentelemetry/test/TestAbstractSuperClass.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.test;
+
+public abstract class TestAbstractSuperClass {
+  protected abstract int abstractMethod();
+
+  public final String finalMethod() {
+    return "42";
+  }
+
+  static int bar() {
+    return 12345;
+  }
+}

muzzle/src/test/java/io/opentelemetry/test/TestInterface.java

@@ -0,0 +1,10 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.test;
+
+public interface TestInterface {
+  void foo();
+}