Simplify HelperInjector to make it easier to switch out the underlying approach

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java

@@ -1,14 +1,17 @@
 package datadog.trace.agent.tooling;
 
 import static datadog.trace.bootstrap.AgentClassLoading.INJECTING_HELPERS;
+import static java.util.Arrays.asList;
 
 import datadog.trace.bootstrap.instrumentation.api.EagerHelper;
+import datadog.trace.util.JDK9ModuleAccess;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
+import java.lang.reflect.AnnotatedElement;
 import java.security.CodeSource;
 import java.security.ProtectionDomain;
-import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -20,7 +23,6 @@
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.dynamic.loading.ClassInjector;
-import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
 import net.bytebuddy.utility.JavaModule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -40,9 +42,9 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
   private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
 
   private final Map<ClassLoader, Boolean> injectedClassLoaders =
-      Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+      Collections.synchronizedMap(new WeakHashMap<>());
 
-  private final List<WeakReference<Object>> helperModules = new CopyOnWriteArrayList<>();
+  private final List<WeakReference<AnnotatedElement>> helperModules = new CopyOnWriteArrayList<>();
 
   /**
    * Construct HelperInjector.
@@ -71,7 +73,7 @@ public HelperInjector(
     this.requestingName = requestingName;
     this.adviceShader = adviceShader;
 
-    this.helperClassNames = new LinkedHashSet<>(Arrays.asList(helperClassNames));
+    this.helperClassNames = new LinkedHashSet<>(asList(helperClassNames));
   }
 
   public HelperInjector(
@@ -132,12 +134,10 @@ public DynamicType.Builder<?> transform(
           final Map<String, byte[]> classnameToBytes = getHelperMap();
           final Map<String, Class<?>> classes = injectClassLoader(classLoader, 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
+          // all datadog helper classes are in the unnamed module
+          // and there's exactly one unnamed module per classloader
           if (JavaModule.isSupported()) {
-            final JavaModule javaModule = JavaModule.ofType(classes.values().iterator().next());
-            helperModules.add(new WeakReference<>(javaModule.unwrap()));
+            helperModules.add(new WeakReference<>(JDK9ModuleAccess.getUnnamedModule(classLoader)));
           }
 
           // forcibly initialize any eager helpers
@@ -177,43 +177,42 @@ private Map<String, Class<?>> injectClassLoader(
       final ClassLoader classLoader, final Map<String, byte[]> classnameToBytes) {
     INJECTING_HELPERS.begin();
     try {
-      ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
-      return new ClassInjector.UsingReflection(classLoader, protectionDomain)
-          .injectRaw(classnameToBytes);
+      if (useAgentCodeSource) {
+        ProtectionDomain protectionDomain = createProtectionDomain(classLoader);
+        return new ClassInjector.UsingReflection(classLoader, protectionDomain)
+            .injectRaw(classnameToBytes);
+      } else {
+        return new ClassInjector.UsingReflection(classLoader).injectRaw(classnameToBytes);
+      }
     } finally {
       INJECTING_HELPERS.end();
     }
   }
 
   private ProtectionDomain createProtectionDomain(final ClassLoader classLoader) {
-    if (useAgentCodeSource) {
-      CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
-      return new ProtectionDomain(codeSource, null, classLoader, null);
-    } else {
-      return ClassLoadingStrategy.NO_PROTECTION_DOMAIN;
-    }
+    CodeSource codeSource = HelperInjector.class.getProtectionDomain().getCodeSource();
+    return new ProtectionDomain(codeSource, null, classLoader, null);
   }
 
   private void ensureModuleCanReadHelperModules(final JavaModule target) {
     if (JavaModule.isSupported() && target != JavaModule.UNSUPPORTED && target.isNamed()) {
-      for (final WeakReference<Object> helperModuleReference : helperModules) {
-        final Object realModule = helperModuleReference.get();
-        if (realModule != null) {
-          final JavaModule helperModule = JavaModule.of(realModule);
-
-          if (!target.canRead(helperModule)) {
-            log.debug("Adding module read from {} to {}", target, helperModule);
-            ClassInjector.UsingInstrumentation.redefineModule(
-                Utils.getInstrumentation(),
-                target,
-                Collections.singleton(helperModule),
-                Collections.<String, Set<JavaModule>>emptyMap(),
-                Collections.<String, Set<JavaModule>>emptyMap(),
-                Collections.<Class<?>>emptySet(),
-                Collections.<Class<?>, List<Class<?>>>emptyMap());
+      AnnotatedElement targetModule = (AnnotatedElement) target.unwrap();
+      Set<AnnotatedElement> extraReads = null;
+      for (final WeakReference<AnnotatedElement> helperModuleReference : helperModules) {
+        final AnnotatedElement helperModule = helperModuleReference.get();
+        if (helperModule != null) {
+          if (!JDK9ModuleAccess.canRead(targetModule, helperModule)) {
+            if (extraReads == null) {
+              extraReads = new HashSet<>();
+            }
+            extraReads.add(helperModule);
           }
         }
       }
+      if (extraReads != null) {
+        log.debug("Adding module reads from {} to {}", targetModule, extraReads);
+        JDK9ModuleAccess.addModuleReads(Utils.getInstrumentation(), targetModule, extraReads);
+      }
     }
   }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java

@@ -1,7 +1,6 @@
 package datadog.trace.agent.tooling.muzzle;
 
 import datadog.trace.agent.tooling.AdviceShader;
-import datadog.trace.agent.tooling.HelperInjector;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
 import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
@@ -72,16 +71,15 @@ public static void assertInstrumentationMuzzled(
     if (assertPass) {
       for (InstrumenterModule module : toBeTested) {
         try {
-          // verify helper injector works
+          // verify helper consistency
           final String[] helperClassNames = module.helperClassNames();
           if (helperClassNames.length > 0) {
-            new HelperInjector(
-                    module.useAgentCodeSource(),
-                    MuzzleVersionScanPlugin.class.getSimpleName(),
-                    createHelperMap(module))
-                .transform(null, null, testApplicationLoader, null, null);
+            HelperClassLoader helperClassLoader = new HelperClassLoader(testApplicationLoader);
+            for (Map.Entry<String, byte[]> helper : createHelperMap(module).entrySet()) {
+              helperClassLoader.injectClass(helper.getKey(), helper.getValue());
+            }
           }
-        } catch (final Exception e) {
+        } catch (final Throwable e) {
           System.err.println(
               "FAILED HELPER INJECTION. Are Helpers being injected in the correct order?");
           System.err.println(e.getMessage());
@@ -107,6 +105,18 @@ private static synchronized List<InstrumenterModule> toBeTested(
     return toBeTested;
   }
 
+  // Exposes ClassLoader.defineClass() to test helper consistency
+  // without requiring java.lang.instrument.Instrumentation agent
+  static final class HelperClassLoader extends ClassLoader {
+    HelperClassLoader(ClassLoader parent) {
+      super(parent);
+    }
+
+    public void injectClass(String name, byte[] bytecode) {
+      defineClass(name, bytecode, 0, bytecode.length);
+    }
+  }
+
   private static Map<String, byte[]> createHelperMap(final InstrumenterModule module)
       throws IOException {
     String[] helperClasses = module.helperClassNames();

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/HelperInjectionTest.groovy

@@ -3,9 +3,6 @@ package datadog.trace.agent.test
 import datadog.trace.agent.tooling.HelperInjector
 import datadog.trace.agent.tooling.Utils
 import datadog.trace.test.util.DDSpecification
-import net.bytebuddy.description.type.TypeDescription
-import net.bytebuddy.dynamic.ClassFileLocator
-import net.bytebuddy.dynamic.loading.ClassInjector
 
 import java.lang.ref.WeakReference
 import java.util.concurrent.atomic.AtomicReference
@@ -16,7 +13,6 @@ import static datadog.trace.test.util.GCUtils.awaitGC
 class HelperInjectionTest extends DDSpecification {
   static final String HELPER_CLASS_NAME = 'datadog.trace.agent.test.HelperClass'
 
-  //@Flaky("awaitGC usage is flaky")
   def "helpers injected to non-delegating classloader"() {
     setup:
     HelperInjector injector = new HelperInjector(false, "test", HELPER_CLASS_NAME)
@@ -45,39 +41,4 @@ class HelperInjectionTest extends DDSpecification {
     then: "HelperInjector doesn't prevent it from being collected"
     null == ref.get()
   }
-
-  //@Flaky("awaitGC usage is flaky")
-  def "check hard references on class injection"() {
-    setup:
-
-    // Copied from HelperInjector:
-    final ClassFileLocator locator =
-      ClassFileLocator.ForClassLoader.of(Utils.getAgentClassLoader())
-    final byte[] classBytes = locator.locate(HELPER_CLASS_NAME).resolve()
-    final TypeDescription typeDesc =
-      new TypeDescription.Latent(
-      HELPER_CLASS_NAME, 0, null, Collections.<TypeDescription.Generic> emptyList())
-
-    AtomicReference<URLClassLoader> emptyLoader = new AtomicReference<>(new URLClassLoader(new URL[0], (ClassLoader) null))
-    AtomicReference<ClassInjector> injector = new AtomicReference<>(new ClassInjector.UsingReflection(emptyLoader.get()))
-    injector.get().inject([(typeDesc): classBytes])
-
-    when:
-    def injectorRef = new WeakReference(injector.get())
-    injector.set(null)
-
-    awaitGC(injectorRef)
-
-    then:
-    null == injectorRef.get()
-
-    when:
-    def loaderRef = new WeakReference(emptyLoader.get())
-    emptyLoader.set(null)
-
-    awaitGC(loaderRef)
-
-    then:
-    null == loaderRef.get()
-  }
 }

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy

@@ -78,9 +78,6 @@ class MuzzleVersionScanPluginTest extends DDSpecification {
 
     expect:
     MuzzleVersionScanPlugin.assertInstrumentationMuzzled(instrumentationLoader, testApplicationLoader, true, null)
-    !helpers.findAll {
-      testApplicationLoader.loadClass(it.name) != null
-    }.isEmpty()
 
     where:
     // spotless:off

internal-api/internal-api-9/src/main/java/datadog/trace/util/JDK9ModuleAccess.java

@@ -0,0 +1,36 @@
+package datadog.trace.util;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.emptySet;
+
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.AnnotatedElement;
+import java.util.Set;
+
+/** Use standard API to work with JPMS modules on Java9+. */
+@SuppressWarnings("Since15")
+public final class JDK9ModuleAccess {
+
+  /** Retrieves a class-loader's unnamed module. */
+  public static AnnotatedElement getUnnamedModule(ClassLoader cl) {
+    return cl.getUnnamedModule();
+  }
+
+  /** Returns {@code true} if the first module can read the second module. */
+  public static boolean canRead(AnnotatedElement module, AnnotatedElement anotherModule) {
+    return ((java.lang.Module) module).canRead((java.lang.Module) anotherModule);
+  }
+
+  /** Adds extra module reads to the given module. */
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  public static void addModuleReads(
+      Instrumentation inst, AnnotatedElement module, Set<AnnotatedElement> extraReads) {
+    inst.redefineModule(
+        (java.lang.Module) module,
+        (Set) extraReads,
+        emptyMap(),
+        emptyMap(),
+        emptySet(),
+        emptyMap());
+  }
+}