Write helpers into module bytecode directly

Author: sarahchen6

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

@@ -130,12 +130,7 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    String[] helperClassNames =
-        InstrumenterModule.loadStaticMuzzleHelperClassNames(
-            Utils.getExtendedClassLoader(), module.getClass().getName());
-    if (null == helperClassNames) {
-      helperClassNames = module.helperClassNames();
-    }
+    String[] helperClassNames = module.helperClassNames();
     if (module.injectHelperDependencies()) {
       helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
     }

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

@@ -113,23 +113,10 @@ public static ReferenceMatcher loadStaticMuzzleReferences(
   }
 
   /**
-   * @return the build-time inferred and manually-declared helper class names captured by {@code
-   *     $Muzzle}, or {@code null} when none are available and fall back to {@link
-   *     #helperClassNames()}.
+   * Optional manual additions to the injected helper set. At build time {@code MuzzleGenerator}
+   * overwrites this with the fully resolved list (inferred + manual), so at runtime it returns
+   * every helper the module injects.
    */
-  public static String[] loadStaticMuzzleHelperClassNames(
-      ClassLoader classLoader, String instrumentationClass) {
-    String muzzleClass = instrumentationClass + "$Muzzle";
-    try {
-      // helper class names captured at build-time; see MuzzleGenerator
-      return (String[])
-          classLoader.loadClass(muzzleClass).getMethod("helperClassNames").invoke(null);
-    } catch (Throwable e) {
-      return null;
-    }
-  }
-
-  /** Optional manual additions to the injected helper set. */
   public String[] helperClassNames() {
     return NO_HELPERS;
   }

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

@@ -77,14 +77,33 @@ public ClassVisitor wrap(
       throw new RuntimeException(e);
     }
 
+    AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
+
+    // Collect the muzzle references from every advice the module defines.
+    Set<String> adviceClasses = new HashSet<>();
+    List<Reference> allReferences = new ArrayList<>();
+    for (Instrumenter instrumenter : module.typeInstrumentations()) {
+      if (instrumenter instanceof Instrumenter.HasMethodAdvice) {
+        Collections.addAll(
+            allReferences,
+            generateReferences(
+                (Instrumenter.HasMethodAdvice) instrumenter, adviceShader, adviceClasses));
+      }
+    }
+
+    String[] orderedHelpers = computeInjectedHelpers(module, allReferences, adviceClasses);
+
     File muzzleClass = new File(targetDir, moduleDefinition.getInternalName() + "$Muzzle.class");
     try {
       muzzleClass.getParentFile().mkdirs();
-      Files.write(muzzleClass.toPath(), generateMuzzleClass(module));
+      Files.write(muzzleClass.toPath(), generateMuzzleClass(module, allReferences, orderedHelpers));
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
-    return classVisitor;
+
+    // Set resolved helpers directly in the module's helperClassNames() so agent reads
+    // them directly without loading the $Muzzle class.
+    return new HelperClassNamesWriter(classVisitor, orderedHelpers);
   }
 
   private static Reference[] generateReferences(
@@ -123,23 +142,8 @@ private static Reference[] generateReferences(
   }
 
   /** This code is generated in a separate side-class. */
-  private byte[] generateMuzzleClass(InstrumenterModule module) {
-
-    AdviceShader adviceShader = AdviceShader.with(module.adviceShading());
-
-    // Collect the muzzle references from every advice the module defines.
-    Set<String> adviceClasses = new HashSet<>();
-    List<Reference> allReferences = new ArrayList<>();
-    for (Instrumenter instrumenter : module.typeInstrumentations()) {
-      if (instrumenter instanceof Instrumenter.HasMethodAdvice) {
-        Collections.addAll(
-            allReferences,
-            generateReferences(
-                (Instrumenter.HasMethodAdvice) instrumenter, adviceShader, adviceClasses));
-      }
-    }
-
-    String[] orderedHelpers = computeInjectedHelpers(module, allReferences, adviceClasses);
+  private byte[] generateMuzzleClass(
+      InstrumenterModule module, List<Reference> allReferences, String[] orderedHelpers) {
 
     // Injected helpers are our own classes, so they don't need to be asserted as library
     // references.
@@ -201,24 +205,46 @@ private byte[] generateMuzzleClass(InstrumenterModule module) {
     mv.visitMaxs(0, 0);
     mv.visitEnd();
 
-    // Generate helperClassNames() with resolved helpers for the agent to read at load time;
-    // skip the method entirely when the module injects nothing.
-    if (orderedHelpers.length > 0) {
-      MethodVisitor hv =
-          cw.visitMethod(
-              Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
-              "helperClassNames",
-              "()[Ljava/lang/String;",
-              null,
-              null);
-      hv.visitCode();
-      writeStrings(hv, orderedHelpers);
-      hv.visitInsn(Opcodes.ARETURN);
-      hv.visitMaxs(0, 0);
-      hv.visitEnd();
+    return cw.toByteArray();
+  }
+
+  /**
+   * Rewrite a module's {@code helperClassNames()} to return the build-time-resolved helper list.
+   */
+  private static final class HelperClassNamesWriter extends ClassVisitor {
+    private static final String HELPER_METHOD = "helperClassNames";
+    private static final String HELPER_DESCRIPTOR = "()[Ljava/lang/String;";
+
+    private final String[] helpers;
+
+    HelperClassNamesWriter(ClassVisitor classVisitor, String[] helpers) {
+      super(Opcodes.ASM7, classVisitor);
+      this.helpers = helpers;
     }
 
-    return cw.toByteArray();
+    @Override
+    public MethodVisitor visitMethod(
+        int access, String name, String descriptor, String signature, String[] exceptions) {
+      // Drop any existing helperClassNames() - resolved version will be re-added in visitEnd.
+      if (HELPER_METHOD.equals(name) && HELPER_DESCRIPTOR.equals(descriptor)) {
+        return null;
+      }
+      return super.visitMethod(access, name, descriptor, signature, exceptions);
+    }
+
+    @Override
+    public void visitEnd() {
+      if (helpers.length > 0) {
+        MethodVisitor mv =
+            super.visitMethod(Opcodes.ACC_PUBLIC, HELPER_METHOD, HELPER_DESCRIPTOR, null, null);
+        mv.visitCode();
+        writeStrings(mv, helpers);
+        mv.visitInsn(Opcodes.ARETURN);
+        mv.visitMaxs(0, 0);
+        mv.visitEnd();
+      }
+      super.visitEnd();
+    }
   }
 
   /** Resolves the ordered set of helper classes to inject for a module. */