Fix native-image builds with build-time JMX registration

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/AnnotationSubstitutionProcessorInstrumentation.java

@@ -34,6 +34,9 @@ public static class FindTargetClassesAdvice {
     public static void onExit(@Advice.Return(readOnly = false) List<Class<?>> result) {
       result.add(Target_com_datadog_profiling_agent_ProcessContext.class);
       result.add(Target_datadog_jctools_util_UnsafeRefArrayAccess.class);
+      result.add(Target_javax_management_MBeanServerFactory.class);
+      result.add(Target_com_sun_jmx_mbeanserver_MXBeanLookup.class);
+      result.add(Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap.class);
       result.add(Target_org_datadog_jmxfetch_App.class);
       result.add(Target_org_datadog_jmxfetch_Status.class);
       result.add(Target_org_datadog_jmxfetch_reporter_JsonReporter.class);

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/GraalNativeImageModule.java

@@ -14,6 +14,9 @@ public String[] helperClassNames() {
     return new String[] {
       packageName + ".Target_com_datadog_profiling_agent_ProcessContext",
       packageName + ".Target_datadog_jctools_util_UnsafeRefArrayAccess",
+      packageName + ".Target_javax_management_MBeanServerFactory",
+      packageName + ".Target_com_sun_jmx_mbeanserver_MXBeanLookup",
+      packageName + ".Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap",
       packageName + ".Target_org_datadog_jmxfetch_App",
       packageName + ".Target_org_datadog_jmxfetch_Status",
       packageName + ".Target_org_datadog_jmxfetch_reporter_JsonReporter"
@@ -29,6 +32,9 @@ public String[] muzzleIgnoredClassNames() {
       // ignore helper class names as usual
       packageName + ".Target_com_datadog_profiling_agent_ProcessContext",
       packageName + ".Target_datadog_jctools_util_UnsafeRefArrayAccess",
+      packageName + ".Target_javax_management_MBeanServerFactory",
+      packageName + ".Target_com_sun_jmx_mbeanserver_MXBeanLookup",
+      packageName + ".Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap",
       packageName + ".Target_org_datadog_jmxfetch_App",
       packageName + ".Target_org_datadog_jmxfetch_Status",
       packageName + ".Target_org_datadog_jmxfetch_reporter_JsonReporter",

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/Target_com_sun_jmx_mbeanserver_MXBeanLookup.java

@@ -0,0 +1,16 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import com.oracle.svm.core.annotate.Alias;
+import com.oracle.svm.core.annotate.RecomputeFieldValue;
+import com.oracle.svm.core.annotate.TargetClass;
+
+/** Starts the native executable with an empty cache of per-server MXBean lookups. */
+@TargetClass(className = "com.sun.jmx.mbeanserver.MXBeanLookup")
+public final class Target_com_sun_jmx_mbeanserver_MXBeanLookup {
+  @Alias
+  @RecomputeFieldValue(
+      kind = RecomputeFieldValue.Kind.NewInstance,
+      declClassName = "com.sun.jmx.mbeanserver.WeakIdentityHashMap",
+      isFinal = true)
+  private static Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap mbscToLookup;
+}

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap.java

@@ -0,0 +1,7 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import com.oracle.svm.core.annotate.TargetClass;
+
+/** Type alias for the JDK-internal map used by MXBeanLookup. */
+@TargetClass(className = "com.sun.jmx.mbeanserver.WeakIdentityHashMap")
+public final class Target_com_sun_jmx_mbeanserver_WeakIdentityHashMap {}

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/Target_javax_management_MBeanServerFactory.java

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import com.oracle.svm.core.annotate.Alias;
+import com.oracle.svm.core.annotate.RecomputeFieldValue;
+import com.oracle.svm.core.annotate.TargetClass;
+import java.util.ArrayList;
+
+/** Keeps MBean servers created by build-time initializers out of the image heap. */
+@TargetClass(className = "javax.management.MBeanServerFactory")
+public final class Target_javax_management_MBeanServerFactory {
+  @Alias
+  @RecomputeFieldValue(kind = RecomputeFieldValue.Kind.NewInstance, declClass = ArrayList.class)
+  private static ArrayList<?> mBeanServerList;
+}

dd-smoke-tests/native-image-jmx/build.gradle

@@ -0,0 +1,14 @@
+plugins {
+  id 'dd-trace-java.module.smoke-test'
+}
+
+description = 'Native-image build-time JMX regression reproducer'
+
+testJvmConstraints {
+  minJavaVersion = JavaVersion.VERSION_21
+  nativeImageCapable = true
+}
+
+tasks.withType(Test).configureEach {
+  maxParallelForks = 1
+}

dd-smoke-tests/native-image-jmx/gradle.lockfile

@@ -0,0 +1,71 @@
+# This is a Gradle generated file for dependency locking.
+# Manual edits can break the build and are not advised.
+# This file is expected to be part of source control.
+# To regenerate this file, run: ./gradlew :dd-smoke-tests:native-image-jmx:dependencies --write-locks
+ch.qos.logback:logback-classic:1.2.13=testCompileClasspath,testRuntimeClasspath
+ch.qos.logback:logback-core:1.2.13=testCompileClasspath,testRuntimeClasspath
+com.github.javaparser:javaparser-core:3.25.6=codenarc
+com.github.spotbugs:spotbugs-annotations:4.10.4=compileClasspath,spotbugs
+com.github.spotbugs:spotbugs:4.10.4=spotbugs
+com.github.stephenc.jcip:jcip-annotations:1.0-1=spotbugs
+com.google.code.findbugs:jsr305:3.0.2=compileClasspath,spotbugs,testCompileClasspath,testRuntimeClasspath
+com.google.code.gson:gson:2.14.0=spotbugs
+com.google.errorprone:error_prone_annotations:2.48.0=spotbugs
+com.thoughtworks.qdox:qdox:1.12.1=codenarc
+commons-io:commons-io:2.21.0=spotbugs
+de.thetaphi:forbiddenapis:3.11=compileClasspath
+io.leangen.geantyref:geantyref:1.3.16=testRuntimeClasspath
+jaxen:jaxen:2.0.6=spotbugs
+net.bytebuddy:byte-buddy-agent:1.12.8=testRuntimeClasspath
+net.bytebuddy:byte-buddy:1.12.8=testRuntimeClasspath
+net.sf.saxon:Saxon-HE:12.10=spotbugs
+org.apache.ant:ant-antlr:1.10.14=codenarc
+org.apache.ant:ant-junit:1.10.14=codenarc
+org.apache.bcel:bcel:6.12.0=spotbugs
+org.apache.commons:commons-lang3:3.20.0=spotbugs
+org.apache.commons:commons-text:1.15.0=spotbugs
+org.apache.logging.log4j:log4j-api:2.26.1=spotbugs
+org.apache.logging.log4j:log4j-core:2.26.1=spotbugs
+org.apiguardian:apiguardian-api:1.1.2=testCompileClasspath
+org.codehaus.groovy:groovy-ant:3.0.23=codenarc
+org.codehaus.groovy:groovy-docgenerator:3.0.23=codenarc
+org.codehaus.groovy:groovy-groovydoc:3.0.23=codenarc
+org.codehaus.groovy:groovy-json:3.0.23=codenarc
+org.codehaus.groovy:groovy-json:3.0.25=testCompileClasspath,testRuntimeClasspath
+org.codehaus.groovy:groovy-templates:3.0.23=codenarc
+org.codehaus.groovy:groovy-xml:3.0.23=codenarc
+org.codehaus.groovy:groovy:3.0.23=codenarc
+org.codehaus.groovy:groovy:3.0.25=testCompileClasspath,testRuntimeClasspath
+org.codenarc:CodeNarc:3.7.0=codenarc
+org.dom4j:dom4j:2.2.0=spotbugs
+org.gmetrics:GMetrics:2.1.0=codenarc
+org.hamcrest:hamcrest:3.0=testCompileClasspath,testRuntimeClasspath
+org.junit.jupiter:junit-jupiter-api:5.14.1=testCompileClasspath,testRuntimeClasspath
+org.junit.jupiter:junit-jupiter-engine:5.14.1=testRuntimeClasspath
+org.junit.jupiter:junit-jupiter-params:5.14.1=testCompileClasspath,testRuntimeClasspath
+org.junit.jupiter:junit-jupiter:5.14.1=testCompileClasspath,testRuntimeClasspath
+org.junit.platform:junit-platform-commons:1.14.1=testCompileClasspath,testRuntimeClasspath
+org.junit.platform:junit-platform-engine:1.14.1=testCompileClasspath,testRuntimeClasspath
+org.junit.platform:junit-platform-launcher:1.14.1=testRuntimeClasspath
+org.junit:junit-bom:5.14.1=testCompileClasspath,testRuntimeClasspath
+org.mockito:mockito-core:4.4.0=testRuntimeClasspath
+org.objenesis:objenesis:3.3=testCompileClasspath,testRuntimeClasspath
+org.opentest4j:opentest4j:1.3.0=testCompileClasspath,testRuntimeClasspath
+org.ow2.asm:asm-analysis:9.10.1=spotbugs
+org.ow2.asm:asm-commons:9.10.1=spotbugs
+org.ow2.asm:asm-tree:9.10.1=spotbugs
+org.ow2.asm:asm-util:9.10.1=spotbugs
+org.ow2.asm:asm:9.10.1=spotbugs
+org.slf4j:jcl-over-slf4j:1.7.30=testCompileClasspath,testRuntimeClasspath
+org.slf4j:jul-to-slf4j:1.7.30=testCompileClasspath,testRuntimeClasspath
+org.slf4j:log4j-over-slf4j:1.7.30=testCompileClasspath,testRuntimeClasspath
+org.slf4j:slf4j-api:1.7.32=testCompileClasspath,testRuntimeClasspath
+org.slf4j:slf4j-api:2.0.17=spotbugsSlf4j
+org.slf4j:slf4j-api:2.0.18=spotbugs
+org.slf4j:slf4j-simple:2.0.17=spotbugsSlf4j
+org.spockframework:spock-bom:2.4-groovy-3.0=testCompileClasspath,testRuntimeClasspath
+org.spockframework:spock-core:2.4-groovy-3.0=testCompileClasspath,testRuntimeClasspath
+org.tabletest:tabletest-junit:1.2.2=testCompileClasspath,testRuntimeClasspath
+org.tabletest:tabletest-parser:1.2.1=testCompileClasspath,testRuntimeClasspath
+org.xmlresolver:xmlresolver:5.3.3=spotbugs
+empty=annotationProcessor,runtimeClasspath,spotbugsPlugins,testAnnotationProcessor

dd-smoke-tests/native-image-jmx/src/main/java/datadog/smoketest/jmx/JmxApplication.java

@@ -0,0 +1,44 @@
+package datadog.smoketest.jmx;
+
+import static java.lang.management.ManagementFactory.getPlatformMBeanServer;
+
+import javax.management.ObjectName;
+import javax.management.StandardMBean;
+
+/** Reproduces the build-time MXBean registration performed by Oracle JDBC. */
+public class JmxApplication {
+  public static void main(String[] args) {
+    if (!BuildTimeRegistration.INITIALIZED) {
+      throw new AssertionError("Build-time registration did not run");
+    }
+    System.out.println("Native application started");
+  }
+
+  public static class BuildTimeRegistration {
+    static final boolean INITIALIZED;
+
+    static {
+      try {
+        getPlatformMBeanServer()
+            .registerMBean(
+                new StandardMBean(new Diagnostics(), DiagnosticsMXBean.class, true),
+                new ObjectName("datadog.reproducer:type=Diagnostics"));
+        System.out.println("Build-time MXBean registered");
+        INITIALIZED = true;
+      } catch (Exception e) {
+        throw new ExceptionInInitializerError(e);
+      }
+    }
+  }
+
+  public interface DiagnosticsMXBean {
+    boolean isEnabled();
+  }
+
+  public static class Diagnostics implements DiagnosticsMXBean {
+    @Override
+    public boolean isEnabled() {
+      return true;
+    }
+  }
+}

dd-smoke-tests/native-image-jmx/src/main/java/datadog/smoketest/jmx/RuntimeJmxApplication.java

@@ -0,0 +1,30 @@
+package datadog.smoketest.jmx;
+
+import static java.lang.management.ManagementFactory.MEMORY_MXBEAN_NAME;
+import static java.lang.management.ManagementFactory.getPlatformMBeanServer;
+import static javax.management.MBeanServerFactory.findMBeanServer;
+
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+import javax.management.openmbean.CompositeData;
+
+/** Checks that discarding builder JMX state preserves runtime JVM metrics. */
+public class RuntimeJmxApplication {
+  public static void main(String[] args) throws Exception {
+    JmxApplication.main(args);
+
+    MBeanServer server = getPlatformMBeanServer();
+    if (server.isRegistered(new ObjectName("datadog.reproducer:type=Diagnostics"))) {
+      throw new AssertionError("The build-time MXBean leaked into the executable");
+    }
+    if (!findMBeanServer(null).contains(server)) {
+      throw new AssertionError("The runtime MBean server was not registered");
+    }
+    CompositeData heapUsage =
+        (CompositeData) server.getAttribute(new ObjectName(MEMORY_MXBEAN_NAME), "HeapMemoryUsage");
+    if ((Long) heapUsage.get("used") < 0) {
+      throw new AssertionError("Cannot read runtime heap usage through JMX");
+    }
+    System.out.println("Runtime JMX metrics available");
+  }
+}

dd-smoke-tests/native-image-jmx/src/test/java/datadog/smoketest/jmx/NativeImageJmxSmokeTest.java

@@ -0,0 +1,108 @@
+package datadog.smoketest.jmx;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+class NativeImageJmxSmokeTest {
+  @Test
+  void buildsWithAgentAfterBuildTimeMxBeanRegistration() throws Exception {
+    Path directory = newBuildDirectory();
+
+    // The control must pass first: the fixture alone must not retain the builder's MBean server.
+    buildAndRun(directory, false, JmxApplication.class);
+    buildAndRun(directory, true, JmxApplication.class);
+  }
+
+  @Test
+  void supportsRuntimeJmxAfterBuildTimeMxBeanRegistration() throws Exception {
+    assertTrue(
+        buildAndRun(newBuildDirectory(), true, RuntimeJmxApplication.class)
+            .contains("Runtime JMX metrics available"));
+  }
+
+  private Path newBuildDirectory() throws Exception {
+    Path buildDirectory = Paths.get(System.getProperty("datadog.smoketest.builddir"));
+    Files.createDirectories(buildDirectory);
+    return Files.createTempDirectory(buildDirectory, "native-image-jmx-");
+  }
+
+  private String buildAndRun(Path directory, boolean withAgent, Class<?> application)
+      throws Exception {
+    String name = withAgent ? "with-agent" : "without-agent";
+    Path executable = directory.resolve(name);
+    Path buildLog = directory.resolve(name + "-build.log");
+    List<String> command =
+        new ArrayList<>(
+            Arrays.asList(
+                Paths.get(System.getProperty("java.home"), "bin", "native-image").toString(),
+                "-O0",
+                "--no-fallback",
+                "--parallelism=4",
+                "-J-Xmx4g",
+                // Match Quarkus's default; the tracer supplies its own runtime exceptions.
+                "--initialize-at-build-time",
+                "-cp",
+                new File(
+                        JmxApplication.class
+                            .getProtectionDomain()
+                            .getCodeSource()
+                            .getLocation()
+                            .toURI())
+                    .getAbsolutePath()));
+    if (withAgent) {
+      command.add("-J-javaagent:" + System.getProperty("datadog.smoketest.agent.shadowJar.path"));
+    }
+    command.add(application.getName());
+    command.add(executable.toString());
+
+    int exitCode = run(command, buildLog, MINUTES.toSeconds(10));
+    String output = new String(Files.readAllBytes(buildLog), UTF_8);
+    assertTrue(
+        output.contains("Build-time MXBean registered"), () -> "Fixture did not run: " + buildLog);
+    assertEquals(
+        0,
+        exitCode,
+        () ->
+            "Native-image build failed ("
+                + name
+                + "). It reproduces when the agent retains the builder's JmxMBeanServer. Log: "
+                + buildLog
+                + "\n"
+                + output);
+
+    Path runLog = directory.resolve(name + "-run.log");
+    assertEquals(0, run(Arrays.asList(executable.toString()), runLog, 30), () -> "See " + runLog);
+    String runtimeOutput = new String(Files.readAllBytes(runLog), UTF_8);
+    assertTrue(runtimeOutput.contains("Native application started"), () -> "See " + runLog);
+    assertFalse(
+        runtimeOutput.contains("Build-time MXBean registered"),
+        "Registration must happen during image building, not application startup");
+    return runtimeOutput;
+  }
+
+  private int run(List<String> command, Path log, long timeoutSeconds) throws Exception {
+    Process process =
+        new ProcessBuilder(command).redirectErrorStream(true).redirectOutput(log.toFile()).start();
+    try {
+      assertTrue(
+          process.waitFor(timeoutSeconds, SECONDS),
+          () -> "Timed out: " + command + "; log: " + log);
+      return process.exitValue();
+    } finally {
+      process.destroyForcibly();
+    }
+  }
+}

settings.gradle.kts

@@ -212,6 +212,7 @@ include(
   ":dd-smoke-tests:lib-injection",
   ":dd-smoke-tests:log-injection",
   ":dd-smoke-tests:maven",
+  ":dd-smoke-tests:native-image-jmx",
   ":dd-smoke-tests:openfeature",
   ":dd-smoke-tests:openfeature-legacy",
   ":dd-smoke-tests:opentracing",