Allow injection of helper bytecode as resources

instrumentation/spring/spring-boot-actuator-autoconfigure-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/actuator/v2_0/SpringBootActuatorInstrumentationModule.java

@@ -12,11 +12,15 @@
 import io.opentelemetry.javaagent.extension.instrumentation.HelperResourceBuilder;
 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.extension.instrumentation.internal.injection.ClassInjector;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
 import java.util.List;
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
-public class SpringBootActuatorInstrumentationModule extends InstrumentationModule {
+public class SpringBootActuatorInstrumentationModule extends InstrumentationModule
+ implements ExperimentalInstrumentationModule {
 
  public SpringBootActuatorInstrumentationModule() {
  super(
@@ -39,14 +43,19 @@ public void registerHelperResources(HelperResourceBuilder helperResourceBuilder)
  // this line will make OpenTelemetryMeterRegistryAutoConfiguration available to all
  // classloaders, so that the bean class loader (different from the instrumented class loader)
  // can load it
- helperResourceBuilder.registerForAllClassLoaders(
- "io/opentelemetry/javaagent/instrumentation/spring/actuator/v2_0/OpenTelemetryMeterRegistryAutoConfiguration.class");
+ if (!isIndyModule()) {
+ // For indy module the proxy-bytecode will be injected as resource by injectClasses()
+ helperResourceBuilder.registerForAllClassLoaders(
+ "io/opentelemetry/javaagent/instrumentation/spring/actuator/v2_0/OpenTelemetryMeterRegistryAutoConfiguration.class");
+ }
  }
 
  @Override
- public boolean isIndyModule() {
- // can not access OpenTelemetryMeterRegistryAutoConfiguration
- return false;
+ public void injectClasses(ClassInjector injector) {
+ injector
+ .proxyBuilder(
+ "io.opentelemetry.javaagent.instrumentation.spring.actuator.v2_0.OpenTelemetryMeterRegistryAutoConfiguration")
+ .inject(InjectionMode.CLASS_AND_RESOURCE);
  }
 
  @Override

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

@@ -10,9 +10,23 @@
  * 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
+ CLASS_ONLY(true, false),
+ RESOURCE_ONLY(false, true),
+ CLASS_AND_RESOURCE(true, true);
 
+ private final boolean injectClass;
+ private final boolean injectResource;
+
+ InjectionMode(boolean injectClass, boolean injectResource) {
+ this.injectClass = injectClass;
+ this.injectResource = injectResource;
+ }
+
+ public boolean shouldInjectClass() {
+ return injectClass;
+ }
+
+ public boolean shouldInjectResource() {
+ return injectResource;
+ }
 }

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

@@ -14,6 +14,8 @@
 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.extension.instrumentation.internal.injection.InjectionMode;
+import io.opentelemetry.javaagent.tooling.HelperClassDefinition;
 import io.opentelemetry.javaagent.tooling.HelperInjector;
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.Utils;
@@ -31,8 +33,10 @@
 import io.opentelemetry.javaagent.tooling.util.NamedMatcher;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import java.lang.instrument.Instrumentation;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationSource;
 import net.bytebuddy.description.type.TypeDescription;
@@ -96,11 +100,13 @@ private AgentBuilder installIndyModule(
  return parentAgentBuilder;
  }
 
- List<String> injectedHelperClassNames = Collections.emptyList();
+ List<String> injectedHelperClassNames;
  if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
  ExperimentalInstrumentationModule experimentalInstrumentationModule =
  (ExperimentalInstrumentationModule) instrumentationModule;
  injectedHelperClassNames = experimentalInstrumentationModule.injectedClassNames();
+ } else {
+ injectedHelperClassNames = Collections.emptyList();
  }
 
  IndyModuleRegistry.registerIndyModule(instrumentationModule);
@@ -113,20 +119,27 @@ private AgentBuilder installIndyModule(
 
  MuzzleMatcher muzzleMatcher = new MuzzleMatcher(logger, instrumentationModule, config);
 
+ Function<ClassLoader, List<HelperClassDefinition>> helperGenerator =
+ cl -> {
+ List<HelperClassDefinition> helpers =
+ new ArrayList<>(injectedClassesCollector.getClassesToInject(cl));
+ for (String helperName : injectedHelperClassNames) {
+ helpers.add(
+ HelperClassDefinition.create(
+ helperName,
+ instrumentationModule.getClass().getClassLoader(),
+ InjectionMode.CLASS_ONLY));
+ }
+ return helpers;
+ };
+
  AgentBuilder.Transformer helperInjector =
  new HelperInjector(
  instrumentationModule.instrumentationName(),
- injectedHelperClassNames,
+ helperGenerator,
  helperResourceBuilder.getResources(),
  instrumentationModule.getClass().getClassLoader(),
  instrumentation);
- AgentBuilder.Transformer indyHelperInjector =
- new HelperInjector(
- instrumentationModule.instrumentationName(),
- injectedClassesCollector.getClassesToInject(),
- Collections.emptyList(),
- instrumentationModule.getClass().getClassLoader(),
- instrumentation);
 
  VirtualFieldImplementationInstaller contextProvider =
  virtualFieldInstallerFactory.create(instrumentationModule);
@@ -137,8 +150,7 @@ private AgentBuilder installIndyModule(
  setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
  .and(muzzleMatcher)
  .transform(new PatchByteCodeVersionTransformer())
- .transform(helperInjector)
- .transform(indyHelperInjector);
+ .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

@@ -9,28 +9,33 @@
 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 io.opentelemetry.javaagent.tooling.HelperClassDefinition;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.pool.TypePool;
 
 public class ClassInjectorImpl implements ClassInjector {
 
  private final InstrumentationModule instrumentationModule;
 
- private final Map<String, Function<ClassLoader, byte[]>> classesToInject;
+ private final List<Function<ClassLoader, HelperClassDefinition>> classesToInject;
 
  private final IndyProxyFactory proxyFactory;
 
  public ClassInjectorImpl(InstrumentationModule module) {
  instrumentationModule = module;
- classesToInject = new HashMap<>();
+ classesToInject = new ArrayList<>();
  proxyFactory = IndyBootstrap.getProxyFactory(module);
  }
 
- public Map<String, Function<ClassLoader, byte[]>> getClassesToInject() {
- return classesToInject;
+ public List<HelperClassDefinition> getClassesToInject(ClassLoader instrumentedCl) {
+ return classesToInject.stream()
+ .map(generator -> generator.apply(instrumentedCl))
+ .collect(Collectors.toList());
  }
 
  @Override
@@ -50,15 +55,12 @@ private class ProxyBuilder implements ProxyInjectionBuilder {
 
  @Override
  public void inject(InjectionMode mode) {
- if (mode != InjectionMode.CLASS_ONLY) {
- throw new UnsupportedOperationException("Not yet implemented");
- }
- classesToInject.put(
- proxyClassName,
+ classesToInject.add(
  cl -> {
  TypePool typePool = IndyModuleTypePool.get(cl, instrumentationModule);
  TypeDescription proxiedType = typePool.describe(classToProxy).resolve();
- return proxyFactory.generateProxy(proxiedType, proxyClassName).getBytes();
+ DynamicType.Unloaded<?> proxy = proxyFactory.generateProxy(proxiedType, proxyClassName);
+ return HelperClassDefinition.create(proxy, mode);
  });
  }
  }

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

@@ -10,6 +10,7 @@
 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;
@@ -90,11 +91,11 @@ static InstrumentationModuleClassLoader createInstrumentationModuleClassloader(
  }
 
  ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
- Map<String, ClassCopySource> injectedClasses =
+ Map<String, BytecodeWithUrl> injectedClasses =
  toInject.stream()
  .collect(
  Collectors.toMap(
- name -> name, name -> ClassCopySource.create(name, agentOrExtensionCl)));
+ name -> name, name -> BytecodeWithUrl.create(name, agentOrExtensionCl)));
 
  return new InstrumentationModuleClassLoader(
  instrumentedClassloader, agentOrExtensionCl, injectedClasses);

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

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.tooling.instrumentation.indy;
 
+import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
 import java.io.IOException;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
@@ -43,13 +44,13 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
  ClassLoader.registerAsParallelCapable();
  }
 
- private static final Map<String, ClassCopySource> ALWAYS_INJECTED_CLASSES =
+ private static final Map<String, BytecodeWithUrl> ALWAYS_INJECTED_CLASSES =
  Collections.singletonMap(
- LookupExposer.class.getName(), ClassCopySource.create(LookupExposer.class).cached());
+ LookupExposer.class.getName(), BytecodeWithUrl.create(LookupExposer.class).cached());
  private static final ProtectionDomain PROTECTION_DOMAIN = getProtectionDomain();
  private static final MethodHandle FIND_PACKAGE_METHOD = getFindPackageMethod();
 
- private final Map<String, ClassCopySource> additionalInjectedClasses;
+ private final Map<String, BytecodeWithUrl> additionalInjectedClasses;
  private final ClassLoader agentOrExtensionCl;
  private volatile MethodHandles.Lookup cachedLookup;
 
@@ -59,14 +60,14 @@ public class InstrumentationModuleClassLoader extends ClassLoader {
  public InstrumentationModuleClassLoader(
  ClassLoader instrumentedCl,
  ClassLoader agentOrExtensionCl,
- Map<String, ClassCopySource> injectedClasses) {
+ Map<String, BytecodeWithUrl> injectedClasses) {
  this(instrumentedCl, agentOrExtensionCl, injectedClasses, false);
  }
 
  InstrumentationModuleClassLoader(
  ClassLoader instrumentedCl,
  ClassLoader agentOrExtensionCl,
- Map<String, ClassCopySource> injectedClasses,
+ Map<String, BytecodeWithUrl> injectedClasses,
  boolean delegateAllToAgent) {
  // agent/extension-classloader is "main"-parent, but class lookup is overridden
  super(agentOrExtensionCl);
@@ -105,7 +106,7 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
 
  // This CL is self-first: Injected class are loaded BEFORE a parent lookup
  if (result == null) {
- ClassCopySource injected = getInjectedClass(name);
+ BytecodeWithUrl injected = getInjectedClass(name);
  if (injected != null) {
  byte[] bytecode =
  bytecodeOverride.get(name) != null
@@ -158,7 +159,7 @@ public URL getResource(String resourceName) {
  return super.getResource(resourceName);
  }
  // for classes use the same precedence as in loadClass
- ClassCopySource injected = getInjectedClass(className);
+ BytecodeWithUrl injected = getInjectedClass(className);
  if (injected != null) {
  return injected.getUrl();
  }
@@ -196,8 +197,8 @@ private static String resourceToClassName(String resourceName) {
  }
 
  @Nullable
- private ClassCopySource getInjectedClass(String name) {
- ClassCopySource alwaysInjected = ALWAYS_INJECTED_CLASSES.get(name);
+ private BytecodeWithUrl getInjectedClass(String name) {
+ BytecodeWithUrl alwaysInjected = ALWAYS_INJECTED_CLASSES.get(name);
  if (alwaysInjected != null) {
  return alwaysInjected;
  }

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

@@ -8,6 +8,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.dummies.Bar;
 import io.opentelemetry.javaagent.tooling.instrumentation.indy.dummies.Foo;
 import java.io.IOException;
@@ -36,9 +37,9 @@ class InstrumentationModuleClassLoaderTest {
 
  @Test
  void checkLookup() throws Throwable {
- Map<String, ClassCopySource> toInject = new HashMap<>();
- toInject.put(Foo.class.getName(), ClassCopySource.create(Foo.class));
- toInject.put(Bar.class.getName(), ClassCopySource.create(Bar.class));
+ Map<String, BytecodeWithUrl> toInject = new HashMap<>();
+ toInject.put(Foo.class.getName(), BytecodeWithUrl.create(Foo.class));
+ toInject.put(Bar.class.getName(), BytecodeWithUrl.create(Bar.class));
 
  ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null);
 
@@ -72,9 +73,9 @@ private static void lookupAndInvokeFoo(InstrumentationModuleClassLoader classLoa
 
  @Test
  void checkInjectedClassesHavePackage() throws Throwable {
- Map<String, ClassCopySource> toInject = new HashMap<>();
- toInject.put(A.class.getName(), ClassCopySource.create(A.class));
- toInject.put(B.class.getName(), ClassCopySource.create(B.class));
+ Map<String, BytecodeWithUrl> toInject = new HashMap<>();
+ toInject.put(A.class.getName(), BytecodeWithUrl.create(A.class));
+ toInject.put(B.class.getName(), BytecodeWithUrl.create(B.class));
  String packageName = A.class.getName().substring(0, A.class.getName().lastIndexOf('.'));
 
  ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null);
@@ -116,8 +117,8 @@ void checkClassLookupPrecedence(@TempDir Path tempDir) throws Exception {
  URLClassLoader moduleSourceCl = new URLClassLoader(new URL[] {moduleJar.toUri().toURL()}, null);
 
  try {
- Map<String, ClassCopySource> toInject = new HashMap<>();
- toInject.put(C.class.getName(), ClassCopySource.create(C.class.getName(), moduleSourceCl));
+ Map<String, BytecodeWithUrl> toInject = new HashMap<>();
+ toInject.put(C.class.getName(), BytecodeWithUrl.create(C.class.getName(), moduleSourceCl));
 
  InstrumentationModuleClassLoader moduleCl =
  new InstrumentationModuleClassLoader(appCl, agentCl, toInject, true);