Support 'aot_training' mode to be used when creating AOT caches

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.

Author: mcculls

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

@@ -0,0 +1,104 @@
+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 {
+
+  @Override
+  public byte[] transform(
+      ClassLoader loader,
+      String className,
+      Class<?> classBeingRedefined,
+      ProtectionDomain pd,
+      byte[] bytecode) {
+
+    // workaround only needed in the system class-loader
+    if (loader == ClassLoader.getSystemClassLoader()) {
+      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/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java

@@ -142,10 +142,17 @@ private static void agentmainImpl(
       recordInstrumentationSource("cmd_line");
     }
 
+    String agentClassName;
+    if ("aot_training".equalsIgnoreCase(agentArgs)) {
+      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);
     }