Skip to content

Commit

Permalink
Allow injection of helper bytecode as resources (#9752)
Browse files Browse the repository at this point in the history
JonasKunz authored Nov 6, 2023
1 parent 431c544 commit 6eb8ae1
Showing 14 changed files with 500 additions and 210 deletions.
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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);
});
}
}
Original file line number Diff line number Diff line change
@@ -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);
Original file line number Diff line number Diff line change
@@ -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;
}
Original file line number Diff line number Diff line change
@@ -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);
Loading

0 comments on commit 6eb8ae1

Please sign in to comment.