Add index for instrumentation helper classes

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

@@ -61,6 +61,7 @@ public final class CombiningTransformerBuilder
 
   private final AgentBuilder agentBuilder;
   private final InstrumenterIndex instrumenterIndex;
+  private final HelpersIndex helpersIndex = HelpersIndex.readIndex();
   private final int knownTransformationCount;
   private final Set<InstrumenterModule.TargetSystem> enabledSystems;
 
@@ -130,9 +131,9 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
 
     adviceShader = AdviceShader.with(module);
 
-    String[] helperClassNames =
-        InstrumenterModule.loadStaticMuzzleHelperClassNames(
-            Utils.getExtendedClassLoader(), module.getClass().getName());
+    // Helper names are resolved at build time and read from the index; fall back to the module for
+    // runtime-configured ones.
+    String[] helperClassNames = helpersIndex.helperClassNames(instrumentationId);
     if (null == helperClassNames) {
       helperClassNames = module.helperClassNames();
     }

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

@@ -0,0 +1,133 @@
+package datadog.trace.agent.tooling;
+
+import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
+import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.net.URL;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Maps each {@link InstrumenterModule} to the helper classes it injects, resolved at build time so
+ * the agent does not have to load every module's {@code $Muzzle} class at install to read them.
+ *
+ * <p>Based on {@link InstrumenterIndex} and its module ordering (both iterate {@link
+ * InstrumenterIndex#loadModules}), so a module's {@code instrumentationId} indexes into this table
+ * directly.
+ */
+final class HelpersIndex {
+  private static final Logger log = LoggerFactory.getLogger(HelpersIndex.class);
+
+  private static final String HELPERS_INDEX_NAME = "helpers.index";
+
+  static final ClassLoader instrumenterClassLoader = Instrumenter.class.getClassLoader();
+
+  private final String[][] helpersByInstrumentationId;
+
+  private HelpersIndex(String[][] helpersByInstrumentationId) {
+    this.helpersByInstrumentationId = helpersByInstrumentationId;
+  }
+
+  /** Build-time-resolved helper class names for the module, or {@code null} if not indexed. */
+  public String[] helperClassNames(int instrumentationId) {
+    if (instrumentationId < 0 || instrumentationId >= helpersByInstrumentationId.length) {
+      return null;
+    }
+    return helpersByInstrumentationId[instrumentationId];
+  }
+
+  public static HelpersIndex readIndex() {
+    URL indexResource = instrumenterClassLoader.getResource(HELPERS_INDEX_NAME);
+    if (null != indexResource) {
+      try (DataInputStream in =
+          new DataInputStream(new BufferedInputStream(indexResource.openStream()))) {
+        int moduleCount = in.readInt();
+        String[][] helpers = new String[moduleCount][];
+        for (int i = 0; i < moduleCount; i++) {
+          String[] names = new String[in.readInt()];
+          for (int j = 0; j < names.length; j++) {
+            names[j] = in.readUTF();
+          }
+          helpers[i] = names;
+        }
+        return new HelpersIndex(helpers);
+      } catch (Throwable e) {
+        log.error("Problem reading {}", HELPERS_INDEX_NAME, e);
+      }
+    }
+    return buildIndex(); // fallback to runtime generation when testing
+  }
+
+  public static HelpersIndex buildIndex() {
+    IndexGenerator indexGenerator = new IndexGenerator();
+    indexGenerator.buildIndex();
+    return new HelpersIndex(indexGenerator.helpers.toArray(new String[0][]));
+  }
+
+  /**
+   * Resolves each module's helpers from its build-time {@code $Muzzle}, falling back to the API.
+   */
+  static String[] resolveHelperClassNames(InstrumenterModule module) {
+    String[] helperClassNames =
+        InstrumenterModule.loadStaticMuzzleHelperClassNames(
+            instrumenterClassLoader, module.getClass().getName());
+    return null != helperClassNames ? helperClassNames : module.helperClassNames();
+  }
+
+  /** Generates the helpers index from known {@link InstrumenterModule}s on the build class-path. */
+  static final class IndexGenerator {
+    final List<String[]> helpers = new ArrayList<>();
+
+    void buildIndex() {
+      log.debug("Generating HelpersIndex");
+      try {
+        for (InstrumenterModule module : InstrumenterIndex.loadModules(instrumenterClassLoader)) {
+          helpers.add(resolveHelperClassNames(module));
+        }
+      } catch (IOException e) {
+        throw new UncheckedIOException("Problem generating HelpersIndex", e);
+      }
+    }
+
+    void writeIndex(Path indexFile) throws IOException {
+      try (DataOutputStream out =
+          new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(indexFile)))) {
+        out.writeInt(helpers.size());
+        for (String[] names : helpers) {
+          out.writeInt(names.length);
+          for (String name : names) {
+            out.writeUTF(name);
+          }
+        }
+      }
+    }
+
+    /**
+     * Called from the 'generateHelpersIndex' task in 'dd-java-agent/instrumentation/build.gradle'.
+     */
+    public static void main(String[] args) throws IOException {
+      if (args.length < 1) {
+        throw new IllegalArgumentException("Expected: index-dir");
+      }
+      Path indexDir = Paths.get(args[0]).toAbsolutePath();
+
+      // satisfy some instrumenters that cache matchers in initializers
+      HierarchyMatchers.registerIfAbsent(HierarchyMatchers.simpleChecks());
+      SharedTypePools.registerIfAbsent(SharedTypePools.simpleCache());
+
+      IndexGenerator indexGenerator = new IndexGenerator();
+      indexGenerator.buildIndex();
+      indexGenerator.writeIndex(indexDir.resolve(HELPERS_INDEX_NAME));
+    }
+  }
+}

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/HelpersIndexTest.java

@@ -0,0 +1,39 @@
+package datadog.trace.agent.tooling;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
+import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
+import java.util.List;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class HelpersIndexTest {
+
+  @BeforeAll
+  static void setup() {
+    // some modules cache matchers in initializers while resolving helpers
+    HierarchyMatchers.registerIfAbsent(HierarchyMatchers.simpleChecks());
+    SharedTypePools.registerIfAbsent(SharedTypePools.simpleCache());
+  }
+
+  @Test
+  void entriesMatchModuleOrderAndResolvedHelpers() throws Exception {
+    List<InstrumenterModule> modules =
+        InstrumenterIndex.loadModules(HelpersIndex.instrumenterClassLoader);
+    HelpersIndex index = HelpersIndex.buildIndex();
+
+    assertTrue(modules.size() > 0, "expected instrumentation modules on the class-path");
+    for (int i = 0; i < modules.size(); i++) {
+      assertArrayEquals(
+          HelpersIndex.resolveHelperClassNames(modules.get(i)),
+          index.helperClassNames(i),
+          "helpers mismatch for module " + modules.get(i).getClass().getName());
+    }
+    // ids outside the module range are not indexed.
+    assertNull(index.helperClassNames(-1));
+    assertNull(index.helperClassNames(modules.size()));
+  }
+}

dd-java-agent/instrumentation/build.gradle

@@ -185,3 +185,9 @@ registerIndexTask(
   'datadog.trace.agent.tooling.KnownTypesIndex$IndexGenerator',
   'Generate known-types.index'
   )
+
+registerIndexTask(
+  'generateHelpersIndex',
+  'datadog.trace.agent.tooling.HelpersIndex$IndexGenerator',
+  'Generate helpers.index'
+  )