Allow multiple invokedynamic InstrumentationModules to share classloaders (#10015)

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AbstractAwsSdkInstrumentationModule.java

@@ -12,24 +12,26 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import java.util.List;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
-abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule {
+abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
 
   protected AbstractAwsSdkInstrumentationModule(String additionalInstrumentationName) {
     super("aws-sdk", "aws-sdk-2.2", additionalInstrumentationName);
   }
 
   @Override
-  public boolean isHelperClass(String className) {
+  public String getModuleGroup() {
-    return className.startsWith("io.opentelemetry.contrib.awsxray.");
+    return "aws-sdk-v2";
   }
 
   @Override
-  public boolean isIndyModule() {
+  public boolean isHelperClass(String className) {
-    return false;
+    return className.startsWith("io.opentelemetry.contrib.awsxray.");
   }
 
   @Override

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AwsSdkInstrumentationModule.java

@@ -10,6 +10,8 @@ import io.opentelemetry.instrumentation.awssdk.v2_2.autoconfigure.TracingExecuti
 import io.opentelemetry.javaagent.extension.instrumentation.HelperResourceBuilder;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ClassInjector;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
 
 @AutoService(InstrumentationModule.class)
 public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
@@ -26,6 +28,14 @@ public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationMo
     helperResourceBuilder.register("software/amazon/awssdk/global/handlers/execution.interceptors");
   }
 
+  @Override
+  public void injectClasses(ClassInjector injector) {
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.instrumentation.awssdk.v2_2.autoconfigure.TracingExecutionInterceptor")
+        .inject(InjectionMode.CLASS_ONLY);
+  }
+
   @Override
   void doTransform(TypeTransformer transformer) {
     // Nothing to transform, this type instrumentation is only used for injecting resources.

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

@@ -36,4 +36,16 @@ public interface ExperimentalInstrumentationModule {
   default List<String> injectedClassNames() {
     return emptyList();
   }
+
+  /**
+   * By default every InstrumentationModule is loaded by an isolated classloader, even if multiple
+   * modules instrument the same application classloader.
+   *
+   * <p>Sometimes this is not desired, e.g. when instrumenting modular libraries such as the AWS
+   * SDK. In such cases the {@link InstrumentationModule}s which want to share a classloader can
+   * return the same group name from this method.
+   */
+  default String getModuleGroup() {
+    return getClass().getName();
+  }
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java

@@ -109,8 +109,6 @@ public final class InstrumentationModuleInstaller {
       injectedHelperClassNames = Collections.emptyList();
     }
 
-    IndyModuleRegistry.registerIndyModule(instrumentationModule);
-
     ClassInjectorImpl injectedClassesCollector = new ClassInjectorImpl(instrumentationModule);
     if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
       ((ExperimentalInstrumentationModule) instrumentationModule)
@@ -149,14 +147,17 @@ public final class InstrumentationModuleInstaller {
       AgentBuilder.Identified.Extendable extendableAgentBuilder =
           setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
               .and(muzzleMatcher)
-              .transform(new PatchByteCodeVersionTransformer())
+              .transform(new PatchByteCodeVersionTransformer());
-              .transform(helperInjector);
 
       // TODO (Jonas): we are not calling
       // contextProvider.rewriteVirtualFieldsCalls(extendableAgentBuilder) anymore
       // As a result the advices should store `VirtualFields` as static variables instead of having
       // the lookup inline
       // We need to update our documentation on that
+      extendableAgentBuilder =
+          IndyModuleRegistry.initializeModuleLoaderOnMatch(
+              instrumentationModule, extendableAgentBuilder);
+      extendableAgentBuilder = extendableAgentBuilder.transform(helperInjector);
       extendableAgentBuilder = contextProvider.injectHelperClasses(extendableAgentBuilder);
       IndyTypeTransformerImpl typeTransformer =
           new IndyTypeTransformerImpl(extendableAgentBuilder, instrumentationModule);

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/MuzzleMatcher.java

@@ -15,6 +15,7 @@ import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.config.AgentConfig;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.IndyModuleRegistry;
+import io.opentelemetry.javaagent.tooling.instrumentation.indy.InstrumentationModuleClassLoader;
 import io.opentelemetry.javaagent.tooling.muzzle.Mismatch;
 import io.opentelemetry.javaagent.tooling.muzzle.ReferenceMatcher;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
@@ -61,12 +62,19 @@ class MuzzleMatcher implements AgentBuilder.RawMatcher {
       ProtectionDomain protectionDomain) {
     if (classLoader == BOOTSTRAP_LOADER) {
       classLoader = Utils.getBootstrapProxy();
-    } else if (instrumentationModule.isIndyModule()) {
-      classLoader =
-          IndyModuleRegistry.getInstrumentationClassloader(
-              instrumentationModule.getClass().getName(), classLoader);
     }
-    return matchCache.computeIfAbsent(classLoader, this::doesMatch);
+    if (instrumentationModule.isIndyModule()) {
+      return matchCache.computeIfAbsent(
+          classLoader,
+          cl -> {
+            InstrumentationModuleClassLoader moduleCl =
+                IndyModuleRegistry.createInstrumentationClassLoaderWithoutRegistration(
+                    instrumentationModule, cl);
+            return doesMatch(moduleCl);
+          });
+    } else {
+      return matchCache.computeIfAbsent(classLoader, this::doesMatch);
+    }
   }
 
   private boolean doesMatch(ClassLoader classLoader) {

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

@@ -10,6 +10,7 @@ import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.C
 import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
 import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ProxyInjectionBuilder;
 import io.opentelemetry.javaagent.tooling.HelperClassDefinition;
+import io.opentelemetry.javaagent.tooling.muzzle.AgentTooling;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -57,7 +58,11 @@ public class ClassInjectorImpl implements ClassInjector {
     public void inject(InjectionMode mode) {
       classesToInject.add(
           cl -> {
-            TypePool typePool = IndyModuleTypePool.get(cl, instrumentationModule);
+            InstrumentationModuleClassLoader moduleCl =
+                IndyModuleRegistry.getInstrumentationClassLoader(instrumentationModule, cl);
+            TypePool typePool =
+                AgentTooling.poolStrategy()
+                    .typePool(AgentTooling.locationStrategy().classFileLocator(moduleCl), moduleCl);
             TypeDescription proxiedType = typePool.describe(classToProxy).resolve();
             DynamicType.Unloaded<?> proxy = proxyFactory.generateProxy(proxiedType, proxyClassName);
             return HelperClassDefinition.create(proxy, mode);

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

@@ -39,7 +39,7 @@ import net.bytebuddy.utility.JavaConstant;
  *       ↑ └───────── IndyBootstrapDispatcher ─ ↑ ──→ └────────────── {@link IndyBootstrap#bootstrap}
  *     Ext/Platform CL               ↑          │                        ╷
  *       ↑                           ╷          │                        ↓
- *     System CL                     ╷          │        {@link IndyModuleRegistry#getInstrumentationClassloader(String, ClassLoader)}
+ *     System CL                     ╷          │        {@link IndyModuleRegistry#getInstrumentationClassLoader(String, ClassLoader)}
  *       ↑                           ╷          │                        ╷
  *     Common               linking of CallSite │                        ╷
  *     ↑    ↑             (on first invocation) │                        ╷
@@ -171,7 +171,7 @@ public class IndyBootstrap {
       }
 
       InstrumentationModuleClassLoader instrumentationClassloader =
-          IndyModuleRegistry.getInstrumentationClassloader(
+          IndyModuleRegistry.getInstrumentationClassLoader(
               moduleClassName, lookup.lookupClass().getClassLoader());
 
       // Advices are not inlined. They are loaded as normal classes by the
@@ -207,7 +207,7 @@ public class IndyBootstrap {
       String methodKind)
       throws NoSuchMethodException, IllegalAccessException, ClassNotFoundException {
     InstrumentationModuleClassLoader instrumentationClassloader =
-        IndyModuleRegistry.getInstrumentationClassloader(
+        IndyModuleRegistry.getInstrumentationClassLoader(
             moduleClassName, lookup.lookupClass().getClassLoader());
 
     Class<?> proxiedClass = instrumentationClassloader.loadClass(proxyClassName);

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

@@ -5,129 +5,124 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
-import io.opentelemetry.instrumentation.api.internal.cache.Cache;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
-import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
-import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
-import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
+import io.opentelemetry.javaagent.tooling.util.ClassLoaderValue;
-import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
-import java.lang.ref.WeakReference;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.stream.Collectors;
 import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.description.method.MethodDescription;
-import net.bytebuddy.matcher.ElementMatcher;
 
 public class IndyModuleRegistry {
 
   private IndyModuleRegistry() {}
 
-  private static final ConcurrentHashMap<String, InstrumentationModule> modulesByName =
+  private static final ConcurrentHashMap<String, InstrumentationModule> modulesByClassName =
       new ConcurrentHashMap<>();
 
   /**
-   * Weakly references the {@link InstrumentationModuleClassLoader}s for a given application
+   * Weakly references the {@link InstrumentationModuleClassLoader}s for a given application class
-   * classloader. We only store weak references to make sure we don't prevent application
+   * loader. The {@link InstrumentationModuleClassLoader} are kept alive by a strong reference from
-   * classloaders from being GCed. The application classloaders will strongly reference the {@link
+   * the instrumented class loader realized via {@link ClassLoaderValue}.
-   * InstrumentationModuleClassLoader} through the invokedynamic callsites.
+   *
+   * <p>The keys of the contained map are the instrumentation module group names, see {@link
+   * ExperimentalInstrumentationModule#getModuleGroup()};
    */
-  private static final ConcurrentHashMap<
+  private static final ClassLoaderValue<Map<String, InstrumentationModuleClassLoader>>
-          InstrumentationModule,
+      instrumentationClassLoaders = new ClassLoaderValue<>();
-          Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>>>
-      instrumentationClassloaders = new ConcurrentHashMap<>();
 
-  public static InstrumentationModuleClassLoader getInstrumentationClassloader(
+  public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
-      String moduleClassName, ClassLoader instrumentedClassloader) {
+      String moduleClassName, ClassLoader instrumentedClassLoader) {
-    InstrumentationModule instrumentationModule = modulesByName.get(moduleClassName);
+    InstrumentationModule instrumentationModule = modulesByClassName.get(moduleClassName);
     if (instrumentationModule == null) {
       throw new IllegalArgumentException(
-          "No module with the class name " + modulesByName + " has been registered!");
+          "No module with the class name " + modulesByClassName + " has been registered!");
     }
-    return getInstrumentationClassloader(instrumentationModule, instrumentedClassloader);
+    return getInstrumentationClassLoader(instrumentationModule, instrumentedClassLoader);
   }
 
-  private static synchronized InstrumentationModuleClassLoader getInstrumentationClassloader(
+  public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
-      InstrumentationModule module, ClassLoader instrumentedClassloader) {
+      InstrumentationModule module, ClassLoader instrumentedClassLoader) {
 
-    Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>> cacheForModule =
+    String groupName = getModuleGroup(module);
-        instrumentationClassloaders.computeIfAbsent(module, (k) -> Cache.weak());
 
-    instrumentedClassloader = maskNullClassLoader(instrumentedClassloader);
+    Map<String, InstrumentationModuleClassLoader> loadersByGroupName =
-    WeakReference<InstrumentationModuleClassLoader> cached =
+        instrumentationClassLoaders.get(instrumentedClassLoader);
-        cacheForModule.get(instrumentedClassloader);
-    if (cached != null) {
-      InstrumentationModuleClassLoader cachedCl = cached.get();
-      if (cachedCl != null) {
-        return cachedCl;
-      }
-    }
-    // We can't directly use "compute-if-absent" here because then for a short time only the
-    // WeakReference will point to the InstrumentationModuleCL
-    InstrumentationModuleClassLoader created =
-        createInstrumentationModuleClassloader(module, instrumentedClassloader);
-    cacheForModule.put(instrumentedClassloader, new WeakReference<>(created));
-    return created;
-  }
 
-  private static final ClassLoader BOOT_LOADER = new ClassLoader() {};
+    if (loadersByGroupName == null) {
-
+      throw new IllegalArgumentException(
-  private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
+          module
-    return classLoader == null ? BOOT_LOADER : classLoader;
+              + " has not been initialized for class loader "
-  }
+              + instrumentedClassLoader
-
+              + " yet");
-  static InstrumentationModuleClassLoader createInstrumentationModuleClassloader(
-      InstrumentationModule module, ClassLoader instrumentedClassloader) {
-
-    Set<String> toInject = new HashSet<>(InstrumentationModuleMuzzle.getHelperClassNames(module));
-    // TODO (Jonas): Make muzzle include advice classes as helper classes
-    // so that we don't have to include them here
-    toInject.addAll(getModuleAdviceNames(module));
-    if (module instanceof ExperimentalInstrumentationModule) {
-      toInject.removeAll(((ExperimentalInstrumentationModule) module).injectedClassNames());
     }
 
+    InstrumentationModuleClassLoader loader = loadersByGroupName.get(groupName);
+    if (loader == null || !loader.hasModuleInstalled(module)) {
+      throw new IllegalArgumentException(
+          module
+              + " has not been initialized for class loader "
+              + instrumentedClassLoader
+              + " yet");
+    }
+
+    return loader;
+  }
+
+  /**
+   * Returns a newly created class loader containing only the provided module. Note that other
+   * modules from the same module group (see {@link #getModuleGroup(InstrumentationModule)}) will
+   * not be installed in this class loader.
+   */
+  public static InstrumentationModuleClassLoader
+      createInstrumentationClassLoaderWithoutRegistration(
+          InstrumentationModule module, ClassLoader instrumentedClassLoader) {
+    // TODO: remove this method and replace usages with a custom TypePool implementation instead
     ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
-    Map<String, BytecodeWithUrl> injectedClasses =
+    InstrumentationModuleClassLoader cl =
-        toInject.stream()
+        new InstrumentationModuleClassLoader(instrumentedClassLoader, agentOrExtensionCl);
-            .collect(
+    cl.installModule(module);
-                Collectors.toMap(
+    return cl;
-                    name -> name, name -> BytecodeWithUrl.create(name, agentOrExtensionCl)));
-
-    return new InstrumentationModuleClassLoader(
-        instrumentedClassloader, agentOrExtensionCl, injectedClasses);
   }
 
-  public static void registerIndyModule(InstrumentationModule module) {
+  public static AgentBuilder.Identified.Extendable initializeModuleLoaderOnMatch(
+      InstrumentationModule module, AgentBuilder.Identified.Extendable agentBuilder) {
     if (!module.isIndyModule()) {
       throw new IllegalArgumentException("Provided module is not an indy module!");
     }
     String moduleName = module.getClass().getName();
-    if (modulesByName.putIfAbsent(moduleName, module) != null) {
+    InstrumentationModule existingRegistration = modulesByClassName.putIfAbsent(moduleName, module);
+    if (existingRegistration != null && existingRegistration != module) {
       throw new IllegalArgumentException(
-          "A module with the class name " + moduleName + " has already been registered!");
+          "A different module with the class name " + moduleName + " has already been registered!");
     }
+    return agentBuilder.transform(
+        (builder, typeDescription, classLoader, javaModule, protectionDomain) -> {
+          initializeModuleLoaderForClassLoader(module, classLoader);
+          return builder;
+        });
   }
 
-  private static Set<String> getModuleAdviceNames(InstrumentationModule module) {
+  private static void initializeModuleLoaderForClassLoader(
-    Set<String> adviceNames = new HashSet<>();
+      InstrumentationModule module, ClassLoader classLoader) {
-    TypeTransformer nameCollector =
-        new TypeTransformer() {
-          @Override
-          public void applyAdviceToMethod(
-              ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
-            adviceNames.add(adviceClassName);
-          }
 
-          @Override
+    ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
-          public void applyTransformer(AgentBuilder.Transformer transformer) {}
+
-        };
+    String groupName = getModuleGroup(module);
-    for (TypeInstrumentation instr : module.typeInstrumentations()) {
+
-      instr.transform(nameCollector);
+    InstrumentationModuleClassLoader moduleCl =
+        instrumentationClassLoaders
+            .computeIfAbsent(classLoader, ConcurrentHashMap::new)
+            .computeIfAbsent(
+                groupName,
+                unused -> new InstrumentationModuleClassLoader(classLoader, agentOrExtensionCl));
+
+    moduleCl.installModule(module);
+  }
+
+  private static String getModuleGroup(InstrumentationModule module) {
+    if (module instanceof ExperimentalInstrumentationModule) {
+      return ((ExperimentalInstrumentationModule) module).getModuleGroup();
     }
-    return adviceNames;
+    return module.getClass().getName();
   }
 }

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

@@ -1,34 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.javaagent.tooling.instrumentation.indy;
-
-import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
-import io.opentelemetry.javaagent.tooling.muzzle.AgentTooling;
-import net.bytebuddy.pool.TypePool;
-
-public class IndyModuleTypePool {
-
-  private IndyModuleTypePool() {}
-
-  /**
-   * Provides a {@link TypePool} which has the same lookup rules for {@link
-   * net.bytebuddy.description.type.TypeDescription}s as {@link InstrumentationModuleClassLoader}
-   * have for classes.
-   *
-   * @param instrumentedCl the classloader being instrumented (e.g. for which the {@link
-   *     InstrumentationModuleClassLoader} is being created).
-   * @param module the {@link InstrumentationModule} performing the instrumentation
-   * @return the type pool, must not be cached!
-   */
-  public static TypePool get(ClassLoader instrumentedCl, InstrumentationModule module) {
-    // TODO: this implementation doesn't allow caching the returned pool and its types
-    // This could be improved by implementing a custom TypePool instead, which delegates to parent
-    // TypePools and mirrors the delegation model of the InstrumentationModuleClassLoader
-    InstrumentationModuleClassLoader dummyCl =
-        IndyModuleRegistry.createInstrumentationModuleClassloader(module, instrumentedCl);
-    return TypePool.Default.of(AgentTooling.locationStrategy().classFileLocator(dummyCl));
-  }
-}

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

@@ -5,7 +5,12 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
 import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
+import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
 import java.io.IOException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
@@ -15,28 +20,35 @@ import java.security.PrivilegedAction;
 import java.security.ProtectionDomain;
 import java.util.Collections;
 import java.util.Enumeration;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.StringMatcher;
 
 /**
- * Classloader used to load the helper classes from {@link
+ * Class loader used to load the helper classes from {@link
  * io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule}s, so that those
  * classes have access to both the agent/extension classes and the instrumented application classes.
  *
- * <p>This classloader implements the following classloading delegation strategy:
+ * <p>This class loader implements the following classloading delegation strategy:
  *
  * <ul>
  *   <li>First, injected classes are considered (usually the helper classes from the
  *       InstrumentationModule)
- *   <li>Next, the classloader looks in the agent or extension classloader, depending on where the
+ *   <li>Next, the class loader looks in the agent or extension class loader, depending on where the
  *       InstrumentationModule comes from
- *   <li>Finally, the instrumented application classloader is checked for the class
+ *   <li>Finally, the instrumented application class loader is checked for the class
  * </ul>
  *
- * <p>In addition, this classloader ensures that the lookup of corresponding .class resources follow
+ * <p>In addition, this class loader ensures that the lookup of corresponding .class resources
- * the same delegation strategy, so that bytecode inspection tools work correctly.
+ * follow the same delegation strategy, so that bytecode inspection tools work correctly.
  */
 public class InstrumentationModuleClassLoader extends ClassLoader {
 
@@ -44,6 +56,8 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
     ClassLoader.registerAsParallelCapable();
   }
 
+  private static final ClassLoader BOOT_LOADER = new ClassLoader() {};
+
   private static final Map<String, BytecodeWithUrl> ALWAYS_INJECTED_CLASSES =
       Collections.singletonMap(
           LookupExposer.class.getName(), BytecodeWithUrl.create(LookupExposer.class).cached());
@@ -54,33 +68,42 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
   private final ClassLoader agentOrExtensionCl;
   private volatile MethodHandles.Lookup cachedLookup;
 
-  private final ClassLoader instrumentedCl;
+  @Nullable private final ClassLoader instrumentedCl;
-  private final boolean delegateAllToAgent;
+
+  /**
+   * Only class names matching this matcher will be attempted to be loaded from the {@link
+   * #agentOrExtensionCl}. If a class is requested and it does not match this matcher, the lookup in
+   * {@link #agentOrExtensionCl} will be skipped.
+   */
+  private final ElementMatcher<String> agentClassNamesMatcher;
+
+  private final Set<InstrumentationModule> installedModules;
 
   public InstrumentationModuleClassLoader(
-      ClassLoader instrumentedCl,
+      ClassLoader instrumentedCl, ClassLoader agentOrExtensionCl) {
-      ClassLoader agentOrExtensionCl,
+    this(
-      Map<String, BytecodeWithUrl> injectedClasses) {
+        instrumentedCl,
-    this(instrumentedCl, agentOrExtensionCl, injectedClasses, false);
+        agentOrExtensionCl,
+        new StringMatcher("io.opentelemetry.javaagent", StringMatcher.Mode.STARTS_WITH));
   }
 
   InstrumentationModuleClassLoader(
-      ClassLoader instrumentedCl,
+      @Nullable ClassLoader instrumentedCl,
       ClassLoader agentOrExtensionCl,
-      Map<String, BytecodeWithUrl> injectedClasses,
+      ElementMatcher<String> classesToLoadFromAgentOrExtensionCl) {
-      boolean delegateAllToAgent) {
+    // agent/extension-class loader is "main"-parent, but class lookup is overridden
-    // agent/extension-classloader is "main"-parent, but class lookup is overridden
     super(agentOrExtensionCl);
-    additionalInjectedClasses = injectedClasses;
+    additionalInjectedClasses = new ConcurrentHashMap<>();
+    installedModules = Collections.newSetFromMap(new ConcurrentHashMap<>());
     this.agentOrExtensionCl = agentOrExtensionCl;
     this.instrumentedCl = instrumentedCl;
-    this.delegateAllToAgent = delegateAllToAgent;
+    this.agentClassNamesMatcher = classesToLoadFromAgentOrExtensionCl;
   }
 
   /**
-   * Provides a Lookup within this classloader. See {@link LookupExposer} for the details.
+   * Provides a Lookup within this class loader. See {@link LookupExposer} for the details.
    *
-   * @return a lookup capable of accessing public types in this classloader
+   * @return a lookup capable of accessing public types in this class loader
    */
   public MethodHandles.Lookup getLookup() {
     if (cachedLookup == null) {
@@ -96,6 +119,62 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
     return cachedLookup;
   }
 
+  public synchronized void installModule(InstrumentationModule module) {
+    if (module.getClass().getClassLoader() != agentOrExtensionCl) {
+      throw new IllegalArgumentException(
+          module.getClass().getName() + " is not loaded by " + agentOrExtensionCl);
+    }
+    if (!installedModules.add(module)) {
+      return;
+    }
+    Map<String, BytecodeWithUrl> classesToInject =
+        getClassesToInject(module).stream()
+            .collect(
+                Collectors.toMap(
+                    className -> className,
+                    className -> BytecodeWithUrl.create(className, agentOrExtensionCl)));
+    installInjectedClasses(classesToInject);
+  }
+
+  public synchronized boolean hasModuleInstalled(InstrumentationModule module) {
+    return installedModules.contains(module);
+  }
+
+  // Visible for testing
+  synchronized void installInjectedClasses(Map<String, BytecodeWithUrl> classesToInject) {
+    classesToInject.forEach(additionalInjectedClasses::putIfAbsent);
+  }
+
+  private static Set<String> getClassesToInject(InstrumentationModule module) {
+    Set<String> toInject = new HashSet<>(InstrumentationModuleMuzzle.getHelperClassNames(module));
+    // TODO (Jonas): Make muzzle include advice classes as helper classes
+    // so that we don't have to include them here
+    toInject.addAll(getModuleAdviceNames(module));
+    if (module instanceof ExperimentalInstrumentationModule) {
+      toInject.removeAll(((ExperimentalInstrumentationModule) module).injectedClassNames());
+    }
+    return toInject;
+  }
+
+  private static Set<String> getModuleAdviceNames(InstrumentationModule module) {
+    Set<String> adviceNames = new HashSet<>();
+    TypeTransformer nameCollector =
+        new TypeTransformer() {
+          @Override
+          public void applyAdviceToMethod(
+              ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
+            adviceNames.add(adviceClassName);
+          }
+
+          @Override
+          public void applyTransformer(AgentBuilder.Transformer transformer) {}
+        };
+    for (TypeInstrumentation instr : module.typeInstrumentations()) {
+      instr.transform(nameCollector);
+    }
+    return adviceNames;
+  }
+
   public static final Map<String, byte[]> bytecodeOverride = new ConcurrentHashMap<>();
 
   @Override
@@ -140,12 +219,12 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
   }
 
   private boolean shouldLoadFromAgent(String dotClassName) {
-    return delegateAllToAgent || dotClassName.startsWith("io.opentelemetry.javaagent");
+    return agentClassNamesMatcher.matches(dotClassName);
   }
 
-  private static Class<?> tryLoad(ClassLoader cl, String name) {
+  private static Class<?> tryLoad(@Nullable ClassLoader cl, String name) {
     try {
-      return cl.loadClass(name);
+      return Class.forName(name, false, cl);
     } catch (ClassNotFoundException e) {
       return null;
     }
@@ -155,7 +234,7 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
   public URL getResource(String resourceName) {
     String className = resourceToClassName(resourceName);
     if (className == null) {
-      // delegate to just the default parent (the agent classloader)
+      // delegate to just the default parent (the agent class loader)
       return super.getResource(resourceName);
     }
     // for classes use the same precedence as in loadClass
@@ -167,7 +246,12 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
     if (fromAgentCl != null) {
       return fromAgentCl;
     }
-    return instrumentedCl.getResource(resourceName);
+
+    if (instrumentedCl != null) {
+      return instrumentedCl.getResource(resourceName);
+    } else {
+      return BOOT_LOADER.getResource(resourceName);
+    }
   }
 
   @Override

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java

@@ -11,6 +11,7 @@ import java.lang.reflect.Field;
 import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.description.modifier.Ownership;
 import net.bytebuddy.description.modifier.Visibility;
@@ -27,6 +28,11 @@ class ClassLoaderMap {
     getClassLoaderData(classLoader, true).put(key, value);
   }
 
+  public static Object computeIfAbsent(
+      ClassLoader classLoader, Object key, Supplier<? extends Object> value) {
+    return getClassLoaderData(classLoader, true).computeIfAbsent(key, unused -> value.get());
+  }
+
   private static Map<Object, Object> getClassLoaderData(
       ClassLoader classLoader, boolean initialize) {
     classLoader = maskNullClassLoader(classLoader);

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java

@@ -5,6 +5,8 @@
 
 package io.opentelemetry.javaagent.tooling.util;
 
+import java.util.function.Supplier;
+
 /**
  * Associate value with a class loader. Added value will behave as if it was stored in a field in
  * the class loader object, meaning that the value can be garbage collected once the class loader is
@@ -21,4 +23,9 @@ public final class ClassLoaderValue<T> {
   public void put(ClassLoader classLoader, T value) {
     ClassLoaderMap.put(classLoader, this, value);
   }
+
+  @SuppressWarnings("unchecked")
+  public T computeIfAbsent(ClassLoader classLoader, Supplier<? extends T> value) {
+    return (T) ClassLoaderMap.computeIfAbsent(classLoader, this, value);
+  }
 }

javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoaderTest.java

@@ -29,6 +29,7 @@ import java.util.jar.JarOutputStream;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.implementation.FixedValue;
+import net.bytebuddy.matcher.ElementMatchers;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -44,9 +45,12 @@ class InstrumentationModuleClassLoaderTest {
     ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null);
 
     InstrumentationModuleClassLoader m1 =
-        new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true);
+        new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any());
+    m1.installInjectedClasses(toInject);
+
     InstrumentationModuleClassLoader m2 =
-        new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true);
+        new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any());
+    m2.installInjectedClasses(toInject);
 
     // MethodHandles.publicLookup() always succeeds on the first invocation
     lookupAndInvokeFoo(m1);
@@ -80,7 +84,8 @@ class InstrumentationModuleClassLoaderTest {
 
     ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null);
     InstrumentationModuleClassLoader m1 =
-        new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true);
+        new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any());
+    m1.installInjectedClasses(toInject);
 
     Class<?> injected = Class.forName(A.class.getName(), true, m1);
     // inject two classes from the same package to trigger errors if we try to redefine the package
@@ -121,7 +126,8 @@ class InstrumentationModuleClassLoaderTest {
       toInject.put(C.class.getName(), BytecodeWithUrl.create(C.class.getName(), moduleSourceCl));
 
       InstrumentationModuleClassLoader moduleCl =
-          new InstrumentationModuleClassLoader(appCl, agentCl, toInject, true);
+          new InstrumentationModuleClassLoader(appCl, agentCl, ElementMatchers.any());
+      moduleCl.installInjectedClasses(toInject);
 
       // Verify precedence for classloading
       Class<?> clA = moduleCl.loadClass(A.class.getName());

javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java

@@ -22,9 +22,22 @@ class ClassLoaderValueTest {
   }
 
   void testClassLoader(ClassLoader classLoader) {
-    ClassLoaderValue<String> classLoaderValue = new ClassLoaderValue<>();
+    ClassLoaderValue<String> value1 = new ClassLoaderValue<>();
-    classLoaderValue.put(classLoader, "value");
+    value1.put(classLoader, "value");
-    assertThat(classLoaderValue.get(classLoader)).isEqualTo("value");
+    assertThat(value1.get(classLoader)).isEqualTo("value");
+
+    ClassLoaderValue<String> value2 = new ClassLoaderValue<>();
+    String value = "value";
+    String ret1 = value2.computeIfAbsent(classLoader, () -> value);
+    String ret2 =
+        value2.computeIfAbsent(
+            classLoader,
+            () -> {
+              throw new IllegalStateException("Shouldn't be invoked");
+            });
+    assertThat(ret1).isSameAs(value);
+    assertThat(ret2).isSameAs(value);
+    assertThat(value2.get(classLoader)).isSameAs(value);
   }
 
   @Test