Support aot_training mode to be used when creating AOT caches

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",