Support generation of multiple types of glue per-module (#29)

Author: mcculls

build.gradle.kts

@@ -48,10 +48,7 @@ tasks.jar {
 tasks.javadoc {
   dependsOn(embed)
   setSource(allSources())
-  exclude(
-    "datadog/instrument/glue",
-    "datadog/instrument/utils/Glue.java",
-    "datadog/instrument/utils/JVM.java")
+  exclude("datadog/instrument/glue", "datadog/instrument/utils/JVM.java")
   var javadocOptions = (options as StandardJavadocDocletOptions)
   if (JavaVersion.current().isJava9Compatible) {
     javadocOptions.addBooleanOption("html5", true)

buildSrc/src/main/kotlin/instrument-glue.gradle.kts

@@ -25,18 +25,18 @@ dependencies {
 }
 
 tasks.register<JavaExec>("generateGlue") {
-  val glue: String by project.extra
+  val glue: List<String> by project.extra
 
   // generate glue files under a consistent packaging location
-  val resourcesDir = generatedGlueResources.dir("datadog/instrument/glue/")
-  val javaDir = generatedGlueJava.dir("datadog/instrument/glue/")
+  val resourcePath = generatedGlueResources.dir("datadog/instrument/glue/")
+  val javaPath = generatedGlueJava.dir("datadog/instrument/glue/")
 
   group = "Build"
-  description = "Generate ${glue}"
-  mainClass = "datadog.instrument.glue.${glue}Generator"
+  description = "Generate Instrumentation Glue"
+  mainClass = "datadog.instrument.glue.GlueGenerator"
   classpath = sourceSets["glue"].runtimeClasspath
-  args = listOf(resourcesDir.toString(), javaDir.toString())
-  outputs.dirs(resourcesDir, javaDir)
+  args = listOf(resourcePath.toString(), javaPath.toString()) + glue
+  outputs.dirs(resourcePath, javaPath)
 }
 
 tasks.processResources {

class-inject/build.gradle.kts

@@ -5,7 +5,7 @@ plugins {
 
 // class-inject generates glue bytecode at build-time to provide access to Unsafe.defineClass
 // we inline this glue as a string constant, then install it at runtime using instrumentation
-extra["glue"] = "DefineClassGlue"
+extra["glue"] = listOf("DefineClassGlue")
 
 tasks.jar {
   // DefineClassGlue only contains large string constants that get inlined into ClassInjector

class-inject/src/glue/java/datadog/instrument/glue/DefineClassGlueGenerator.java

@@ -6,16 +6,15 @@
 
 package datadog.instrument.glue;
 
-import static datadog.instrument.utils.Glue.classHeader;
-import static datadog.instrument.utils.Glue.packBytecode;
+import static datadog.instrument.glue.GlueGenerator.classHeader;
+import static datadog.instrument.glue.GlueGenerator.packBytecode;
 import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
 import static org.objectweb.asm.Opcodes.*;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiFunction;
@@ -351,11 +350,8 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     return cw.toByteArray();
   }
 
-  public static void main(String[] args) throws IOException {
-    if (args.length != 2) {
-      throw new IllegalArgumentException("Expected: resources-dir java-dir");
-    }
-    Path defineClassGlue = Paths.get(args[1], "DefineClassGlue.java");
+  public static void generateGlue(Path resourcePath, Path javaPath) throws IOException {
+    Path defineClassGlue = javaPath.resolve("DefineClassGlue.java");
     List<String> lines = new ArrayList<>();
     classHeader(lines, "DefineClassGlue");
     lines.add("  /** Glue Id */");

utils/src/main/java/datadog/instrument/glue/Glue.java

@@ -0,0 +1,65 @@
+/*
+ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+ * This product includes software developed at Datadog (https://www.datadoghq.com/).
+ * Copyright 2025-Present Datadog, Inc.
+ */
+
+package datadog.instrument.glue;
+
+import static java.nio.charset.StandardCharsets.UTF_16BE;
+
+import datadog.instrument.utils.JVM;
+import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
+import java.util.MissingResourceException;
+import java.util.Objects;
+
+/** Methods for loading instrumentation glue bytecode from string literals and resource files. */
+public final class Glue {
+
+  private static final String GLUE_RESOURCE_PREFIX = "/datadog/instrument/glue/";
+
+  private static final int BUFFER_SIZE = 8192;
+
+  private Glue() {}
+
+  /**
+   * Unpacks a string literal produced by {@link GlueGenerator#packBytecode} back into bytecode.
+   *
+   * @param bytecode the packed bytecode
+   * @return the unpacked bytecode
+   */
+  public static byte[] unpackBytecode(String bytecode) {
+    return bytecode.getBytes(UTF_16BE);
+  }
+
+  /**
+   * Loads glue bytecode with the given name from the given host.
+   *
+   * @param host the host class used to load the glue resource
+   * @param glueName the glue resource containing the bytecode
+   * @return the glue bytecode
+   * @throws MissingResourceException if the bytecode cannot be read
+   */
+  @SuppressWarnings({"Since15"})
+  public static byte[] loadBytecode(Class<?> host, String glueName) {
+    String glueResource = GLUE_RESOURCE_PREFIX + glueName;
+    try (InputStream is = Objects.requireNonNull(host.getResourceAsStream(glueResource))) {
+      if (JVM.atLeastJava(9)) {
+        return is.readAllBytes();
+      } else {
+        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
+          int bytesRead;
+          byte[] buf = new byte[BUFFER_SIZE];
+          while ((bytesRead = is.read(buf, 0, BUFFER_SIZE)) != -1) {
+            os.write(buf, 0, bytesRead);
+          }
+          return os.toByteArray();
+        }
+      }
+    } catch (Throwable e) {
+      String detail = "Cannot load " + glueResource + " from " + host + ": " + e;
+      throw new MissingResourceException(detail, host.getName(), glueResource);
+    }
+  }
+}

…/java/datadog/instrument/utils/Glue.java …tadog/instrument/glue/GlueGenerator.javautils/src/main/java/datadog/instrument/utils/Glue.java renamed to utils/src/main/java/datadog/instrument/glue/GlueGenerator.java utils/src/main/java/datadog/instrument/utils/Glue.java renamed to utils/src/main/java/datadog/instrument/glue/GlueGenerator.java

@@ -4,20 +4,31 @@
  * Copyright 2025-Present Datadog, Inc.
  */
 
-package datadog.instrument.utils;
+package datadog.instrument.glue;
 
-import static java.nio.charset.StandardCharsets.UTF_16BE;
-
-import java.io.ByteArrayOutputStream;
-import java.io.InputStream;
+import java.lang.reflect.Method;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.List;
-import java.util.MissingResourceException;
-import java.util.Objects;
 
-/** Methods for packing glue bytecode into strings that can be stored in the constant pool. */
-public final class Glue {
+/** Generates instrumentation glue for projects that apply the {@code instrument-glue} plugin. */
+public final class GlueGenerator {
+
+  private GlueGenerator() {}
 
-  private Glue() {}
+  /** Entry-point for generating instrumentation glue into resource files or Java source code. */
+  public static void main(String[] args) throws ReflectiveOperationException {
+    if (args.length < 3) {
+      throw new IllegalArgumentException("Expected: resource-path java-path glue-name...");
+    }
+    Path resourcePath = Paths.get(args[0]);
+    Path javaPath = Paths.get(args[1]);
+    for (int i = 2; i < args.length; i++) {
+      Class<?> generatorClass = Class.forName("datadog.instrument.glue." + args[i] + "Generator");
+      Method generateGlue = generatorClass.getMethod("generateGlue", Path.class, Path.class);
+      generateGlue.invoke(null, resourcePath, javaPath);
+    }
+  }
 
   /**
    * Writes the source header of a class to hold the string representation of glue bytecode.
@@ -63,48 +74,4 @@ public static void packBytecode(List<String> lines, byte[] bytecode) {
     }
     lines.add(buf + "\";");
   }
-
-  /**
-   * Unpacks a string literal produced by {@link #packBytecode} back into the original bytecode.
-   *
-   * @param bytecode the packed bytecode
-   * @return the unpacked bytecode
-   */
-  public static byte[] unpackBytecode(String bytecode) {
-    return bytecode.getBytes(UTF_16BE);
-  }
-
-  private static final String GLUE_RESOURCE_PREFIX = "/datadog/instrument/glue/";
-
-  private static final int BUFFER_SIZE = 8192;
-
-  /**
-   * Loads glue bytecode with the given name from the given host.
-   *
-   * @param host the host class used to load the glue resource
-   * @param glueName the glue resource containing the bytecode
-   * @return the glue bytecode
-   * @throws MissingResourceException if the bytecode cannot be read
-   */
-  @SuppressWarnings({"Since15"})
-  public static byte[] loadBytecode(Class<?> host, String glueName) {
-    String glueResource = GLUE_RESOURCE_PREFIX + glueName;
-    try (InputStream is = Objects.requireNonNull(host.getResourceAsStream(glueResource))) {
-      if (JVM.atLeastJava(9)) {
-        return is.readAllBytes();
-      } else {
-        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
-          int bytesRead;
-          byte[] buf = new byte[BUFFER_SIZE];
-          while ((bytesRead = is.read(buf, 0, BUFFER_SIZE)) != -1) {
-            os.write(buf, 0, bytesRead);
-          }
-          return os.toByteArray();
-        }
-      }
-    } catch (Throwable e) {
-      String detail = "Cannot load " + glueResource + " from " + host + ": " + e;
-      throw new MissingResourceException(detail, host.getName(), glueResource);
-    }
-  }
 }

…a/datadog/instrument/utils/GlueTest.java …g/instrument/glue/GlueGeneratorTest.javautils/src/test/java/datadog/instrument/utils/GlueTest.java renamed to utils/src/test/java/datadog/instrument/glue/GlueGeneratorTest.java utils/src/test/java/datadog/instrument/utils/GlueTest.java renamed to utils/src/test/java/datadog/instrument/glue/GlueGeneratorTest.java

@@ -1,4 +1,4 @@
-package datadog.instrument.utils;
+package datadog.instrument.glue;
 
 import static java.util.Collections.singletonList;
 import static javax.tools.ToolProvider.getSystemJavaCompiler;
@@ -10,38 +10,37 @@
 import java.net.URI;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.MissingResourceException;
 import javax.tools.JavaCompiler;
 import javax.tools.SimpleJavaFileObject;
 import javax.tools.StandardJavaFileManager;
 import javax.tools.StandardLocation;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
-class GlueTest {
-
-  private static final byte[] exampleBytecode = {
-    -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
-    97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
-    116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 7, 69, 120, 97, 109, 112, 108, 101, 1, 0, 4, 67,
-    111, 100, 101, 0, 32, 0, 7, 0, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 5, 0, 6, 0, 1, 0, 9, 0, 0, 0, 17,
-    0, 1, 0, 1, 0, 0, 0, 5, 42, -73, 0, 1, -79, 0, 0, 0, 0, 0, 0
-  };
+class GlueGeneratorTest {
 
   @Test
-  void glueGeneration(@TempDir File classesDir) throws Exception {
+  void glueGenerator(@TempDir File resourceDir, @TempDir File javaDir, @TempDir File classesDir)
+      throws Exception {
+
+    // TestGlueGenerator generates both resource and Java examples of glue
+    GlueGenerator.main(new String[] {resourceDir.getPath(), javaDir.getPath(), "TestGlue"});
+
+    File testGlueResourceFile = new File(resourceDir, "test.glue");
+    File testGlueJavaFile = new File(javaDir, "TestGlue.java");
 
+    assertTrue(testGlueResourceFile.isFile());
+    assertTrue(testGlueJavaFile.isFile());
+
+    // check we can compile the generated Java source
     JavaCompiler compiler = getSystemJavaCompiler();
     StandardJavaFileManager fileManager = compiler.getStandardFileManager(null, null, null);
     fileManager.setLocation(StandardLocation.CLASS_OUTPUT, singletonList(classesDir));
 
-    List<String> glueSource = new ArrayList<>();
-    Glue.classHeader(glueSource, "TestGlue");
-    glueSource.add("  String BYTECODE =");
-    Glue.packBytecode(glueSource, exampleBytecode);
-    glueSource.add("}");
+    List<String> glueSource = Files.readAllLines(testGlueJavaFile.toPath());
 
     List<InMemoryJavaFile> compilationUnits =
         singletonList(new InMemoryJavaFile("datadog/instrument/glue/TestGlue.java", glueSource));
@@ -57,7 +56,16 @@ void glueGeneration(@TempDir File classesDir) throws Exception {
       packedBytecode = (String) testGlueClass.getField("BYTECODE").get(null);
     }
 
-    assertArrayEquals(exampleBytecode, Glue.unpackBytecode(packedBytecode));
+    byte[] expectedBytecode = {
+      -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
+      97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
+      116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 7, 69, 120, 97, 109, 112, 108, 101, 1, 0, 4, 67,
+      111, 100, 101, 0, 32, 0, 7, 0, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 5, 0, 6, 0, 1, 0, 9, 0, 0, 0, 17,
+      0, 1, 0, 1, 0, 0, 0, 5, 42, -73, 0, 1, -79, 0, 0, 0, 0, 0, 0
+    };
+
+    assertArrayEquals(expectedBytecode, Files.readAllBytes(testGlueResourceFile.toPath()));
+    assertArrayEquals(expectedBytecode, Glue.unpackBytecode(packedBytecode));
   }
 
   @Test
@@ -73,20 +81,17 @@ void paddingRequired() {
 
     assertThrows(
         IllegalStateException.class,
-        () -> Glue.packBytecode(new ArrayList<>(), exampleBytecodeNeedsPadding),
+        () -> GlueGenerator.packBytecode(new ArrayList<>(), exampleBytecodeNeedsPadding),
         "Bytecode length is not even; requires padding");
   }
 
   @Test
-  void glueResource() {
-    byte[] loadedBytecode = Glue.loadBytecode(GlueTest.class, "test.glue");
-    assertArrayEquals(exampleBytecode, loadedBytecode);
-  }
+  void notEnoughArgs() {
 
-  @Test
-  void missingGlueResource() {
     assertThrows(
-        MissingResourceException.class, () -> Glue.loadBytecode(GlueTest.class, "missing.glue"));
+        IllegalArgumentException.class,
+        () -> GlueGenerator.main(new String[0]),
+        "Expected: resource-path java-path glue-name...");
   }
 
   static class InMemoryJavaFile extends SimpleJavaFileObject {

utils/src/test/java/datadog/instrument/glue/GlueTest.java

@@ -0,0 +1,30 @@
+package datadog.instrument.glue;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.MissingResourceException;
+import org.junit.jupiter.api.Test;
+
+class GlueTest {
+
+  @Test
+  void glueResource() {
+
+    final byte[] expectedBytecode = {
+      -54, -2, -70, -66, 0, 0, 0, 52, 0, 10, 10, 0, 2, 0, 3, 7, 0, 4, 12, 0, 5, 0, 6, 1, 0, 16, 106,
+      97, 118, 97, 47, 108, 97, 110, 103, 47, 79, 98, 106, 101, 99, 116, 1, 0, 6, 60, 105, 110, 105,
+      116, 62, 1, 0, 3, 40, 41, 86, 7, 0, 8, 1, 0, 7, 69, 120, 97, 109, 112, 108, 101, 1, 0, 4, 67,
+      111, 100, 101, 0, 32, 0, 7, 0, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 5, 0, 6, 0, 1, 0, 9, 0, 0, 0, 17,
+      0, 1, 0, 1, 0, 0, 0, 5, 42, -73, 0, 1, -79, 0, 0, 0, 0, 0, 0
+    };
+
+    assertArrayEquals(expectedBytecode, Glue.loadBytecode(GlueTest.class, "test.glue"));
+  }
+
+  @Test
+  void missingGlueResource() {
+    assertThrows(
+        MissingResourceException.class, () -> Glue.loadBytecode(GlueTest.class, "missing.glue"));
+  }
+}