Make helper-name resolution lazy

Author: sarahchen6

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

@@ -26,6 +26,7 @@
 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;
@@ -130,18 +131,20 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    String[] helperClassNames = module.getAllHelperClassNames();
-    if (module.injectHelperDependencies()) {
-      helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
-    }
+    // Resolve helper names lazily, so agent install doesn't load every module's generated $Muzzle
+    // class.
     helperTransformer =
-        helperClassNames.length > 0
-            ? new HelperTransformer(
-                module.useAgentCodeSource(),
-                adviceShader,
-                module.getClass().getSimpleName(),
-                helperClassNames)
-            : null;
+        new HelperTransformer(
+            module.useAgentCodeSource(),
+            adviceShader,
+            module.getClass().getSimpleName(),
+            () -> {
+              String[] helperClassNames = module.getAllHelperClassNames();
+              if (module.injectHelperDependencies()) {
+                helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
+              }
+              return helperClassNames;
+            });
 
     postProcessor = module.postProcessor();
 
@@ -385,7 +388,7 @@ static final class HelperTransformer extends HelperInjector implements AgentBuil
         boolean useAgentCodeSource,
         AdviceShader adviceShader,
         String requestingName,
-        String... helperClassNames) {
+        Supplier<String[]> helperClassNames) {
       super(useAgentCodeSource, adviceShader, requestingName, helperClassNames);
     }
   }

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

@@ -21,6 +21,7 @@
 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;
@@ -42,7 +43,10 @@ public class HelperInjector implements Instrumenter.TransformingAdvice {
   private final AdviceShader adviceShader;
   private final String requestingName;
 
-  private final Set<String> helperClassNames;
+  // 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 Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
 
   private final Map<ClassLoader, Boolean> injectedClassLoaders =
@@ -77,9 +81,27 @@ 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,
@@ -88,11 +110,27 @@ public HelperInjector(
     this.requestingName = requestingName;
     this.adviceShader = null;
 
+    this.helperClassNamesSupplier = null;
     helperClassNames = helperMap.keySet();
     dynamicTypeMap.putAll(helperMap);
   }
 
-  private Map<String, byte[]> getHelperMap() throws IOException {
+  /** 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 {
     if (dynamicTypeMap.isEmpty()) {
       final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
       for (String helperName : helperClassNames) {
@@ -117,6 +155,7 @@ 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(
@@ -135,7 +174,7 @@ public DynamicType.Builder<?> transform(
                 String.join(",", helperClassNames));
           }
 
-          final Map<String, byte[]> classnameToBytes = getHelperMap();
+          final Map<String, byte[]> classnameToBytes = getHelperMap(helperClassNames);
           final Collection<Class<?>> classes = injectClassLoader(classLoader, classnameToBytes);
 
           // all datadog helper classes are in the unnamed module