Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow multiple invokedynamic InstrumentationModules to share classloaders #10015

Merged
merged 17 commits into from
Feb 2, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Switched module registration to happen on first class match instead o…
…f on invokedynamic linkage.
  • Loading branch information
JonasKunz committed Dec 5, 2023
commit e7c265440a8abb575f47411c4be13b49495f1bb6
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,6 @@ private AgentBuilder installIndyModule(
injectedHelperClassNames = Collections.emptyList();
}

IndyModuleRegistry.registerIndyModule(instrumentationModule);

ClassInjectorImpl injectedClassesCollector = new ClassInjectorImpl(instrumentationModule);
if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
((ExperimentalInstrumentationModule) instrumentationModule)
Expand Down Expand Up @@ -157,6 +155,8 @@ private AgentBuilder installIndyModule(
// 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.registerModuleOnMatch(instrumentationModule, extendableAgentBuilder);
extendableAgentBuilder = contextProvider.injectHelperClasses(extendableAgentBuilder);
IndyTypeTransformerImpl typeTransformer =
new IndyTypeTransformerImpl(extendableAgentBuilder, instrumentationModule);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ public boolean matches(
classLoader = Utils.getBootstrapProxy();
} else if (instrumentationModule.isIndyModule()) {
classLoader =
IndyModuleRegistry.getInstrumentationClassloader(
instrumentationModule.getClass().getName(), classLoader);
IndyModuleRegistry.createInstrumentationClassloaderWithoutRegistration(
instrumentationModule, classLoader);
}
return matchCache.computeIfAbsent(classLoader, this::doesMatch);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,11 @@
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
import java.lang.ref.WeakReference;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import net.bytebuddy.agent.builder.AgentBuilder;

public class IndyModuleRegistry {

Expand All @@ -18,6 +22,66 @@ private IndyModuleRegistry() {}
private static final ConcurrentHashMap<String, InstrumentationModule> modulesByName =
new ConcurrentHashMap<>();

private static class ModuleRegistration {

private final WeakReference<ClassLoader> instrumentedCl;

private final Set<InstrumentationModule> registeredModules =
Collections.newSetFromMap(new ConcurrentHashMap<>());

private WeakReference<InstrumentationModuleClassLoader> modulesClassloader = null;

private ModuleRegistration(ClassLoader instrumentedCl) {
this.instrumentedCl = new WeakReference<>(instrumentedCl);
}

private boolean isRegistered(InstrumentationModule module) {
return registeredModules.contains(module);
}

public void register(InstrumentationModule module) {
if (registeredModules.contains(module)) {
return;
}
synchronized (this) {
if (registeredModules.add(module)) {
if (modulesClassloader != null) {
InstrumentationModuleClassLoader modulesCl = modulesClassloader.get();
if (modulesCl != null) {
// Classloader has already been created, at the module to it
modulesCl.installModule(module);
}
}
}
}
}

public InstrumentationModuleClassLoader getOrCreateModulesClassloader() {
if (modulesClassloader != null) {
InstrumentationModuleClassLoader modulesCl = modulesClassloader.get();
if (modulesCl != null) {
return modulesCl;
}
}
synchronized (this) {
if (modulesClassloader != null) {
InstrumentationModuleClassLoader modulesCl = modulesClassloader.get();
if (modulesCl != null) {
return modulesCl;
}
}
ClassLoader instrumented = instrumentedCl.get();
if (instrumented == null) {
throw new IllegalStateException("Target classloader has been GCed!");
}
InstrumentationModuleClassLoader modulesCl =
createInstrumentationModuleClassloader(registeredModules, instrumented);
modulesClassloader = new WeakReference<>(modulesCl);
return modulesCl;
}
}
}

/**
* Weakly references the {@link InstrumentationModuleClassLoader}s for a given application
* classloader. We only store weak references to make sure we don't prevent application
Expand All @@ -27,8 +91,7 @@ private IndyModuleRegistry() {}
* <p>The keys of this map are the instrumentation module group names, see {@link
* ExperimentalInstrumentationModule#getModuleGroup()};
*/
private static final ConcurrentHashMap<
String, Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>>>
private static final ConcurrentHashMap<String, Cache<ClassLoader, ModuleRegistration>>
instrumentationClassloaders = new ConcurrentHashMap<>();

public static InstrumentationModuleClassLoader getInstrumentationClassloader(
Expand All @@ -45,51 +108,100 @@ private static synchronized InstrumentationModuleClassLoader getInstrumentationC
InstrumentationModule module, ClassLoader instrumentedClassloader) {

String groupName = getModuleGroup(module);
Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>> cacheForGroup =
instrumentationClassloaders.computeIfAbsent(groupName, (k) -> Cache.weak());
Cache<ClassLoader, ModuleRegistration> cacheForGroup =
instrumentationClassloaders.get(groupName);
if (cacheForGroup == null) {
throw new IllegalArgumentException(module + " has not been registered yet");
}

instrumentedClassloader = maskNullClassLoader(instrumentedClassloader);
WeakReference<InstrumentationModuleClassLoader> cached =
cacheForGroup.get(instrumentedClassloader);
if (cached != null) {
InstrumentationModuleClassLoader cachedCl = cached.get();
if (cachedCl != null) {
cachedCl.installModule(module);
return cachedCl;
ModuleRegistration registrations = cacheForGroup.get(instrumentedClassloader);
if (registrations == null || !registrations.isRegistered(module)) {
throw new IllegalArgumentException(
module + " has not been registered for Classloader " + instrumentedClassloader);
}
return registrations.getOrCreateModulesClassloader();
}

/**
* Returns a classloader which behaves exactly as if you would first register this module and
* afterwards call {@link #getInstrumentationClassloader(String, ClassLoader)}.
*
* <p>The difference is that the registration actually does not happen, therefore this call is
* side-effect free.
*
* <p>The returned classloader will load the provided module AND all other previously registered
* for the given instrumentedClassloader instrumentation modules with the same group.
*/
public static InstrumentationModuleClassLoader
createInstrumentationClassloaderWithoutRegistration(
InstrumentationModule module, ClassLoader instrumentedClassloader) {
// TODO: remove this method and replace usages with a custom TypePool implementation instead
String groupName = getModuleGroup(module);

instrumentedClassloader = maskNullClassLoader(instrumentedClassloader);

Set<InstrumentationModule> modules = new HashSet<>();
modules.add(module);
Cache<ClassLoader, ModuleRegistration> cacheForGroup =
instrumentationClassloaders.get(groupName);
if (cacheForGroup != null) {
ModuleRegistration registration = cacheForGroup.get(instrumentedClassloader);
if (registration != null) {
modules.addAll(registration.registeredModules);
}
}
// 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);
cacheForGroup.put(instrumentedClassloader, new WeakReference<>(created));
return created;

return createInstrumentationModuleClassloader(modules, instrumentedClassloader);
}

private static final ClassLoader BOOT_LOADER = new ClassLoader() {};
private static final ClassLoader BOOT_LOADER = new ClassLoader(null) {};

private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
return classLoader == null ? BOOT_LOADER : classLoader;
}

static InstrumentationModuleClassLoader createInstrumentationModuleClassloader(
InstrumentationModule module, ClassLoader instrumentedClassloader) {
ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
Set<InstrumentationModule> modules, ClassLoader instrumentedClassloader) {
if (modules.isEmpty()) {
throw new IllegalArgumentException("Must provide at least one module!");
}
ClassLoader agentOrExtensionCl = modules.iterator().next().getClass().getClassLoader();
InstrumentationModuleClassLoader moduleCl =
new InstrumentationModuleClassLoader(instrumentedClassloader, agentOrExtensionCl);
moduleCl.installModule(module);

modules.forEach(moduleCl::installModule);
return moduleCl;
}

public static void registerIndyModule(InstrumentationModule module) {
public static AgentBuilder.Identified.Extendable registerModuleOnMatch(
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 = modulesByName.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) -> {
// this causes the classloader to be created and kept alive when the first class matches
registerModuleForClassloader(module, classLoader);
return builder;
});
}

private static void registerModuleForClassloader(
InstrumentationModule module, ClassLoader classLoader) {
String groupName = getModuleGroup(module);
Cache<ClassLoader, ModuleRegistration> cacheForGroup =
instrumentationClassloaders.computeIfAbsent(groupName, (k) -> Cache.weak());
classLoader = maskNullClassLoader(classLoader);
ModuleRegistration registrations =
cacheForGroup.computeIfAbsent(classLoader, ModuleRegistration::new);
registrations.register(module);
}

private static String getModuleGroup(InstrumentationModule module) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ public static TypePool get(ClassLoader instrumentedCl, InstrumentationModule mod
// 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);
IndyModuleRegistry.createInstrumentationClassloaderWithoutRegistration(
module, instrumentedCl);
return TypePool.Default.of(AgentTooling.locationStrategy().classFileLocator(dummyCl));
}
}