Add automatic check for AOT training mode, based on CDS.isDumpingArchive()

Author: mcculls

dd-java-agent/build.gradle

@@ -23,24 +23,35 @@ tasks.named("processResources") {
   dependsOn(includedJarFileTree)
 }
 
-// The special pre-check should be compiled with Java 6 to detect unsupported Java versions
-// and prevent issues for users that still using them.
 sourceSets {
+  // The special pre-check must be compiled with Java 6 to detect unsupported
+  // Java versions and prevent issues for users that still using them.
   "main_java6" {
     java.srcDirs "${project.projectDir}/src/main/java6"
   }
+  // Additional checks that use the Java 11 API.
+  "main_java11" {
+    java.srcDirs "${project.projectDir}/src/main/java11"
+  }
   main.resources.srcDir(includedAgentDir)
 }
 
 def java6CompileTask = tasks.named("compileMain_java6Java") {
   configureCompiler(it, 8, JavaVersion.VERSION_1_6)
 }
 
+def java11CompileTask = tasks.named("compileMain_java11Java") {
+  configureCompiler(it, 11)
+}
+
 tasks.named("compileJava") {
   dependsOn(java6CompileTask)
+  dependsOn(java11CompileTask)
 }
 
 dependencies {
+  implementation sourceSets.main_java11.output
+  main_java11CompileOnly libs.forbiddenapis
   main_java6CompileOnly libs.forbiddenapis
   testImplementation sourceSets.main_java6.output
 }
@@ -249,6 +260,8 @@ includeShadowJar(traceShadowJar, 'trace', includedJarFileTree)
 tasks.named("shadowJar", ShadowJar) {
   // Include AgentPreCheck compiled with Java 6.
   from sourceSets.main_java6.output
+  // Include additional checks compiled with Java 11.
+  from sourceSets.main_java11.output
 
   generalShadowJarConfig(it)
 

dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java

@@ -143,7 +143,7 @@ private static void agentmainImpl(
     }
 
     String agentClassName;
-    if ("aot_training".equalsIgnoreCase(agentArgs)) {
+    if (isAotTraining(agentArgs, inst)) {
       agentClassName = "datadog.trace.bootstrap.aot.TrainingAgent";
     } else {
       agentClassName = "datadog.trace.bootstrap.Agent";
@@ -429,4 +429,15 @@ private static void checkJarManifestMainClassIsThis(final URL jarUrl) throws IOE
             + jarUrl
             + "'. Make sure you don't have this .class-file anywhere, besides dd-java-agent.jar");
   }
+
+  /** Returns {@code true} if the JVM is training, i.e. writing to a CDS/AOT archive. */
+  private static boolean isAotTraining(String agentArgs, Instrumentation inst) {
+    if (!JavaVirtualMachine.isJavaVersionAtLeast(25)) {
+      return false; // agent doesn't support training mode before Java 25
+    } else if ("aot_training".equalsIgnoreCase(agentArgs)) {
+      return true; // training mode explicitly enabled via -javaagent
+    } else {
+      return AdvancedAgentChecks.isAotTraining(inst); // check JVM status
+    }
+  }
 }

dd-java-agent/src/main/java11/datadog/trace/bootstrap/AdvancedAgentChecks.java

@@ -0,0 +1,39 @@
+package datadog.trace.bootstrap;
+
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.emptySet;
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonMap;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.lang.instrument.Instrumentation;
+import java.lang.reflect.Method;
+
+/** Additional agent checks that require Java 11+. */
+public final class AdvancedAgentChecks {
+
+  /** Returns {@code true} if the JVM is writing to a CDS/AOT archive, i.e. is in training mode. */
+  @SuppressForbidden
+  public static boolean isAotTraining(Instrumentation inst) {
+    try {
+      Class<?> cds = Class.forName("jdk.internal.misc.CDS");
+
+      // ensure the module containing CDS exports it to our unnamed module
+      Module cdsModule = cds.getModule();
+      Module unnamedModule = AdvancedAgentChecks.class.getClassLoader().getUnnamedModule();
+      inst.redefineModule(
+          cdsModule,
+          emptySet(),
+          singletonMap("jdk.internal.misc", singleton(unnamedModule)),
+          emptyMap(),
+          emptySet(),
+          emptyMap());
+
+      // if the JVM is writing to a CDS/AOT archive then it's in training mode
+      Method isDumpingArchive = cds.getMethod("isDumpingArchive");
+      return (boolean) isDumpingArchive.invoke(null);
+    } catch (Throwable ignore) {
+      return false; // if we don't have access then assume we're not training
+    }
+  }
+}