Add capability for invokedynamic InstrumentationModules to inject proxies (#9565)

Author: JonasKunz

instrumentation/log4j/log4j-context-data/log4j-context-data-2.17/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/log4j/contextdata/v2_17/Log4j2InstrumentationModule.java

@@ -14,12 +14,16 @@
 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.extension.instrumentation.internal.injection.ClassInjector;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
 import java.util.List;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
-public class Log4j2InstrumentationModule extends InstrumentationModule {
+public class Log4j2InstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
   public Log4j2InstrumentationModule() {
     super("log4j-context-data", "log4j-context-data-2.17");
   }
@@ -31,8 +35,11 @@ public void registerHelperResources(HelperResourceBuilder helperResourceBuilder)
   }
 
   @Override
-  public boolean isIndyModule() {
-    return false;
+  public void injectClasses(ClassInjector injector) {
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.instrumentation.log4j.contextdata.v2_17.OpenTelemetryContextDataProvider")
+        .inject(InjectionMode.CLASS_ONLY);
   }
 
   @Override

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal;
+
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ClassInjector;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public interface ExperimentalInstrumentationModule {
+
+  /**
+   * 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/ClassInjector.java

@@ -0,0 +1,39 @@
+/*
+ * 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);
+
+  /**
+   * Same as invoking {@link #proxyBuilder(String, String)}, but the resulting proxy will have the
+   * same name as the proxied class.
+   *
+   * @param classToProxy the fully qualified name of the class for which a proxy will be generated
+   * @return a builder for further customizing and injecting the proxy
+   */
+  default ProxyInjectionBuilder proxyBuilder(String classToProxy) {
+    return proxyBuilder(classToProxy, classToProxy);
+  }
+}

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.ExperimentalInstrumentationModule;
 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 ExperimentalInstrumentationModule) {
+      ((ExperimentalInstrumentationModule) 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();
+          });
+    }
+  }
+}