Remove synchronization and static helper map
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -8,7 +8,6 @@ import java.io.File;
import java.io.IOException;
import java.lang.ref.WeakReference;
import java.security.SecureClassLoader;
-import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
@@ -18,6 +17,7 @@ import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.CopyOnWriteArrayList;
import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.builder.AgentBuilder.Transformer;
import net.bytebuddy.description.type.TypeDescription;
@@ -34,12 +34,11 @@ public class HelperInjector implements Transformer {
new SecureClassLoader(null) {};
private final Set<String> helperClassNames;
- private Map<TypeDescription, byte[]> helperMap = null;
- private final WeakMap<ClassLoader, Boolean> injectedClassLoaders = newWeakMap();
+ private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
- // Neither Module nor WeakReference implements equals or hashcode so using a list
- private final List<WeakReference<Object>> helperModules = new ArrayList<>();
+ private final WeakMap<ClassLoader, Boolean> injectedClassLoaders = newWeakMap();
+ private final List<WeakReference<Object>> helperModules = new CopyOnWriteArrayList<>();
/**
* Construct HelperInjector.
*
@@ -55,13 +54,7 @@ public class HelperInjector implements Transformer {
public HelperInjector(final Map<String, byte[]> helperMap) {
helperClassNames = helperMap.keySet();
- this.helperMap = new LinkedHashMap<>(helperClassNames.size());
- for (final String helperName : helperClassNames) {
- final TypeDescription typeDesc =
- new TypeDescription.Latent(
- helperName, 0, null, Collections.<TypeDescription.Generic>emptyList());
- this.helperMap.put(typeDesc, helperMap.get(helperName));
- }
+ dynamicTypeMap.putAll(helperMap);
}
public static HelperInjector forDynamicTypes(final Collection<DynamicType.Unloaded<?>> helpers) {
@@ -72,20 +65,22 @@ public class HelperInjector implements Transformer {
return new HelperInjector(bytes);
}
- private synchronized Map<TypeDescription, byte[]> getHelperMap() throws IOException {
- if (helperMap == null) {
- helperMap = new LinkedHashMap<>(helperClassNames.size());
- for (final String helperName : helperClassNames) {
+ private Map<String, byte[]> getHelperMap() throws IOException {
+ if (dynamicTypeMap.isEmpty()) {
+ final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
+
final ClassFileLocator locator =
ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader());
- final byte[] classBytes = locator.locate(helperName).resolve();
- final TypeDescription typeDesc =
- new TypeDescription.Latent(
- helperName, 0, null, Collections.<TypeDescription.Generic>emptyList());
- helperMap.put(typeDesc, classBytes);
+
+ for (final String helperClassName : helperClassNames) {
+ final byte[] classBytes = locator.locate(helperClassName).resolve();
+ classnameToBytes.put(helperClassName, classBytes);
}
+
+ return classnameToBytes;
+ } else {
+ return dynamicTypeMap;
}
- return helperMap;
}
@Override
@@ -95,56 +90,54 @@ public class HelperInjector implements Transformer {
ClassLoader classLoader,
final JavaModule module) {
if (!helperClassNames.isEmpty()) {
- synchronized (this) {
if (classLoader == BOOTSTRAP_CLASSLOADER) {
classLoader = BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
}
+
if (!injectedClassLoaders.containsKey(classLoader)) {
try {
- final Map<TypeDescription, byte[]> helperMap = getHelperMap();
log.debug("Injecting classes onto classloader {} -> {}", classLoader, helperClassNames);
+
+ final Map<String, byte[]> classnameToBytes = getHelperMap();
+ final Map<String, Class<?>> classes;
if (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER) {
- final Map<TypeDescription, Class<?>> injected =
+ classes =
ClassInjector.UsingInstrumentation.of(
new File(System.getProperty("java.io.tmpdir")),
ClassInjector.UsingInstrumentation.Target.BOOTSTRAP,
AgentInstaller.getInstrumentation())
- .inject(helperMap);
- for (final TypeDescription desc : injected.keySet()) {
- final Class<?> injectedClass =
- Class.forName(desc.getName(), false, Utils.getBootstrapProxy());
- if (JavaModule.isSupported()) {
- helperModules.add(new WeakReference<>(JavaModule.ofType(injectedClass).unwrap()));
- }
- }
+ .injectRaw(classnameToBytes);
} else {
- final Map<TypeDescription, Class<?>> classMap =
- new ClassInjector.UsingReflection(classLoader).inject(helperMap);
- if (JavaModule.isSupported()) {
- for (final Class<?> injectedClass : classMap.values()) {
- helperModules.add(new WeakReference<>(JavaModule.ofType(injectedClass).unwrap()));
- }
+ classes = new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
}
+
+ // All datadog helper classes are in the unnamed module
+ // And there's exactly one unnamed module per classloader
+ // Use the module of the first class for convenience
+ if (JavaModule.isSupported()) {
+ final JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
+ helperModules.add(new WeakReference<>(javaModule.unwrap()));
}
} catch (final Exception e) {
- log.error(
- "Error preparing helpers for "
- + typeDescription
- + ". Failed to inject helper classes into instance "
- + classLoader
- + " of type "
- + (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER
+ final String classLoaderType =
+ classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER
? "<bootstrap>"
- : classLoader.getClass().getName()),
+ : classLoader.getClass().getName();
+
+ log.error(
+ "Error preparing helpers for {}. Failed to inject helper classes into instance {} of type {}",
+ typeDescription,
+ classLoader,
+ classLoaderType,
e);
throw new RuntimeException(e);
}
+
injectedClassLoaders.put(classLoader, true);
}
ensureModuleCanReadHelperModules(module);
}
- }
return builder;
}

Remove synchronization and static helper map