Merge pull request DataDog#1126 from DataDog/landerson/helper-injector-optimization

randomanderson · web-flow · commit f77cce0b79ab · 2019-12-05T10:35:43.000-05:00
HelperInjector optimization
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -8,7 +8,6 @@
 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.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 HelperInjector(final String... helperClassNames) {
 
   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 static HelperInjector forDynamicTypes(final Collection<DynamicType.Unload
     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) {
-        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);
+  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());
+
+      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,55 +90,53 @@ public DynamicType.Builder<?> transform(
       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);
-            if (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER) {
-              final Map<TypeDescription, Class<?>> injected =
-                  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()));
-                }
-              }
-            } 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()));
-                }
-              }
-            }
-          } 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
-                        ? "<bootstrap>"
-                        : classLoader.getClass().getName()),
-                e);
-            throw new RuntimeException(e);
+      if (classLoader == BOOTSTRAP_CLASSLOADER) {
+        classLoader = BOOTSTRAP_CLASSLOADER_PLACEHOLDER;
+      }
+
+      if (!injectedClassLoaders.containsKey(classLoader)) {
+        try {
+          log.debug("Injecting classes onto classloader {} -> {}", classLoader, helperClassNames);
+
+          final Map<String, byte[]> classnameToBytes = getHelperMap();
+          final Map<String, Class<?>> classes;
+          if (classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER) {
+            classes =
+                ClassInjector.UsingInstrumentation.of(
+                        new File(System.getProperty("java.io.tmpdir")),
+                        ClassInjector.UsingInstrumentation.Target.BOOTSTRAP,
+                        AgentInstaller.getInstrumentation())
+                    .injectRaw(classnameToBytes);
+          } else {
+            classes = new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
           }
-          injectedClassLoaders.put(classLoader, true);
+
+          // 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) {
+          final String classLoaderType =
+              classLoader == BOOTSTRAP_CLASSLOADER_PLACEHOLDER
+                  ? "<bootstrap>"
+                  : 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);
         }
 
-        ensureModuleCanReadHelperModules(module);
+        injectedClassLoaders.put(classLoader, true);
       }
+
+      ensureModuleCanReadHelperModules(module);
     }
     return builder;
   }
