Add capability for invokedynamic InstrumentationModules to inject proxies

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

@@ -27,11 +27,6 @@ public boolean isHelperClass(String className) {
     return className.startsWith("io.opentelemetry.contrib.awsxray.");
   }
 
-  @Override
-  public boolean isIndyModule() {
-    return false;
-  }
-
   @Override
   public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
     // We don't actually transform it but want to make sure we only apply the instrumentation when

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

@@ -10,13 +10,49 @@
 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.ExtendedInstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 
 @AutoService(InstrumentationModule.class)
-public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
+public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationModule
+    implements ExtendedInstrumentationModule {
   public AwsSdkInstrumentationModule() {
     super("aws-sdk-2.2-core");
   }
 
+  @Override
+  @SuppressWarnings("unchecked")
+  public List<String> getAdditionalHelperClassNames() {
+    if (isIndyModule()) {
+      // With the invokedynamic approach, the SnsInstrumentationModule and SqsInstrumentationModule
+      // are not required anymore, because we don't inject the helpers which reference potentially
+      // missing classes
+      // Instead, those are loaded by the InstrumentationModuleClassloader and LinkageErrors are
+      // caught just like when using those classes as library instrumentation
+      List<String> helpers = new ArrayList<>();
+      InstrumentationModule[] modules = {
+        new SnsInstrumentationModule(), new SqsInstrumentationModule()
+      };
+      for (InstrumentationModule include : modules) {
+        try {
+          List<String> moduleRefs =
+              (List<String>)
+                  include.getClass().getDeclaredMethod("getMuzzleHelperClassNames").invoke(include);
+          helpers.addAll(moduleRefs);
+        } catch (Exception e) {
+          throw new IllegalStateException(e);
+        }
+      }
+      return helpers;
+    } else {
+      return Collections.emptyList();
+    }
+  }
+
   /**
    * Injects resource file with reference to our {@link TracingExecutionInterceptor} to allow SDK's
    * service loading mechanism to pick it up.
@@ -26,6 +62,15 @@ 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",
+            "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/injection/ClassInjector.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal.injection;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public interface ClassInjector {
+
+  /**
+   * Create a builder for a proxy class which will be injected into the instrumented {@link
+   * ClassLoader}. The generated proxy will delegate to the original class, which is loaded in a
+   * separate classloader.
+   *
+   * <p>This removes the need for the proxied class and its dependencies to be visible (just like
+   * Advices) to the instrumented ClassLoader.
+   *
+   * @param classToProxy the fully qualified name of the class for which a proxy will be generated
+   * @param newProxyName the fully qualified name to use for the generated proxy
+   * @return a builder for further customizing the proxy. {@link
+   *     ProxyInjectionBuilder#inject(InjectionMode)} must be called to actually inject the proxy.
+   */
+  ProxyInjectionBuilder proxyBuilder(String classToProxy, String newProxyName);
+}

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal.injection;
+
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public interface ExtendedInstrumentationModule {
+
+  /**
+   * Only functional for Modules where {@link InstrumentationModule#isIndyModule()} returns {@code
+   * true}.
+   *
+   * <p>Normally, helper and advice classes are loaded in a child classloader of the instrumented
+   * classloader. This method allows to inject classes directly into the instrumented classloader
+   * instead.
+   *
+   * @param injector the builder for injecting classes
+   */
+  default void injectClasses(ClassInjector injector) {}
+}

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal.injection;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public enum InjectionMode {
+  CLASS_ONLY
+  // TODO: implement the modes RESOURCE_ONLY and CLASS_AND_RESOURCE
+  // This will require a custom URL implementation for byte arrays, similar to how bytebuddy's
+  // ByteArrayClassLoader does it
+
+}

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

@@ -0,0 +1,15 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal.injection;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public interface ProxyInjectionBuilder {
+
+  void inject(InjectionMode mode);
+}

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

@@ -13,13 +13,15 @@
 
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ExtendedInstrumentationModule;
 import io.opentelemetry.javaagent.tooling.HelperInjector;
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.javaagent.tooling.bytebuddy.LoggingFailSafeMatcher;
 import io.opentelemetry.javaagent.tooling.config.AgentConfig;
 import io.opentelemetry.javaagent.tooling.field.VirtualFieldImplementationInstaller;
 import io.opentelemetry.javaagent.tooling.field.VirtualFieldImplementationInstallerFactory;
+import io.opentelemetry.javaagent.tooling.instrumentation.indy.ClassInjectorImpl;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.IndyModuleRegistry;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.IndyTypeTransformerImpl;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.PatchByteCodeVersionTransformer;
@@ -78,8 +80,25 @@ private AgentBuilder installIndyModule(
 
     IndyModuleRegistry.registerIndyModule(instrumentationModule);
 
+    HelperResourceBuilderImpl helperResourceBuilder = new HelperResourceBuilderImpl();
+    instrumentationModule.registerHelperResources(helperResourceBuilder);
+
+    ClassInjectorImpl injectedClassesCollector = new ClassInjectorImpl(instrumentationModule);
+    if (instrumentationModule instanceof ExtendedInstrumentationModule) {
+      ((ExtendedInstrumentationModule) instrumentationModule)
+          .injectClasses(injectedClassesCollector);
+    }
+
+    AgentBuilder.Transformer helperInjector =
+        new HelperInjector(
+            instrumentationModule.instrumentationName(),
+            injectedClassesCollector.getClassesToInject(),
+            helperResourceBuilder.getResources(),
+            instrumentationModule.getClass().getClassLoader(),
+            instrumentation);
+
     // TODO (Jonas): Adapt MuzzleMatcher to use the same type lookup strategy as the
-    // InstrumentationModuleClassLoader
+    // InstrumentationModuleClassLoader (see IndyModuleTypePool)
     // MuzzleMatcher muzzleMatcher = new MuzzleMatcher(logger, instrumentationModule, config);
     VirtualFieldImplementationInstaller contextProvider =
         virtualFieldInstallerFactory.create(instrumentationModule);
@@ -88,7 +107,8 @@ private AgentBuilder installIndyModule(
     for (TypeInstrumentation typeInstrumentation : instrumentationModule.typeInstrumentations()) {
       AgentBuilder.Identified.Extendable extendableAgentBuilder =
           setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
-              .transform(new PatchByteCodeVersionTransformer());
+              .transform(new PatchByteCodeVersionTransformer())
+              .transform(helperInjector);
 
       // TODO (Jonas): we are not calling
       // contextProvider.rewriteVirtualFieldsCalls(extendableAgentBuilder) anymore

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

@@ -0,0 +1,65 @@
+/*
+ * 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.extension.instrumentation.internal.injection.ClassInjector;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ProxyInjectionBuilder;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.pool.TypePool;
+
+public class ClassInjectorImpl implements ClassInjector {
+
+  private final InstrumentationModule instrumentationModule;
+
+  private final Map<String, Function<ClassLoader, byte[]>> classesToInject;
+
+  private final IndyProxyFactory proxyFactory;
+
+  public ClassInjectorImpl(InstrumentationModule module) {
+    instrumentationModule = module;
+    classesToInject = new HashMap<>();
+    proxyFactory = IndyBootstrap.getProxyFactory(module);
+  }
+
+  public Map<String, Function<ClassLoader, byte[]>> getClassesToInject() {
+    return classesToInject;
+  }
+
+  @Override
+  public ProxyInjectionBuilder proxyBuilder(String classToProxy, String newProxyName) {
+    return new ProxyBuilder(classToProxy, newProxyName);
+  }
+
+  private class ProxyBuilder implements ProxyInjectionBuilder {
+
+    private final String classToProxy;
+    private final String proxyClassName;
+
+    ProxyBuilder(String classToProxy, String proxyClassName) {
+      this.classToProxy = classToProxy;
+      this.proxyClassName = proxyClassName;
+    }
+
+    @Override
+    public void inject(InjectionMode mode) {
+      if (mode != InjectionMode.CLASS_ONLY) {
+        throw new UnsupportedOperationException("Not yet implemented");
+      }
+      classesToInject.put(
+          proxyClassName,
+          cl -> {
+            TypePool typePool = IndyModuleTypePool.get(cl, instrumentationModule);
+            TypeDescription proxiedType = typePool.describe(classToProxy).resolve();
+            return proxyFactory.generateProxy(proxiedType, proxyClassName).getBytes();
+          });
+    }
+  }
+}