|
5 | 5 |
|
6 | 6 | package io.opentelemetry.javaagent.tooling.instrumentation.indy;
|
7 | 7 |
|
8 |
| -import io.opentelemetry.instrumentation.api.internal.cache.Cache; |
9 | 8 | import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
|
10 |
| -import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
11 |
| -import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
12 | 9 | 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; |
17 | 11 | import java.util.Map;
|
18 |
| -import java.util.Set; |
19 | 12 | import java.util.concurrent.ConcurrentHashMap;
|
20 |
| -import java.util.stream.Collectors; |
21 | 13 | import net.bytebuddy.agent.builder.AgentBuilder;
|
22 |
| -import net.bytebuddy.description.method.MethodDescription; |
23 |
| -import net.bytebuddy.matcher.ElementMatcher; |
24 | 14 |
|
25 | 15 | public class IndyModuleRegistry {
|
26 | 16 |
|
27 | 17 | private IndyModuleRegistry() {}
|
28 | 18 |
|
29 |
| - private static final ConcurrentHashMap<String, InstrumentationModule> modulesByName = |
| 19 | + private static final ConcurrentHashMap<String, InstrumentationModule> modulesByClassName = |
30 | 20 | new ConcurrentHashMap<>();
|
31 | 21 |
|
32 | 22 | /**
|
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()}; |
37 | 29 | */
|
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); |
46 | 36 | if (instrumentationModule == null) {
|
47 | 37 | 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!"); |
49 | 39 | }
|
50 |
| - return getInstrumentationClassloader(instrumentationModule, instrumentedClassloader); |
| 40 | + return getInstrumentationClassLoader(instrumentationModule, instrumentedClassLoader); |
51 | 41 | }
|
52 | 42 |
|
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); |
55 | 47 |
|
56 |
| - Cache<ClassLoader, WeakReference<InstrumentationModuleClassLoader>> cacheForModule = |
57 |
| - instrumentationClassloaders.computeIfAbsent(module, (k) -> Cache.weak()); |
| 48 | + Map<String, InstrumentationModuleClassLoader> loadersByGroupName = |
| 49 | + instrumentationClassLoaders.get(instrumentedClassLoader); |
58 | 50 |
|
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"); |
67 | 57 | }
|
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 |
| - } |
75 | 58 |
|
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 | + } |
77 | 67 |
|
78 |
| - private static ClassLoader maskNullClassLoader(ClassLoader classLoader) { |
79 |
| - return classLoader == null ? BOOT_LOADER : classLoader; |
| 68 | + return loader; |
80 | 69 | }
|
81 | 70 |
|
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 |
93 | 80 | 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; |
102 | 85 | }
|
103 | 86 |
|
104 |
| - public static void registerIndyModule(InstrumentationModule module) { |
| 87 | + public static AgentBuilder.Identified.Extendable initializeModuleLoaderOnMatch( |
| 88 | + InstrumentationModule module, AgentBuilder.Identified.Extendable agentBuilder) { |
105 | 89 | if (!module.isIndyModule()) {
|
106 | 90 | throw new IllegalArgumentException("Provided module is not an indy module!");
|
107 | 91 | }
|
108 | 92 | 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) { |
110 | 95 | 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!"); |
112 | 97 | }
|
| 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); |
113 | 120 | }
|
114 | 121 |
|
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(); |
130 | 125 | }
|
131 |
| - return adviceNames; |
| 126 | + return module.getClass().getName(); |
132 | 127 | }
|
133 | 128 | }
|
0 commit comments