|
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