Skip to content

Commit 980d8ea

Browse files
authored
Allow multiple invokedynamic InstrumentationModules to share classloaders (#10015)
1 parent 147b3e8 commit 980d8ea

File tree

14 files changed

+285
-170
lines changed

14 files changed

+285
-170
lines changed

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AbstractAwsSdkInstrumentationModule.java

+7-5
Original file line numberDiff line numberDiff line change
@@ -12,24 +12,26 @@
1212
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
1313
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
1414
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
15+
import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
1516
import java.util.List;
1617
import net.bytebuddy.description.type.TypeDescription;
1718
import net.bytebuddy.matcher.ElementMatcher;
1819

19-
abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule {
20+
abstract class AbstractAwsSdkInstrumentationModule extends InstrumentationModule
21+
implements ExperimentalInstrumentationModule {
2022

2123
protected AbstractAwsSdkInstrumentationModule(String additionalInstrumentationName) {
2224
super("aws-sdk", "aws-sdk-2.2", additionalInstrumentationName);
2325
}
2426

2527
@Override
26-
public boolean isHelperClass(String className) {
27-
return className.startsWith("io.opentelemetry.contrib.awsxray.");
28+
public String getModuleGroup() {
29+
return "aws-sdk-v2";
2830
}
2931

3032
@Override
31-
public boolean isIndyModule() {
32-
return false;
33+
public boolean isHelperClass(String className) {
34+
return className.startsWith("io.opentelemetry.contrib.awsxray.");
3335
}
3436

3537
@Override

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/AwsSdkInstrumentationModule.java

+10
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import io.opentelemetry.javaagent.extension.instrumentation.HelperResourceBuilder;
1111
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
1212
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
13+
import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ClassInjector;
14+
import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
1315

1416
@AutoService(InstrumentationModule.class)
1517
public class AwsSdkInstrumentationModule extends AbstractAwsSdkInstrumentationModule {
@@ -26,6 +28,14 @@ public void registerHelperResources(HelperResourceBuilder helperResourceBuilder)
2628
helperResourceBuilder.register("software/amazon/awssdk/global/handlers/execution.interceptors");
2729
}
2830

31+
@Override
32+
public void injectClasses(ClassInjector injector) {
33+
injector
34+
.proxyBuilder(
35+
"io.opentelemetry.instrumentation.awssdk.v2_2.autoconfigure.TracingExecutionInterceptor")
36+
.inject(InjectionMode.CLASS_ONLY);
37+
}
38+
2939
@Override
3040
void doTransform(TypeTransformer transformer) {
3141
// Nothing to transform, this type instrumentation is only used for injecting resources.

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

+12
Original file line numberDiff line numberDiff line change
@@ -36,4 +36,16 @@ default void injectClasses(ClassInjector injector) {}
3636
default List<String> injectedClassNames() {
3737
return emptyList();
3838
}
39+
40+
/**
41+
* By default every InstrumentationModule is loaded by an isolated classloader, even if multiple
42+
* modules instrument the same application classloader.
43+
*
44+
* <p>Sometimes this is not desired, e.g. when instrumenting modular libraries such as the AWS
45+
* SDK. In such cases the {@link InstrumentationModule}s which want to share a classloader can
46+
* return the same group name from this method.
47+
*/
48+
default String getModuleGroup() {
49+
return getClass().getName();
50+
}
3951
}

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

+5-4
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,6 @@ private AgentBuilder installIndyModule(
109109
injectedHelperClassNames = Collections.emptyList();
110110
}
111111

112-
IndyModuleRegistry.registerIndyModule(instrumentationModule);
113-
114112
ClassInjectorImpl injectedClassesCollector = new ClassInjectorImpl(instrumentationModule);
115113
if (instrumentationModule instanceof ExperimentalInstrumentationModule) {
116114
((ExperimentalInstrumentationModule) instrumentationModule)
@@ -149,14 +147,17 @@ private AgentBuilder installIndyModule(
149147
AgentBuilder.Identified.Extendable extendableAgentBuilder =
150148
setTypeMatcher(agentBuilder, instrumentationModule, typeInstrumentation)
151149
.and(muzzleMatcher)
152-
.transform(new PatchByteCodeVersionTransformer())
153-
.transform(helperInjector);
150+
.transform(new PatchByteCodeVersionTransformer());
154151

155152
// TODO (Jonas): we are not calling
156153
// contextProvider.rewriteVirtualFieldsCalls(extendableAgentBuilder) anymore
157154
// As a result the advices should store `VirtualFields` as static variables instead of having
158155
// the lookup inline
159156
// We need to update our documentation on that
157+
extendableAgentBuilder =
158+
IndyModuleRegistry.initializeModuleLoaderOnMatch(
159+
instrumentationModule, extendableAgentBuilder);
160+
extendableAgentBuilder = extendableAgentBuilder.transform(helperInjector);
160161
extendableAgentBuilder = contextProvider.injectHelperClasses(extendableAgentBuilder);
161162
IndyTypeTransformerImpl typeTransformer =
162163
new IndyTypeTransformerImpl(extendableAgentBuilder, instrumentationModule);

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

+13-5
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import io.opentelemetry.javaagent.tooling.Utils;
1616
import io.opentelemetry.javaagent.tooling.config.AgentConfig;
1717
import io.opentelemetry.javaagent.tooling.instrumentation.indy.IndyModuleRegistry;
18+
import io.opentelemetry.javaagent.tooling.instrumentation.indy.InstrumentationModuleClassLoader;
1819
import io.opentelemetry.javaagent.tooling.muzzle.Mismatch;
1920
import io.opentelemetry.javaagent.tooling.muzzle.ReferenceMatcher;
2021
import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
@@ -61,12 +62,19 @@ public boolean matches(
6162
ProtectionDomain protectionDomain) {
6263
if (classLoader == BOOTSTRAP_LOADER) {
6364
classLoader = Utils.getBootstrapProxy();
64-
} else if (instrumentationModule.isIndyModule()) {
65-
classLoader =
66-
IndyModuleRegistry.getInstrumentationClassloader(
67-
instrumentationModule.getClass().getName(), classLoader);
6865
}
69-
return matchCache.computeIfAbsent(classLoader, this::doesMatch);
66+
if (instrumentationModule.isIndyModule()) {
67+
return matchCache.computeIfAbsent(
68+
classLoader,
69+
cl -> {
70+
InstrumentationModuleClassLoader moduleCl =
71+
IndyModuleRegistry.createInstrumentationClassLoaderWithoutRegistration(
72+
instrumentationModule, cl);
73+
return doesMatch(moduleCl);
74+
});
75+
} else {
76+
return matchCache.computeIfAbsent(classLoader, this::doesMatch);
77+
}
7078
}
7179

7280
private boolean doesMatch(ClassLoader classLoader) {

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
1111
import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ProxyInjectionBuilder;
1212
import io.opentelemetry.javaagent.tooling.HelperClassDefinition;
13+
import io.opentelemetry.javaagent.tooling.muzzle.AgentTooling;
1314
import java.util.ArrayList;
1415
import java.util.List;
1516
import java.util.function.Function;
@@ -57,7 +58,11 @@ private class ProxyBuilder implements ProxyInjectionBuilder {
5758
public void inject(InjectionMode mode) {
5859
classesToInject.add(
5960
cl -> {
60-
TypePool typePool = IndyModuleTypePool.get(cl, instrumentationModule);
61+
InstrumentationModuleClassLoader moduleCl =
62+
IndyModuleRegistry.getInstrumentationClassLoader(instrumentationModule, cl);
63+
TypePool typePool =
64+
AgentTooling.poolStrategy()
65+
.typePool(AgentTooling.locationStrategy().classFileLocator(moduleCl), moduleCl);
6166
TypeDescription proxiedType = typePool.describe(classToProxy).resolve();
6267
DynamicType.Unloaded<?> proxy = proxyFactory.generateProxy(proxiedType, proxyClassName);
6368
return HelperClassDefinition.create(proxy, mode);

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939
* ↑ └───────── IndyBootstrapDispatcher ─ ↑ ──→ └────────────── {@link IndyBootstrap#bootstrap}
4040
* Ext/Platform CL ↑ │ ╷
4141
* ↑ ╷ │ ↓
42-
* System CL ╷ │ {@link IndyModuleRegistry#getInstrumentationClassloader(String, ClassLoader)}
42+
* System CL ╷ │ {@link IndyModuleRegistry#getInstrumentationClassLoader(String, ClassLoader)}
4343
* ↑ ╷ │ ╷
4444
* Common linking of CallSite │ ╷
4545
* ↑ ↑ (on first invocation) │ ╷
@@ -171,7 +171,7 @@ private static ConstantCallSite bootstrapAdvice(
171171
}
172172

173173
InstrumentationModuleClassLoader instrumentationClassloader =
174-
IndyModuleRegistry.getInstrumentationClassloader(
174+
IndyModuleRegistry.getInstrumentationClassLoader(
175175
moduleClassName, lookup.lookupClass().getClassLoader());
176176

177177
// Advices are not inlined. They are loaded as normal classes by the
@@ -207,7 +207,7 @@ private static ConstantCallSite bootstrapProxyMethod(
207207
String methodKind)
208208
throws NoSuchMethodException, IllegalAccessException, ClassNotFoundException {
209209
InstrumentationModuleClassLoader instrumentationClassloader =
210-
IndyModuleRegistry.getInstrumentationClassloader(
210+
IndyModuleRegistry.getInstrumentationClassLoader(
211211
moduleClassName, lookup.lookupClass().getClassLoader());
212212

213213
Class<?> proxiedClass = instrumentationClassloader.loadClass(proxyClassName);

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

+81-86
Original file line numberDiff line numberDiff line change
@@ -5,129 +5,124 @@
55

66
package io.opentelemetry.javaagent.tooling.instrumentation.indy;
77

8-
import io.opentelemetry.instrumentation.api.internal.cache.Cache;
98
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
10-
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
11-
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
129
import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
13-
import io.opentelemetry.javaagent.tooling.BytecodeWithUrl;
14-
import io.opentelemetry.javaagent.tooling.muzzle.InstrumentationModuleMuzzle;
15-
import java.lang.ref.WeakReference;
16-
import java.util.HashSet;
10+
import io.opentelemetry.javaagent.tooling.util.ClassLoaderValue;
1711
import java.util.Map;
18-
import java.util.Set;
1912
import java.util.concurrent.ConcurrentHashMap;
20-
import java.util.stream.Collectors;
2113
import net.bytebuddy.agent.builder.AgentBuilder;
22-
import net.bytebuddy.description.method.MethodDescription;
23-
import net.bytebuddy.matcher.ElementMatcher;
2414

2515
public class IndyModuleRegistry {
2616

2717
private IndyModuleRegistry() {}
2818

29-
private static final ConcurrentHashMap<String, InstrumentationModule> modulesByName =
19+
private static final ConcurrentHashMap<String, InstrumentationModule> modulesByClassName =
3020
new ConcurrentHashMap<>();
3121

3222
/**
33-
* Weakly references the {@link InstrumentationModuleClassLoader}s for a given application
34-
* classloader. We only store weak references to make sure we don't prevent application
35-
* classloaders from being GCed. The application classloaders will strongly reference the {@link
36-
* InstrumentationModuleClassLoader} through the invokedynamic callsites.
23+
* Weakly references the {@link InstrumentationModuleClassLoader}s for a given application class
24+
* loader. The {@link InstrumentationModuleClassLoader} are kept alive by a strong reference from
25+
* the instrumented class loader realized via {@link ClassLoaderValue}.
26+
*
27+
* <p>The keys of the contained map are the instrumentation module group names, see {@link
28+
* ExperimentalInstrumentationModule#getModuleGroup()};
3729
*/
38-
private static final ConcurrentHashMap<
39-
InstrumentationModule,
40-
Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>>>
41-
instrumentationClassloaders = new ConcurrentHashMap<>();
42-
43-
public static InstrumentationModuleClassLoader getInstrumentationClassloader(
44-
String moduleClassName, ClassLoader instrumentedClassloader) {
45-
InstrumentationModule instrumentationModule = modulesByName.get(moduleClassName);
30+
private static final ClassLoaderValue<Map<String, InstrumentationModuleClassLoader>>
31+
instrumentationClassLoaders = new ClassLoaderValue<>();
32+
33+
public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
34+
String moduleClassName, ClassLoader instrumentedClassLoader) {
35+
InstrumentationModule instrumentationModule = modulesByClassName.get(moduleClassName);
4636
if (instrumentationModule == null) {
4737
throw new IllegalArgumentException(
48-
"No module with the class name " + modulesByName + " has been registered!");
38+
"No module with the class name " + modulesByClassName + " has been registered!");
4939
}
50-
return getInstrumentationClassloader(instrumentationModule, instrumentedClassloader);
40+
return getInstrumentationClassLoader(instrumentationModule, instrumentedClassLoader);
5141
}
5242

53-
private static synchronized InstrumentationModuleClassLoader getInstrumentationClassloader(
54-
InstrumentationModule module, ClassLoader instrumentedClassloader) {
43+
public static InstrumentationModuleClassLoader getInstrumentationClassLoader(
44+
InstrumentationModule module, ClassLoader instrumentedClassLoader) {
45+
46+
String groupName = getModuleGroup(module);
5547

56-
Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>> cacheForModule =
57-
instrumentationClassloaders.computeIfAbsent(module, (k) -> Cache.weak());
48+
Map<String, InstrumentationModuleClassLoader> loadersByGroupName =
49+
instrumentationClassLoaders.get(instrumentedClassLoader);
5850

59-
instrumentedClassloader = maskNullClassLoader(instrumentedClassloader);
60-
WeakReference<InstrumentationModuleClassLoader> cached =
61-
cacheForModule.get(instrumentedClassloader);
62-
if (cached != null) {
63-
InstrumentationModuleClassLoader cachedCl = cached.get();
64-
if (cachedCl != null) {
65-
return cachedCl;
66-
}
51+
if (loadersByGroupName == null) {
52+
throw new IllegalArgumentException(
53+
module
54+
+ " has not been initialized for class loader "
55+
+ instrumentedClassLoader
56+
+ " yet");
6757
}
68-
// We can't directly use "compute-if-absent" here because then for a short time only the
69-
// WeakReference will point to the InstrumentationModuleCL
70-
InstrumentationModuleClassLoader created =
71-
createInstrumentationModuleClassloader(module, instrumentedClassloader);
72-
cacheForModule.put(instrumentedClassloader, new WeakReference<>(created));
73-
return created;
74-
}
7558

76-
private static final ClassLoader BOOT_LOADER = new ClassLoader() {};
59+
InstrumentationModuleClassLoader loader = loadersByGroupName.get(groupName);
60+
if (loader == null || !loader.hasModuleInstalled(module)) {
61+
throw new IllegalArgumentException(
62+
module
63+
+ " has not been initialized for class loader "
64+
+ instrumentedClassLoader
65+
+ " yet");
66+
}
7767

78-
private static ClassLoader maskNullClassLoader(ClassLoader classLoader) {
79-
return classLoader == null ? BOOT_LOADER : classLoader;
68+
return loader;
8069
}
8170

82-
static InstrumentationModuleClassLoader createInstrumentationModuleClassloader(
83-
InstrumentationModule module, ClassLoader instrumentedClassloader) {
84-
85-
Set<String> toInject = new HashSet<>(InstrumentationModuleMuzzle.getHelperClassNames(module));
86-
// TODO (Jonas): Make muzzle include advice classes as helper classes
87-
// so that we don't have to include them here
88-
toInject.addAll(getModuleAdviceNames(module));
89-
if (module instanceof ExperimentalInstrumentationModule) {
90-
toInject.removeAll(((ExperimentalInstrumentationModule) module).injectedClassNames());
91-
}
92-
71+
/**
72+
* Returns a newly created class loader containing only the provided module. Note that other
73+
* modules from the same module group (see {@link #getModuleGroup(InstrumentationModule)}) will
74+
* not be installed in this class loader.
75+
*/
76+
public static InstrumentationModuleClassLoader
77+
createInstrumentationClassLoaderWithoutRegistration(
78+
InstrumentationModule module, ClassLoader instrumentedClassLoader) {
79+
// TODO: remove this method and replace usages with a custom TypePool implementation instead
9380
ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
94-
Map<String, BytecodeWithUrl> injectedClasses =
95-
toInject.stream()
96-
.collect(
97-
Collectors.toMap(
98-
name -> name, name -> BytecodeWithUrl.create(name, agentOrExtensionCl)));
99-
100-
return new InstrumentationModuleClassLoader(
101-
instrumentedClassloader, agentOrExtensionCl, injectedClasses);
81+
InstrumentationModuleClassLoader cl =
82+
new InstrumentationModuleClassLoader(instrumentedClassLoader, agentOrExtensionCl);
83+
cl.installModule(module);
84+
return cl;
10285
}
10386

104-
public static void registerIndyModule(InstrumentationModule module) {
87+
public static AgentBuilder.Identified.Extendable initializeModuleLoaderOnMatch(
88+
InstrumentationModule module, AgentBuilder.Identified.Extendable agentBuilder) {
10589
if (!module.isIndyModule()) {
10690
throw new IllegalArgumentException("Provided module is not an indy module!");
10791
}
10892
String moduleName = module.getClass().getName();
109-
if (modulesByName.putIfAbsent(moduleName, module) != null) {
93+
InstrumentationModule existingRegistration = modulesByClassName.putIfAbsent(moduleName, module);
94+
if (existingRegistration != null && existingRegistration != module) {
11095
throw new IllegalArgumentException(
111-
"A module with the class name " + moduleName + " has already been registered!");
96+
"A different module with the class name " + moduleName + " has already been registered!");
11297
}
98+
return agentBuilder.transform(
99+
(builder, typeDescription, classLoader, javaModule, protectionDomain) -> {
100+
initializeModuleLoaderForClassLoader(module, classLoader);
101+
return builder;
102+
});
103+
}
104+
105+
private static void initializeModuleLoaderForClassLoader(
106+
InstrumentationModule module, ClassLoader classLoader) {
107+
108+
ClassLoader agentOrExtensionCl = module.getClass().getClassLoader();
109+
110+
String groupName = getModuleGroup(module);
111+
112+
InstrumentationModuleClassLoader moduleCl =
113+
instrumentationClassLoaders
114+
.computeIfAbsent(classLoader, ConcurrentHashMap::new)
115+
.computeIfAbsent(
116+
groupName,
117+
unused -> new InstrumentationModuleClassLoader(classLoader, agentOrExtensionCl));
118+
119+
moduleCl.installModule(module);
113120
}
114121

115-
private static Set<String> getModuleAdviceNames(InstrumentationModule module) {
116-
Set<String> adviceNames = new HashSet<>();
117-
TypeTransformer nameCollector =
118-
new TypeTransformer() {
119-
@Override
120-
public void applyAdviceToMethod(
121-
ElementMatcher<? super MethodDescription> methodMatcher, String adviceClassName) {
122-
adviceNames.add(adviceClassName);
123-
}
124-
125-
@Override
126-
public void applyTransformer(AgentBuilder.Transformer transformer) {}
127-
};
128-
for (TypeInstrumentation instr : module.typeInstrumentations()) {
129-
instr.transform(nameCollector);
122+
private static String getModuleGroup(InstrumentationModule module) {
123+
if (module instanceof ExperimentalInstrumentationModule) {
124+
return ((ExperimentalInstrumentationModule) module).getModuleGroup();
130125
}
131-
return adviceNames;
126+
return module.getClass().getName();
132127
}
133128
}

0 commit comments

Comments
 (0)