diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/PlaceholderTraceInterceptor.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/PlaceholderTraceInterceptor.java new file mode 100644 index 00000000000..5b18ea5f075 --- /dev/null +++ b/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 onTraceComplete( + Collection trace) { + return trace; // maintain original trace + } + + @Override + public int priority() { + return 0; + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java index bc5c24cb94d..07c5984775e 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java +++ b/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. * - *

Note this transformation is not persisted, so in production the original method is invoked. + *

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. + * + *

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: + * + *

+ */ + 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,6 +87,9 @@ 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 @@ -75,7 +97,7 @@ public MethodVisitor visitMethod( } } - /** 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,11 +112,11 @@ 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 { @@ -102,4 +124,33 @@ public void visitMethodInsn( } } } + + /** 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); + } + } + } }