Support creating AOT caches for Spring apps with custom TraceInterceptors

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

@@ -0,0 +1,23 @@
+package datadog.trace.bootstrap.aot;
+
+import datadog.trace.api.interceptor.MutableSpan;
+import datadog.trace.api.interceptor.TraceInterceptor;
+import java.util.Collection;
+
+/** Placeholder {@link TraceInterceptor} used to temporarily work around an AOT bug. */
+public final class PlaceholderTraceInterceptor implements TraceInterceptor {
+  public static final PlaceholderTraceInterceptor INSTANCE = new PlaceholderTraceInterceptor();
+
+  private PlaceholderTraceInterceptor() {}
+
+  @Override
+  public Collection<? extends MutableSpan> onTraceComplete(
+      Collection<? extends MutableSpan> trace) {
+    return trace; // maintain original trace
+  }
+
+  @Override
+  public int priority() {
+    return 0;
+  }
+}

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

@@ -2,8 +2,12 @@
 
 import static datadog.instrument.asm.Opcodes.ACC_ABSTRACT;
 import static datadog.instrument.asm.Opcodes.ACC_NATIVE;
+import static datadog.instrument.asm.Opcodes.ARETURN;
 import static datadog.instrument.asm.Opcodes.ASM9;
+import static datadog.instrument.asm.Opcodes.GETSTATIC;
+import static datadog.instrument.asm.Opcodes.ICONST_1;
 import static datadog.instrument.asm.Opcodes.INVOKEINTERFACE;
+import static datadog.instrument.asm.Opcodes.POP;
 import static datadog.instrument.asm.Opcodes.POP2;
 
 import datadog.instrument.asm.ClassReader;
@@ -23,11 +27,19 @@
  * 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.
+ * <p>Likewise, custom {@code TraceInterceptor} return values are replaced with simple placeholders
+ * from the boot class-path which are guaranteed not to trigger the AOT bug.
+ *
+ * <p>Note these transformations are not persisted, so in production the original code is used.
  */
 final class TraceApiTransformer implements ClassFileTransformer {
   private static final ClassLoader SYSTEM_CLASS_LOADER = ClassLoader.getSystemClassLoader();
 
+  static final String TRACE_INTERCEPTOR = "datadog/trace/api/interceptor/TraceInterceptor";
+
+  static final String PLACEHOLDER_TRACE_INTERCEPTOR =
+      "datadog/trace/bootstrap/aot/PlaceholderTraceInterceptor";
+
   @Override
   public byte[] transform(
       ClassLoader loader,
@@ -42,7 +54,7 @@ public byte[] transform(
         ClassReader cr = new ClassReader(bytecode);
         ClassWriter cw = new ClassWriter(cr, 0);
         AtomicBoolean modified = new AtomicBoolean();
-        cr.accept(new CallerPatch(cw, modified), 0);
+        cr.accept(new TraceInterceptorPatch(cw, modified), 0);
         // only return something when we've modified the bytecode
         if (modified.get()) {
           return cw.toByteArray();
@@ -54,11 +66,18 @@ public byte[] transform(
     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 {
+  /**
+   * Patches certain references to {@code TraceInterceptor} to workaround AOT bug:
+   *
+   * <ul>
+   *   <li>removes direct calls to {@code Tracer.addTraceInterceptor()}
+   *   <li>replaces {@code TraceInterceptor} return values with placeholders
+   * </ul>
+   */
+  static final class TraceInterceptorPatch extends ClassVisitor {
     private final AtomicBoolean modified;
 
-    CallerPatch(ClassVisitor cv, AtomicBoolean modified) {
+    TraceInterceptorPatch(ClassVisitor cv, AtomicBoolean modified) {
       super(ASM9, cv);
       this.modified = modified;
     }
@@ -68,14 +87,17 @@ 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) {
+        if (descriptor.endsWith(")L" + TRACE_INTERCEPTOR + ";")) {
+          mv = new ReturnPatch(mv, modified);
+        }
         return new InvokePatch(mv, modified);
       } else {
         return mv; // no need to patch abstract/native methods
       }
     }
   }
 
-  /** Removes calls to {@link datadog.trace.api.Tracer#addTraceInterceptor}. */
+  /** Removes direct calls to {@code Tracer.addTraceInterceptor()}. */
   static final class InvokePatch extends MethodVisitor {
     private final AtomicBoolean modified;
 
@@ -90,16 +112,45 @@ public void visitMethodInsn(
       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
+          && ("(L" + TRACE_INTERCEPTOR + ";)Z").equals(descriptor)) {
+        // discard tracer and trace interceptor argument from call stack
         mv.visitInsn(POP2);
-        // push true return value
-        mv.visitLdcInsn(true);
+        // push substitute return value (true)
+        mv.visitInsn(ICONST_1);
         // flag that we've modified the bytecode
         modified.set(true);
       } else {
         super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
       }
     }
   }
+
+  /** Replaces custom {@code TraceInterceptor} return values with placeholders. */
+  static final class ReturnPatch extends MethodVisitor {
+    private final AtomicBoolean modified;
+
+    ReturnPatch(MethodVisitor mv, AtomicBoolean modified) {
+      super(ASM9, mv);
+      this.modified = modified;
+    }
+
+    @Override
+    public void visitInsn(int opcode) {
+      if (ARETURN == opcode) {
+        // discard old return value from call stack
+        mv.visitInsn(POP);
+        // push our placeholder interceptor instead
+        mv.visitFieldInsn(
+            GETSTATIC,
+            PLACEHOLDER_TRACE_INTERCEPTOR,
+            "INSTANCE",
+            "L" + PLACEHOLDER_TRACE_INTERCEPTOR + ";");
+        mv.visitInsn(ARETURN);
+        // flag that we've modified the bytecode
+        modified.set(true);
+      } else {
+        super.visitInsn(opcode);
+      }
+    }
+  }
 }