From 5881daaac1783299cf1c3dd0104bfac8eb55e0f0 Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Thu, 30 Jul 2026 14:32:43 -0400 Subject: [PATCH 1/2] Add index for instrumentation helper classes --- .../tooling/CombiningTransformerBuilder.java | 7 +- .../trace/agent/tooling/HelpersIndex.java | 133 ++++++++++++++++++ .../trace/agent/tooling/HelpersIndexTest.java | 39 +++++ dd-java-agent/instrumentation/build.gradle | 6 + 4 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java create mode 100644 dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/HelpersIndexTest.java diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 415d4886503..7d53eaf1d73 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/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 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(); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java new file mode 100644 index 00000000000..78f3e00f4f1 --- /dev/null +++ b/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. + * + *

Entries are positional and share {@link InstrumenterIndex}'s 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 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)); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/HelpersIndexTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/HelpersIndexTest.java new file mode 100644 index 00000000000..95a88a4cc18 --- /dev/null +++ b/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 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())); + } +} diff --git a/dd-java-agent/instrumentation/build.gradle b/dd-java-agent/instrumentation/build.gradle index 4b2b6588492..f5525950e0a 100644 --- a/dd-java-agent/instrumentation/build.gradle +++ b/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' + ) From 1e713920adc0c84e8cac44d4124e075534a56de1 Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Fri, 31 Jul 2026 13:10:30 -0400 Subject: [PATCH 2/2] Reword comment --- .../main/java/datadog/trace/agent/tooling/HelpersIndex.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java index 78f3e00f4f1..f2009807db8 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelpersIndex.java @@ -21,9 +21,9 @@ * 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. * - *

Entries are positional and share {@link InstrumenterIndex}'s module ordering (both iterate - * {@link InstrumenterIndex#loadModules}), so a module's {@code instrumentationId} indexes into this - * table directly. + *

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);