From 9b38c2399d7d9a26ba04be970dd3d49e584bff30 Mon Sep 17 00:00:00 2001 From: Jonas Kunz Date: Fri, 2 Feb 2024 14:35:54 +0100 Subject: [PATCH] Allow multiple invokedynamic InstrumentationModules to share classloaders (#10015) --- .../AbstractAwsSdkInstrumentationModule.java | 12 +- .../v2_2/AwsSdkInstrumentationModule.java | 10 ++ .../ExperimentalInstrumentationModule.java | 12 ++ .../InstrumentationModuleInstaller.java | 9 +- .../instrumentation/MuzzleMatcher.java | 18 +- .../indy/ClassInjectorImpl.java | 7 +- .../instrumentation/indy/IndyBootstrap.java | 6 +- .../indy/IndyModuleRegistry.java | 167 +++++++++--------- .../indy/IndyModuleTypePool.java | 34 ---- .../InstrumentationModuleClassLoader.java | 134 +++++++++++--- .../tooling/util/ClassLoaderMap.java | 6 + .../tooling/util/ClassLoaderValue.java | 7 + .../InstrumentationModuleClassLoaderTest.java | 14 +- .../tooling/util/ClassLoaderValueTest.java | 19 +- 14 files changed, 285 insertions(+), 170 deletions(-) delete mode 100644 javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleTypePool.java diff --git a/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AbstractAwsSdkInstrumentationModule.java b/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AbstractAwsSdkInstrumentationModule.java index 5dc5f2ab9d4f..3138429d4fc9 100644 --- a/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AbstractAwsSdkInstrumentationModule.java +++ b/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 diff --git a/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AwsSdkInstrumentationModule.java b/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AwsSdkInstrumentationModule.java index a6bb33a49b77..18d9ca5bc0cb 100644 --- a/instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AwsSdkInstrumentationModule.java +++ b/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. diff --git a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java b/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java index b4e90e5ab91f..b2378d752434 100644 --- a/javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/internal/ExperimentalInstrumentationModule.java +++ b/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 injectedClassNames() { return emptyList(); } + + /** + * By default every InstrumentationModule is loaded by an isolated classloader, even if multiple + * modules instrument the same application classloader. + * + *

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(); + } } diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java index 147fd73b0b3c..b4ea4b948639 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java +++ b/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); diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/MuzzleMatcher.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/MuzzleMatcher.java index 03488bf921a3..fd9f84924e5f 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/MuzzleMatcher.java +++ b/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) { diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/ClassInjectorImpl.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/ClassInjectorImpl.java index bc365b17eec4..ad901577ebc0 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/ClassInjectorImpl.java +++ b/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); diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyBootstrap.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyBootstrap.java index 4574caa6ed8e..9e23725cf8d3 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyBootstrap.java +++ b/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); diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleRegistry.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleRegistry.java index b2c56763cfa4..7eaf7f3511b6 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleRegistry.java +++ b/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 modulesByName = + private static final ConcurrentHashMap 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}. + * + *

The keys of the contained map are the instrumentation module group names, see {@link + * ExperimentalInstrumentationModule#getModuleGroup()}; */ - private static final ConcurrentHashMap< - InstrumentationModule, - Cache>> - instrumentationClassloaders = new ConcurrentHashMap<>(); - - public static InstrumentationModuleClassLoader getInstrumentationClassloader( - String moduleClassName, ClassLoader instrumentedClassloader) { - InstrumentationModule instrumentationModule = modulesByName.get(moduleClassName); + private static final ClassLoaderValue> + 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> cacheForModule = - instrumentationClassloaders.computeIfAbsent(module, (k) -> Cache.weak()); + Map loadersByGroupName = + instrumentationClassLoaders.get(instrumentedClassLoader); - instrumentedClassloader = maskNullClassLoader(instrumentedClassloader); - WeakReference 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 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 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 getModuleAdviceNames(InstrumentationModule module) { - Set adviceNames = new HashSet<>(); - TypeTransformer nameCollector = - new TypeTransformer() { - @Override - public void applyAdviceToMethod( - ElementMatcher 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(); } } diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleTypePool.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleTypePool.java deleted file mode 100644 index 7cf63528a870..000000000000 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/IndyModuleTypePool.java +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.tooling.muzzle.AgentTooling; -import net.bytebuddy.pool.TypePool; - -public class IndyModuleTypePool { - - private IndyModuleTypePool() {} - - /** - * Provides a {@link TypePool} which has the same lookup rules for {@link - * net.bytebuddy.description.type.TypeDescription}s as {@link InstrumentationModuleClassLoader} - * have for classes. - * - * @param instrumentedCl the classloader being instrumented (e.g. for which the {@link - * InstrumentationModuleClassLoader} is being created). - * @param module the {@link InstrumentationModule} performing the instrumentation - * @return the type pool, must not be cached! - */ - public static TypePool get(ClassLoader instrumentedCl, InstrumentationModule module) { - // TODO: this implementation doesn't allow caching the returned pool and its types - // This could be improved by implementing a custom TypePool instead, which delegates to parent - // TypePools and mirrors the delegation model of the InstrumentationModuleClassLoader - InstrumentationModuleClassLoader dummyCl = - IndyModuleRegistry.createInstrumentationModuleClassloader(module, instrumentedCl); - return TypePool.Default.of(AgentTooling.locationStrategy().classFileLocator(dummyCl)); - } -} diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java index 4cf5859fd5f2..fada8911eb71 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java +++ b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoader.java @@ -5,7 +5,12 @@ package io.opentelemetry.javaagent.tooling.instrumentation.indy; +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.io.IOException; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; @@ -15,28 +20,35 @@ import java.security.ProtectionDomain; import java.util.Collections; import java.util.Enumeration; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Collectors; import javax.annotation.Nullable; +import net.bytebuddy.agent.builder.AgentBuilder; +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.matcher.ElementMatcher; +import net.bytebuddy.matcher.StringMatcher; /** - * Classloader used to load the helper classes from {@link + * Class loader used to load the helper classes from {@link * io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule}s, so that those * classes have access to both the agent/extension classes and the instrumented application classes. * - *

This classloader implements the following classloading delegation strategy: + *

This class loader implements the following classloading delegation strategy: * *

* - *

In addition, this classloader ensures that the lookup of corresponding .class resources follow - * the same delegation strategy, so that bytecode inspection tools work correctly. + *

In addition, this class loader ensures that the lookup of corresponding .class resources + * follow the same delegation strategy, so that bytecode inspection tools work correctly. */ public class InstrumentationModuleClassLoader extends ClassLoader { @@ -44,6 +56,8 @@ public class InstrumentationModuleClassLoader extends ClassLoader { ClassLoader.registerAsParallelCapable(); } + private static final ClassLoader BOOT_LOADER = new ClassLoader() {}; + private static final Map ALWAYS_INJECTED_CLASSES = Collections.singletonMap( LookupExposer.class.getName(), BytecodeWithUrl.create(LookupExposer.class).cached()); @@ -54,33 +68,42 @@ public class InstrumentationModuleClassLoader extends ClassLoader { private final ClassLoader agentOrExtensionCl; private volatile MethodHandles.Lookup cachedLookup; - private final ClassLoader instrumentedCl; - private final boolean delegateAllToAgent; + @Nullable private final ClassLoader instrumentedCl; + + /** + * Only class names matching this matcher will be attempted to be loaded from the {@link + * #agentOrExtensionCl}. If a class is requested and it does not match this matcher, the lookup in + * {@link #agentOrExtensionCl} will be skipped. + */ + private final ElementMatcher agentClassNamesMatcher; + + private final Set installedModules; public InstrumentationModuleClassLoader( - ClassLoader instrumentedCl, - ClassLoader agentOrExtensionCl, - Map injectedClasses) { - this(instrumentedCl, agentOrExtensionCl, injectedClasses, false); + ClassLoader instrumentedCl, ClassLoader agentOrExtensionCl) { + this( + instrumentedCl, + agentOrExtensionCl, + new StringMatcher("io.opentelemetry.javaagent", StringMatcher.Mode.STARTS_WITH)); } InstrumentationModuleClassLoader( - ClassLoader instrumentedCl, + @Nullable ClassLoader instrumentedCl, ClassLoader agentOrExtensionCl, - Map injectedClasses, - boolean delegateAllToAgent) { - // agent/extension-classloader is "main"-parent, but class lookup is overridden + ElementMatcher classesToLoadFromAgentOrExtensionCl) { + // agent/extension-class loader is "main"-parent, but class lookup is overridden super(agentOrExtensionCl); - additionalInjectedClasses = injectedClasses; + additionalInjectedClasses = new ConcurrentHashMap<>(); + installedModules = Collections.newSetFromMap(new ConcurrentHashMap<>()); this.agentOrExtensionCl = agentOrExtensionCl; this.instrumentedCl = instrumentedCl; - this.delegateAllToAgent = delegateAllToAgent; + this.agentClassNamesMatcher = classesToLoadFromAgentOrExtensionCl; } /** - * Provides a Lookup within this classloader. See {@link LookupExposer} for the details. + * Provides a Lookup within this class loader. See {@link LookupExposer} for the details. * - * @return a lookup capable of accessing public types in this classloader + * @return a lookup capable of accessing public types in this class loader */ public MethodHandles.Lookup getLookup() { if (cachedLookup == null) { @@ -96,6 +119,62 @@ public MethodHandles.Lookup getLookup() { return cachedLookup; } + public synchronized void installModule(InstrumentationModule module) { + if (module.getClass().getClassLoader() != agentOrExtensionCl) { + throw new IllegalArgumentException( + module.getClass().getName() + " is not loaded by " + agentOrExtensionCl); + } + if (!installedModules.add(module)) { + return; + } + Map classesToInject = + getClassesToInject(module).stream() + .collect( + Collectors.toMap( + className -> className, + className -> BytecodeWithUrl.create(className, agentOrExtensionCl))); + installInjectedClasses(classesToInject); + } + + public synchronized boolean hasModuleInstalled(InstrumentationModule module) { + return installedModules.contains(module); + } + + // Visible for testing + synchronized void installInjectedClasses(Map classesToInject) { + classesToInject.forEach(additionalInjectedClasses::putIfAbsent); + } + + private static Set getClassesToInject(InstrumentationModule module) { + Set 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()); + } + return toInject; + } + + private static Set getModuleAdviceNames(InstrumentationModule module) { + Set adviceNames = new HashSet<>(); + TypeTransformer nameCollector = + new TypeTransformer() { + @Override + public void applyAdviceToMethod( + ElementMatcher methodMatcher, String adviceClassName) { + adviceNames.add(adviceClassName); + } + + @Override + public void applyTransformer(AgentBuilder.Transformer transformer) {} + }; + for (TypeInstrumentation instr : module.typeInstrumentations()) { + instr.transform(nameCollector); + } + return adviceNames; + } + public static final Map bytecodeOverride = new ConcurrentHashMap<>(); @Override @@ -140,12 +219,12 @@ protected Class loadClass(String name, boolean resolve) throws ClassNotFoundE } private boolean shouldLoadFromAgent(String dotClassName) { - return delegateAllToAgent || dotClassName.startsWith("io.opentelemetry.javaagent"); + return agentClassNamesMatcher.matches(dotClassName); } - private static Class tryLoad(ClassLoader cl, String name) { + private static Class tryLoad(@Nullable ClassLoader cl, String name) { try { - return cl.loadClass(name); + return Class.forName(name, false, cl); } catch (ClassNotFoundException e) { return null; } @@ -155,7 +234,7 @@ private static Class tryLoad(ClassLoader cl, String name) { public URL getResource(String resourceName) { String className = resourceToClassName(resourceName); if (className == null) { - // delegate to just the default parent (the agent classloader) + // delegate to just the default parent (the agent class loader) return super.getResource(resourceName); } // for classes use the same precedence as in loadClass @@ -167,7 +246,12 @@ public URL getResource(String resourceName) { if (fromAgentCl != null) { return fromAgentCl; } - return instrumentedCl.getResource(resourceName); + + if (instrumentedCl != null) { + return instrumentedCl.getResource(resourceName); + } else { + return BOOT_LOADER.getResource(resourceName); + } } @Override diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java index 9da9c3ac6fa3..aedadef84eac 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java +++ b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderMap.java @@ -11,6 +11,7 @@ import java.util.Collections; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; import net.bytebuddy.ByteBuddy; import net.bytebuddy.description.modifier.Ownership; import net.bytebuddy.description.modifier.Visibility; @@ -27,6 +28,11 @@ public static void put(ClassLoader classLoader, Object key, Object value) { getClassLoaderData(classLoader, true).put(key, value); } + public static Object computeIfAbsent( + ClassLoader classLoader, Object key, Supplier value) { + return getClassLoaderData(classLoader, true).computeIfAbsent(key, unused -> value.get()); + } + private static Map getClassLoaderData( ClassLoader classLoader, boolean initialize) { classLoader = maskNullClassLoader(classLoader); diff --git a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java index 04df8da95615..6288d8f97217 100644 --- a/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java +++ b/javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValue.java @@ -5,6 +5,8 @@ package io.opentelemetry.javaagent.tooling.util; +import java.util.function.Supplier; + /** * Associate value with a class loader. Added value will behave as if it was stored in a field in * the class loader object, meaning that the value can be garbage collected once the class loader is @@ -21,4 +23,9 @@ public T get(ClassLoader classLoader) { public void put(ClassLoader classLoader, T value) { ClassLoaderMap.put(classLoader, this, value); } + + @SuppressWarnings("unchecked") + public T computeIfAbsent(ClassLoader classLoader, Supplier value) { + return (T) ClassLoaderMap.computeIfAbsent(classLoader, this, value); + } } diff --git a/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoaderTest.java b/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoaderTest.java index 7fbfd45e2826..5f242fe719c7 100644 --- a/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoaderTest.java +++ b/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/instrumentation/indy/InstrumentationModuleClassLoaderTest.java @@ -29,6 +29,7 @@ import net.bytebuddy.ByteBuddy; import net.bytebuddy.dynamic.ClassFileLocator; import net.bytebuddy.implementation.FixedValue; +import net.bytebuddy.matcher.ElementMatchers; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -44,9 +45,12 @@ void checkLookup() throws Throwable { ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null); InstrumentationModuleClassLoader m1 = - new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true); + new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any()); + m1.installInjectedClasses(toInject); + InstrumentationModuleClassLoader m2 = - new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true); + new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any()); + m2.installInjectedClasses(toInject); // MethodHandles.publicLookup() always succeeds on the first invocation lookupAndInvokeFoo(m1); @@ -80,7 +84,8 @@ void checkInjectedClassesHavePackage() throws Throwable { ClassLoader dummyParent = new URLClassLoader(new URL[] {}, null); InstrumentationModuleClassLoader m1 = - new InstrumentationModuleClassLoader(dummyParent, dummyParent, toInject, true); + new InstrumentationModuleClassLoader(dummyParent, dummyParent, ElementMatchers.any()); + m1.installInjectedClasses(toInject); Class injected = Class.forName(A.class.getName(), true, m1); // inject two classes from the same package to trigger errors if we try to redefine the package @@ -121,7 +126,8 @@ void checkClassLookupPrecedence(@TempDir Path tempDir) throws Exception { toInject.put(C.class.getName(), BytecodeWithUrl.create(C.class.getName(), moduleSourceCl)); InstrumentationModuleClassLoader moduleCl = - new InstrumentationModuleClassLoader(appCl, agentCl, toInject, true); + new InstrumentationModuleClassLoader(appCl, agentCl, ElementMatchers.any()); + moduleCl.installInjectedClasses(toInject); // Verify precedence for classloading Class clA = moduleCl.loadClass(A.class.getName()); diff --git a/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java b/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java index b4a3cac845e5..db8ba91f92e6 100644 --- a/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java +++ b/javaagent-tooling/src/test/java/io/opentelemetry/javaagent/tooling/util/ClassLoaderValueTest.java @@ -22,9 +22,22 @@ void testValue() { } void testClassLoader(ClassLoader classLoader) { - ClassLoaderValue classLoaderValue = new ClassLoaderValue<>(); - classLoaderValue.put(classLoader, "value"); - assertThat(classLoaderValue.get(classLoader)).isEqualTo("value"); + ClassLoaderValue value1 = new ClassLoaderValue<>(); + value1.put(classLoader, "value"); + assertThat(value1.get(classLoader)).isEqualTo("value"); + + ClassLoaderValue value2 = new ClassLoaderValue<>(); + String value = "value"; + String ret1 = value2.computeIfAbsent(classLoader, () -> value); + String ret2 = + value2.computeIfAbsent( + classLoader, + () -> { + throw new IllegalStateException("Shouldn't be invoked"); + }); + assertThat(ret1).isSameAs(value); + assertThat(ret2).isSameAs(value); + assertThat(value2.get(classLoader)).isSameAs(value); } @Test