Use source file to determine helpers

Author: sarahchen6

buildSrc/src/main/kotlin/datadog/gradle/plugin/instrument/BuildTimeInstrumentationPlugin.kt

@@ -41,7 +41,8 @@ import org.gradle.kotlin.dsl.withType
  *
  * Requirements for ByteBuddy plugins:
  * 1. Must implement [net.bytebuddy.build.Plugin]
- * 2. Must have a constructor accepting a [java.io.File] parameter (target directory)
+ * 2. Must have a constructor accepting a [java.io.File] (target directory), or a two-`File`
+ *    `(sourceDirectory, targetDirectory)` constructor when the plugin also needs the source folder
  * 3. Plugin classes must be available on `buildTimeInstrumentationPlugin` configuration
  *
  * @see BuildTimeInstrumentationExtension

buildSrc/src/main/kotlin/datadog/gradle/plugin/instrument/ByteBuddyInstrumenter.kt

@@ -33,7 +33,13 @@ object ByteBuddyInstrumenter {
       val factories = plugins.map {
         try {
           val pluginClass = instrumentingLoader.loadClass(it).asSubclass(Plugin::class.java)
-          val loadedPlugin = pluginClass.getConstructor(File::class.java).newInstance(targetDirectory)
+          // Prefer a (source, target) constructor; fall back to the single target-directory one.
+          val loadedPlugin = try {
+            pluginClass.getConstructor(File::class.java, File::class.java)
+              .newInstance(sourceDirectory, targetDirectory)
+          } catch (_: NoSuchMethodException) {
+            pluginClass.getConstructor(File::class.java).newInstance(targetDirectory)
+          }
           Plugin.Factory.Simple(loadedPlugin)
         } catch (throwable: Throwable) {
           throw IllegalStateException("Cannot resolve plugin: $it", throwable)

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

@@ -5,15 +5,17 @@
 
 /**
  * Classifies a referenced class as an injectable tracer helper, a bootstrap class, or a library
- * class — similar to OpenTelemetry's {@code HelperClassPredicate#isHelperClass}; however, the main
- * signal here is whether the class was compiled from the instrumentation subproject's own output
- * ({@code ownOutput}). {@link #HELPER_PREFIXES} additionally covers helpers that live in other
- * tracer subprojects.
+ * class — similar to OpenTelemetry's {@code HelperClassPredicate#isHelperClass}. The primary signal
+ * is {@code ownOutput}: a class the instrumentation subproject compiled itself.
+ *
+ * <p>A subproject only injects helpers it owns; a helper owned by another subproject must be
+ * declared explicitly via {@code helperClassNames()}. {@link #HELPER_PREFIXES} lists the shared
+ * infrastructure subprojects that are not owned by a specific subproject and so are always treated
+ * as helpers.
  */
 public final class HelperClassPredicate {
 
   static final String[] HELPER_PREFIXES = {
-    "datadog.trace.instrumentation.",
     "datadog.opentelemetry.shim.",
     "datadog.trace.agent.tooling.iast.",
     "datadog.trace.agent.tooling.nativeimage.",

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

@@ -34,9 +34,16 @@
 
 /** Generates a 'Muzzle' side-class for each {@link InstrumenterModule}. */
 public class MuzzleGenerator implements AsmVisitorWrapper {
+  /**
+   * The engine's source folder, fully populated before processing. OwnOutput is decided against it
+   * rather than the target folder, which is written incrementally as classes are processed.
+   */
+  private final File sourceDir;
+
   private final File targetDir;
 
-  public MuzzleGenerator(File targetDir) {
+  public MuzzleGenerator(File sourceDir, File targetDir) {
+    this.sourceDir = sourceDir;
     this.targetDir = targetDir;
   }
 
@@ -246,12 +253,12 @@ private byte[] generateMuzzleClass(InstrumenterModule module) {
 
   /** {@code true} if the class was compiled from this instrumentation subproject's own output. */
   private boolean isOwnOutput(String className) {
-    return new File(targetDir, className.replace('.', '/') + ".class").isFile();
+    return new File(sourceDir, className.replace('.', '/') + ".class").isFile();
   }
 
   /** Adds the nested classes ({@code Foo$Bar}, {@code Foo$1}, ...) of an own-output helper. */
   private void addNestedClasses(String className, Set<String> helperClasses) {
-    File classFile = new File(targetDir, className.replace('.', '/') + ".class");
+    File classFile = new File(sourceDir, className.replace('.', '/') + ".class");
     File dir = classFile.getParentFile();
     if (dir == null || !dir.isDirectory()) {
       return;

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

@@ -25,10 +25,16 @@ public class MuzzleGradlePlugin extends Plugin.ForElementMatcher {
     HierarchyMatchers.registerIfAbsent(HierarchyMatchers.simpleChecks());
   }
 
+  private final File sourceDir;
   private final File targetDir;
 
-  public MuzzleGradlePlugin(File targetDir) {
+  /**
+   * @param sourceDir the engine's source folder, used to decide which helpers this subproject owns
+   * @param targetDir the engine's target folder, where {@code $Muzzle} classes are written
+   */
+  public MuzzleGradlePlugin(File sourceDir, File targetDir) {
     super(concreteClass().and(extendsClass(named(InstrumenterModule.class.getName()))));
+    this.sourceDir = sourceDir;
     this.targetDir = targetDir;
   }
 
@@ -37,7 +43,7 @@ public DynamicType.Builder<?> apply(
       final DynamicType.Builder<?> builder,
       final TypeDescription typeDescription,
       final ClassFileLocator classFileLocator) {
-    return builder.visit(new MuzzleGenerator(targetDir));
+    return builder.visit(new MuzzleGenerator(sourceDir, targetDir));
   }
 
   @Override

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/muzzle/HelperClassPredicateTest.java

@@ -18,10 +18,14 @@ private static HelperClassPredicate predicateWithOwnOutput(String... ownOutput)
   }
 
   @Test
-  void instrumentationPackageIsHelperEvenWithoutOwnOutput() {
-    // classes under the instrumentation package are ours regardless of where they were compiled
-    HelperClassPredicate predicate = predicateWithOwnOutput();
-    assertTrue(predicate.isHelperClass("datadog.trace.instrumentation.foo.FooDecorator"));
+  void instrumentationPackageIsHelperOnlyWhenOwnOutput() {
+    // a module's own instrumentation-package class is a helper via ownOutput
+    HelperClassPredicate own =
+        predicateWithOwnOutput("datadog.trace.instrumentation.foo.FooDecorator");
+    assertTrue(own.isHelperClass("datadog.trace.instrumentation.foo.FooDecorator"));
+    // a helper compiled by a different instrumentation subproject is not adopted just by reference
+    HelperClassPredicate other = predicateWithOwnOutput();
+    assertFalse(other.isHelperClass("datadog.trace.instrumentation.servlet.ServletBlockingHelper"));
   }
 
   @Test