From 3dc6d899faab824fac48bfd21c7d8d59cd9cf54e Mon Sep 17 00:00:00 2001 From: Sarah Chen Date: Mon, 3 Aug 2026 15:09:24 -0400 Subject: [PATCH] Write helpers into module bytecode directly --- .../tooling/CombiningTransformerBuilder.java | 7 +- .../agent/tooling/InstrumenterModule.java | 19 +--- .../agent/tooling/muzzle/MuzzleGenerator.java | 96 ++++++++++++------- 3 files changed, 65 insertions(+), 57 deletions(-) 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..5a7e4b9c3df 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 @@ -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); } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java index 833d573cfde..86981af75a0 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java +++ b/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; } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java index 50cc27526de..22ddf4c61c2 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleGenerator.java +++ b/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 adviceClasses = new HashSet<>(); + List 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 adviceClasses = new HashSet<>(); - List 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 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. */