Allow multiple invokedynamic InstrumentationModules to share classloaders

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 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) {
-    return className.startsWith("io.opentelemetry.contrib.awsxray.");
+  public String getModuleGroup() {
+    return "aws-sdk-v2";
   }
 
   @Override
-  public boolean isIndyModule() {
-    return false;
+  public boolean isHelperClass(String className) {
+    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.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 void registerHelperResources(HelperResourceBuilder helperResourceBuilder)
     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 @@ default void injectClasses(ClassInjector injector) {}
   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 @@ private AgentBuilder installIndyModule(
       injectedHelperClassNames = Collections.emptyList();
     }
 
-    IndyModuleRegistry.registerIndyModule(instrumentationModule);
-
     ClassInjectorImpl injectedClassesCollector = new ClassInjectorImpl(instrumentationModule);
     if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
       ((ExperimentalInstrumentationModule) instrumentationModule)
@@ -149,14 +147,17 @@ private AgentBuilder installIndyModule(
       AgentBuilder.Identified.Extendable extendableAgentBuilder =
           setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
               .and(muzzleMatcher)
-              .transform(new PatchByteCodeVersionTransformer())
-              .transform(helperInjector);
+              .transform(new PatchByteCodeVersionTransformer());
 
       // 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.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 @@ public boolean matches(
       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.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 @@ private class ProxyBuilder implements ProxyInjectionBuilder {
     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 @@
  *       ↑ └───────── 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 @@ private static ConstantCallSite bootstrapAdvice(
       }
 
       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 @@ private static ConstantCallSite bootstrapProxyMethod(
       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.muzzle.InstrumentationModuleMuzzle;
-import java.lang.ref.WeakReference;
-import java.util.HashSet;
+import io.opentelemetry.javaagent.tooling.util.ClassLoaderValue;
 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
-   * classloader. We only store weak references to make sure we don't prevent application
-   * classloaders from being GCed. The application classloaders will strongly reference the {@link
-   * InstrumentationModuleClassLoader} through the invokedynamic callsites.
+   * Weakly references the {@link InstrumentationModuleClassLoader}s for a given application class
+   * loader. The {@link InstrumentationModuleClassLoader} are kept alive by a strong reference from
+   * the instrumented class loader realized via {@link ClassLoaderValue}.
+   *
+   * <p>The keys of the contained map are the instrumentation module group names, see {@link
+   * ExperimentalInstrumentationModule#getModuleGroup()};
    */
-  private static final ConcurrentHashMap<
-          InstrumentationModule,
-          Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>>>
-      instrumentationClassloaders = new ConcurrentHashMap<>();
-
-  public static InstrumentationModuleClassLoader getInstrumentationClassloader(
-      String moduleClassName, ClassLoader instrumentedClassloader) {
-    InstrumentationModule instrumentationModule = modulesByName.get(moduleClassName);
+  private static final ClassLoaderValue<Map<String, InstrumentationModuleClassLoader>>
+      instrumentationClassLoaders = new ClassLoaderValue<>();
+
+  public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
+      String moduleClassName, ClassLoader instrumentedClassLoader) {
+    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(
-      InstrumentationModule module, ClassLoader instrumentedClassloader) {
+  public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
+      InstrumentationModule module, ClassLoader instrumentedClassLoader) {
+
+    String groupName = getModuleGroup(module);
 
-    Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>> cacheForModule =
-        instrumentationClassloaders.computeIfAbsent(module, (k) -> Cache.weak());
+    Map<String, InstrumentationModuleClassLoader> loadersByGroupName =
+        instrumentationClassLoaders.get(instrumentedClassLoader);
 
-    instrumentedClassloader = maskNullClassLoader(instrumentedClassloader);
-    WeakReference<InstrumentationModuleClassLoader> cached =
-        cacheForModule.get(instrumentedClassloader);
-    if (cached != null) {
-      InstrumentationModuleClassLoader cachedCl = cached.get();
-      if (cachedCl != null) {
-        return cachedCl;
-      }
+    if (loadersByGroupName == null) {
+      throw new IllegalArgumentException(
+          module
+              + " has not been initialized for class loader "
+              + instrumentedClassLoader
+              + " yet");
     }
-    // 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() {};
+    InstrumentationModuleClassLoader loader = loadersByGroupName.get(groupName);
+    if (loader == null || !loader.hasModuleInstalled(module)) {
+      throw new IllegalArgumentException(
+          module
+              + " has not been initialized for class loader "
+              + instrumentedClassLoader
+              + " yet");
+    }
 
-  private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
-    return classLoader == null ? BOOT_LOADER : classLoader;
+    return loader;
   }
 
-  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());
-    }
-
+  /**
+   * 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 =
-        toInject.stream()
-            .collect(
-                Collectors.toMap(
-                    name -> name, name -> BytecodeWithUrl.create(name, agentOrExtensionCl)));
-
-    return new InstrumentationModuleClassLoader(
-        instrumentedClassloader, agentOrExtensionCl, injectedClasses);
+    InstrumentationModuleClassLoader cl =
+        new InstrumentationModuleClassLoader(instrumentedClassLoader, agentOrExtensionCl);
+    cl.installModule(module);
+    return cl;
   }
 
-  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 void initializeModuleLoaderForClassLoader(
+      InstrumentationModule module, ClassLoader classLoader) {
+
+    ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
+
+    String groupName = getModuleGroup(module);
+
+    InstrumentationModuleClassLoader moduleCl =
+        instrumentationClassLoaders
+            .computeIfAbsent(classLoader, ConcurrentHashMap::new)
+            .computeIfAbsent(
+                groupName,
+                unused -> new InstrumentationModuleClassLoader(classLoader, agentOrExtensionCl));
+
+    moduleCl.installModule(module);
   }
 
-  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);
+  private static String getModuleGroup(InstrumentationModule module) {
+    if (module instanceof ExperimentalInstrumentationModule) {
+      return ((ExperimentalInstrumentationModule) module).getModuleGroup();
     }
-    return adviceNames;
+    return module.getClass().getName();
   }
 }