Support aot_training mode to be used when creating AOT caches (#10166)

In this mode only the agent jar is added to the boot classpath, no services are started.

Also workaround potential AOT bug where TraceInterceptor is mistakenly restored from the
system class-loader in production, even though it was visible from the boot class-loader
during training, resulting in LinkageErrors.

Any call to Tracer.addTraceInterceptor from application code in the system class-loader
appears to trigger this bug. The workaround is to replace these calls during training
with opcodes that pop the tracer and argument, and push the expected return value.

This transformation is not persisted, so in production the original method is invoked.

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

* Automatic AOT training detection can be turned off with DD_DETECT_AOT_TRAINING_MODE=false

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java

@@ -0,0 +1,105 @@
+package datadog.trace.bootstrap.aot;
+
+import static datadog.instrument.asm.Opcodes.ACC_ABSTRACT;
+import static datadog.instrument.asm.Opcodes.ACC_NATIVE;
+import static datadog.instrument.asm.Opcodes.ASM9;
+import static datadog.instrument.asm.Opcodes.INVOKEINTERFACE;
+import static datadog.instrument.asm.Opcodes.POP2;
+
+import datadog.instrument.asm.ClassReader;
+import datadog.instrument.asm.ClassVisitor;
+import datadog.instrument.asm.ClassWriter;
+import datadog.instrument.asm.MethodVisitor;
+import java.lang.instrument.ClassFileTransformer;
+import java.security.ProtectionDomain;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Workaround a potential AOT bug where {@link datadog.trace.api.interceptor.TraceInterceptor} is
+ * mistakenly restored from the system class-loader in production, even though it was visible from
+ * the boot class-loader during training, resulting in {@link LinkageError}s.
+ *
+ * <p>Any call to {@link datadog.trace.api.Tracer#addTraceInterceptor} from application code in the
+ * system class-loader appears to trigger this bug. The workaround is to replace these calls during
+ * training with opcodes that pop the tracer and argument, and push the expected return value.
+ *
+ * <p>Note this transformation is not persisted, so in production the original method is invoked.
+ */
+final class TraceApiTransformer implements ClassFileTransformer {
+  private static final ClassLoader SYSTEM_CLASS_LOADER = ClassLoader.getSystemClassLoader();
+
+  @Override
+  public byte[] transform(
+      ClassLoader loader,
+      String className,
+      Class<?> classBeingRedefined,
+      ProtectionDomain pd,
+      byte[] bytecode) {
+
+    // workaround only needed in the system class-loader
+    if (loader == SYSTEM_CLASS_LOADER) {
+      try {
+        ClassReader cr = new ClassReader(bytecode);
+        ClassWriter cw = new ClassWriter(cr, 0);
+        AtomicBoolean modified = new AtomicBoolean();
+        cr.accept(new CallerPatch(cw, modified), 0);
+        // only return something when we've modified the bytecode
+        if (modified.get()) {
+          return cw.toByteArray();
+        }
+      } catch (Throwable ignore) {
+        // skip this class
+      }
+    }
+    return null; // tells the JVM to keep the original bytecode
+  }
+
+  /** Patches callers of {@link datadog.trace.api.Tracer#addTraceInterceptor}. */
+  static final class CallerPatch extends ClassVisitor {
+    private final AtomicBoolean modified;
+
+    CallerPatch(ClassVisitor cv, AtomicBoolean modified) {
+      super(ASM9, cv);
+      this.modified = modified;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(
+        int access, String name, String descriptor, String signature, String[] exceptions) {
+      MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
+      if ((access & (ACC_ABSTRACT | ACC_NATIVE)) == 0) {
+        return new InvokePatch(mv, modified);
+      } else {
+        return mv; // no need to patch abstract/native methods
+      }
+    }
+  }
+
+  /** Removes calls to {@link datadog.trace.api.Tracer#addTraceInterceptor}. */
+  static final class InvokePatch extends MethodVisitor {
+    private final AtomicBoolean modified;
+
+    InvokePatch(MethodVisitor mv, AtomicBoolean modified) {
+      super(ASM9, mv);
+      this.modified = modified;
+    }
+
+    @Override
+    public void visitMethodInsn(
+        int opcode, String owner, String name, String descriptor, boolean isInterface) {
+      if (INVOKEINTERFACE == opcode
+          && "datadog/trace/api/Tracer".equals(owner)
+          && "addTraceInterceptor".equals(name)
+          && "(Ldatadog/trace/api/interceptor/TraceInterceptor;)Z".equals(descriptor)) {
+        // pop tracer and trace interceptor argument from call stack
+        mv.visitInsn(POP2);
+        // push true return value
+        mv.visitLdcInsn(true);
+        // flag that we've modified the bytecode
+        modified.set(true);
+      } else {
+        super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+      }
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TrainingAgent.java

@@ -0,0 +1,19 @@
+package datadog.trace.bootstrap.aot;
+
+import java.lang.instrument.Instrumentation;
+import java.net.URL;
+
+/** Prepares the agent for Ahead-of-Time training. */
+public final class TrainingAgent {
+  public static void start(
+      final Object bootstrapInitTelemetry,
+      final Instrumentation inst,
+      final URL agentJarURL,
+      final String agentArgs) {
+
+    // apply TraceInterceptor LinkageError workaround
+    inst.addTransformer(new TraceApiTransformer());
+
+    // don't start services, they won't be cached as they use a custom classloader
+  }
+}

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

@@ -142,10 +142,17 @@ private static void agentmainImpl(
       recordInstrumentationSource("cmd_line");
     }
 
+    String agentClassName;
+    if (isAotTraining(agentArgs, inst)) {
+      agentClassName = "datadog.trace.bootstrap.aot.TrainingAgent";
+    } else {
+      agentClassName = "datadog.trace.bootstrap.Agent";
+    }
+
     final URL agentJarURL = installAgentJar(inst);
     final Class<?> agentClass;
     try {
-      agentClass = Class.forName("datadog.trace.bootstrap.Agent", true, null);
+      agentClass = Class.forName(agentClassName, true, null);
     } catch (ClassNotFoundException | LinkageError e) {
       throw new IllegalStateException("Unable to load DD Java Agent.", e);
     }
@@ -422,4 +429,17 @@ 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 if ("false".equalsIgnoreCase(EnvironmentVariables.get("DD_DETECT_AOT_TRAINING_MODE"))) {
+      return false; // detection of training mode disabled via DD_DETECT_AOT_TRAINING_MODE=false
+    } 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
+    }
+  }
+}

metadata/supported-configurations.json

@@ -105,6 +105,14 @@
         "aliases": []
       }
     ],
+    "DD_DETECT_AOT_TRAINING_MODE": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": null,
+        "aliases": []
+      }
+    ],
     "DD_API_KEY": [
       {
         "version": "A",