Do not capture entire module instanc at prepareInstrumentation

Author: sarahchen6

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

@@ -131,16 +131,25 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    // Resolve helper names lazily, so agent install doesn't load every module's generated $Muzzle
-    // class.
+    // 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();
     helperTransformer =
         new HelperTransformer(
             module.useAgentCodeSource(),
             adviceShader,
             module.getClass().getSimpleName(),
             () -> {
-              String[] helperClassNames = module.getAllHelperClassNames();
-              if (module.injectHelperDependencies()) {
+              String[] helperClassNames =
+                  InstrumenterModule.loadStaticMuzzleHelperClassNames(
+                      Utils.getExtendedClassLoader(), instrumentationClass);
+              if (null == helperClassNames) {
+                helperClassNames = declaredHelperClassNames;
+              }
+              if (injectHelperDependencies) {
                 helperClassNames = HelperScanner.withClassDependencies(helperClassNames);
               }
               return helperClassNames;

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

@@ -113,26 +113,18 @@ public static ReferenceMatcher loadStaticMuzzleReferences(
   }
 
   /**
-   * @return the full, load-ordered set of helper classes to inject into the user's classloader: the
-   *     helpers inferred at build time from the advice/helper byte-code, unioned with any {@link
-   *     #helperClassNames() manually-declared} additions. Falls back to {@link #helperClassNames()}
-   *     when no generated {@code $Muzzle} is available.
+   * @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()}.
    */
-  public final String[] getAllHelperClassNames() {
-    String[] generated =
-        loadStaticMuzzleHelperClassNames(getClass().getClassLoader(), getClass().getName());
-    return null != generated ? generated : helperClassNames();
-  }
-
-  static String[] loadStaticMuzzleHelperClassNames(
+  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) {
-      // no generated helper list: caller falls back to helperClassNames()
       return null;
     }
   }