Undo lazy loading

Author: sarahchen6

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java

@@ -26,7 +26,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Supplier;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.AsmVisitorWrapper;
@@ -131,29 +130,24 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    // Resolve helper names lazily, and capture only GC-neutral values (not the module instance)
-    // so that the instrumenter can still be unloaded after setup, and load $Muzzle via the
-    // extended class loader like the muzzle check does.
-    final String instrumentationClass = module.getClass().getName();
-    final String[] declaredHelperClassNames = module.helperClassNames();
-    final boolean injectHelperDependencies = module.injectHelperDependencies();
+    // Resolve helper names eagerly at install. Only the resolved String[] is captured.
+    String[] helperClassNames =
+        InstrumenterModule.loadStaticMuzzleHelperClassNames(
+            Utils.getExtendedClassLoader(), module.getClass().getName());
+    if (null == helperClassNames) {
+      helperClassNames = module.helperClassNames();
+    }
+    if (module.injectHelperDependencies()) {
+      helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
+    }
     helperTransformer =
-        new HelperTransformer(
-            module.useAgentCodeSource(),
-            adviceShader,
-            module.getClass().getSimpleName(),
-            () -> {
-              String[] helperClassNames =
-                  InstrumenterModule.loadStaticMuzzleHelperClassNames(
-                      Utils.getExtendedClassLoader(), instrumentationClass);
-              if (null == helperClassNames) {
-                helperClassNames = declaredHelperClassNames;
-              }
-              if (injectHelperDependencies) {
-                helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
-              }
-              return helperClassNames;
-            });
+        helperClassNames.length > 0
+            ? new HelperTransformer(
+                module.useAgentCodeSource(),
+                adviceShader,
+                module.getClass().getSimpleName(),
+                helperClassNames)
+            : null;
 
     postProcessor = module.postProcessor();
 
@@ -397,7 +391,7 @@ static final class HelperTransformer extends HelperInjector implements AgentBuil
         boolean useAgentCodeSource,
         AdviceShader adviceShader,
         String requestingName,
-        Supplier<String[]> helperClassNames) {
+        String... helperClassNames) {
       super(useAgentCodeSource, adviceShader, requestingName, helperClassNames);
     }
   }

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

@@ -21,7 +21,6 @@
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
-import java.util.function.Supplier;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.dynamic.DynamicType;
@@ -43,10 +42,7 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
   private final AdviceShader adviceShader;
   private final String requestingName;
 
-  // Helper names are resolved lazily: when a supplier is provided, resolution is deferred to the
-  // first transform() so that we don't load the generated $Muzzle class at agent install.
-  private final Supplier<String[]> helperClassNamesSupplier;
-  private volatile Set<String> helperClassNames;
+  private final Set<String> helperClassNames;
   private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
 
   private final Map<ClassLoader, Boolean> injectedClassLoaders =
@@ -81,27 +77,9 @@ public HelperInjector(
     this.requestingName = requestingName;
     this.adviceShader = adviceShader;
 
-    this.helperClassNamesSupplier = null;
     this.helperClassNames = new LinkedHashSet<>(asList(helperClassNames));
   }
 
-  /**
-   * Construct HelperInjector whose helper names are resolved lazily on first {@link #transform}, to
-   * avoid resolving them (which may load the generated {@code $Muzzle} class) at agent install.
-   */
-  public HelperInjector(
-      final boolean useAgentCodeSource,
-      final AdviceShader adviceShader,
-      final String requestingName,
-      final Supplier<String[]> helperClassNamesSupplier) {
-    this.useAgentCodeSource = useAgentCodeSource;
-    this.requestingName = requestingName;
-    this.adviceShader = adviceShader;
-
-    this.helperClassNamesSupplier = helperClassNamesSupplier;
-    this.helperClassNames = null;
-  }
-
   public HelperInjector(
       final boolean useAgentCodeSource,
       final String requestingName,
@@ -110,27 +88,11 @@ public HelperInjector(
     this.requestingName = requestingName;
     this.adviceShader = null;
 
-    this.helperClassNamesSupplier = null;
     helperClassNames = helperMap.keySet();
     dynamicTypeMap.putAll(helperMap);
   }
 
-  /** Resolves helper class names, deferring to the supplier when one's provided. */
-  private Set<String> resolveHelperClassNames() {
-    Set<String> names = helperClassNames;
-    if (names == null) {
-      synchronized (this) {
-        names = helperClassNames;
-        if (names == null) {
-          names = new LinkedHashSet<>(asList(helperClassNamesSupplier.get()));
-          helperClassNames = names;
-        }
-      }
-    }
-    return names;
-  }
-
-  private Map<String, byte[]> getHelperMap(final Set<String> helperClassNames) throws IOException {
+  private Map<String, byte[]> getHelperMap() throws IOException {
     if (dynamicTypeMap.isEmpty()) {
       final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
       for (String helperName : helperClassNames) {
@@ -155,7 +117,6 @@ public DynamicType.Builder<?> transform(
       ClassLoader classLoader,
       final JavaModule module,
       final ProtectionDomain pd) {
-    final Set<String> helperClassNames = resolveHelperClassNames();
     if (!helperClassNames.isEmpty()) {
       if (classLoader == null) {
         throw new UnsupportedOperationException(
@@ -174,7 +135,7 @@ public DynamicType.Builder<?> transform(
                 String.join(",", helperClassNames));
           }
 
-          final Map<String, byte[]> classnameToBytes = getHelperMap(helperClassNames);
+          final Map<String, byte[]> classnameToBytes = getHelperMap();
           final Collection<Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
 
           // all datadog helper classes are in the unnamed module