From 9ecadb335620b91cd44f18f272806a6843c26278 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Sun, 24 Nov 2024 21:28:51 +0100 Subject: [PATCH] Use a new method visitor just for ctos --- .../plugin/csi/impl/AdviceGeneratorImpl.java | 35 +++-- .../csi/impl/AdviceGeneratorTest.groovy | 4 +- .../bytebuddy/csi/CallSiteTransformer.java | 125 +++++++++++------- .../tooling/bytebuddy/csi/CallSiteUtils.java | 8 +- .../agent/tooling/csi/CallSiteAdvice.java | 7 +- .../csi/CallSiteInstrumentationTest.groovy | 4 +- .../tooling/csi/CallSiteUtilsTest.groovy | 2 +- 7 files changed, 104 insertions(+), 81 deletions(-) diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java index c966d683634..04c7a9be704 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java +++ b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java @@ -286,31 +286,26 @@ private static void writeStackOperations(final AdviceSpecification advice, final body.addStatement(dupMethod); } else { final MethodCallExpr dupMethod = new MethodCallExpr().setScope(new NameExpr("handler")); - if (advice.isConstructor() && advice instanceof AfterSpecification) { - dupMethod.setName("dupConstructor"); + if (advice.includeThis()) { + dupMethod.setName("dupInvoke"); + dupMethod.addArgument(new NameExpr("owner")); dupMethod.addArgument(new NameExpr("descriptor")); } else { - if (advice.includeThis()) { - dupMethod.setName("dupInvoke"); - dupMethod.addArgument(new NameExpr("owner")); - dupMethod.addArgument(new NameExpr("descriptor")); + dupMethod.setName("dupParameters"); + dupMethod.addArgument(new NameExpr("descriptor")); + } + String mode = "COPY"; + if (allArgsSpec != null) { + if (advice instanceof AfterSpecification) { + mode = advice.isConstructor() ? "PREPEND_ARRAY_CTOR" : "PREPEND_ARRAY"; } else { - dupMethod.setName("dupParameters"); - dupMethod.addArgument(new NameExpr("descriptor")); - } - String mode = "COPY"; - if (allArgsSpec != null) { - if (advice instanceof AfterSpecification) { - mode = "PREPEND_ARRAY"; - } else { - mode = "APPEND_ARRAY"; - } + mode = "APPEND_ARRAY"; } - dupMethod.addArgument( - new FieldAccessExpr() - .setScope(new TypeExpr(new ClassOrInterfaceType().setName(STACK_DUP_MODE_CLASS))) - .setName(mode)); } + dupMethod.addArgument( + new FieldAccessExpr() + .setScope(new TypeExpr(new ClassOrInterfaceType().setName(STACK_DUP_MODE_CLASS))) + .setName(mode)); body.addStatement(dupMethod); } } diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy index db18005f8e4..a0b3e15e2cc 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy +++ b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy @@ -140,7 +140,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { advices(0) { pointcut('java/net/URL', '', '(Ljava/lang/String;)V') statements( - 'handler.dupConstructor(descriptor);', + 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);', 'handler.method(opcode, owner, name, descriptor, isInterface);', 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;");', ) @@ -441,7 +441,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest { advices(0) { pointcut('java/lang/StringBuilder', '', '(Ljava/lang/String;)V') statements( - 'handler.dupConstructor(descriptor);', + 'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);', 'handler.method(opcode, owner, name, descriptor, isInterface);', 'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");', 'handler.instruction(Opcodes.CHECKCAST, "java/lang/StringBuilder");' diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java index bfa2a75befe..991523069c1 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java @@ -118,15 +118,15 @@ public MethodVisitor visitMethod( final String[] exceptions) { final MethodVisitor delegated = super.visitMethod(access, name, descriptor, signature, exceptions); - return new CallSiteMethodVisitor(advices, delegated); + return "".equals(name) + ? new CallSiteCtorMethodVisitor(advices, delegated) + : new CallSiteMethodVisitor(advices, delegated); } } private static class CallSiteMethodVisitor extends MethodVisitor implements CallSiteAdvice.MethodHandler { private final Advices advices; - private final Deque newInvocations = new LinkedList<>(); - private StackDupMode ctorDupMode; private CallSiteMethodVisitor( @Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) { @@ -134,25 +134,6 @@ private CallSiteMethodVisitor( this.advices = advices; } - @Override - public void visitEnd() { - super.visitEnd(); - if (!newInvocations.isEmpty()) { - LOGGER.debug( - SEND_TELEMETRY, - "There is an issue handling NEW bytecodes, remaining types {}", - newInvocations); - } - } - - @Override - public void visitTypeInsn(final int opcode, final String type) { - if (opcode == Opcodes.NEW) { - newInvocations.addLast(type); - } - super.visitTypeInsn(opcode, type); - } - @Override public void visitMethodInsn( final int opcode, @@ -160,23 +141,12 @@ public void visitMethodInsn( final String name, final String descriptor, final boolean isInterface) { - try { - if (opcode == Opcodes.INVOKESPECIAL && "".equals(name)) { - if (owner.equals(newInvocations.peekLast())) { - newInvocations.removeLast(); - ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_NEW_CTOR; - } else { - ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR; - } - } - CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor); - if (advice instanceof InvokeAdvice) { - ((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface); - } else { - mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface); - } - } finally { - ctorDupMode = null; + + CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor); + if (advice instanceof InvokeAdvice) { + ((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface); + } else { + mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface); } } @@ -239,10 +209,6 @@ public void method( @Override public void advice(String owner, String name, String descriptor) { - if (ctorDupMode == StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR) { - // append this to the stack after super call - mv.visitIntInsn(Opcodes.ALOAD, 0); - } mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false); } @@ -255,11 +221,6 @@ public void invokeDynamic( mv.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); } - @Override - public void dupConstructor(final String methodDescriptor) { - dupParameters(methodDescriptor, ctorDupMode); - } - @Override public void dupParameters(final String methodDescriptor, final StackDupMode mode) { final Type method = Type.getMethodType(methodDescriptor); @@ -319,4 +280,72 @@ private Type[] methodParamTypesWithThis(String owner, String methodDescriptor) { return methodParameterTypesWithThis; } } + + private static class CallSiteCtorMethodVisitor extends CallSiteMethodVisitor { + + private final Deque newInvocations = new LinkedList<>(); + private boolean isSuperCall = false; + + private CallSiteCtorMethodVisitor( + @Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) { + super(advices, delegated); + } + + @Override + public void visitEnd() { + super.visitEnd(); + if (!newInvocations.isEmpty()) { + LOGGER.debug( + SEND_TELEMETRY, + "There is an issue handling NEW bytecodes, remaining types {}", + newInvocations); + } + } + + @Override + public void visitTypeInsn(final int opcode, final String type) { + if (opcode == Opcodes.NEW) { + newInvocations.addLast(type); + } + super.visitTypeInsn(opcode, type); + } + + @Override + public void visitMethodInsn( + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + try { + if (opcode == Opcodes.INVOKESPECIAL && "".equals(name)) { + if (owner.equals(newInvocations.peekLast())) { + newInvocations.removeLast(); + isSuperCall = false; + } else { + // no new before call to + isSuperCall = true; + } + } + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + } finally { + isSuperCall = false; + } + } + + @Override + public void advice(final String owner, final String name, final String descriptor) { + if (isSuperCall) { + // append this to the stack after super call + mv.visitIntInsn(Opcodes.ALOAD, 0); + } + mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false); + } + + @Override + public void dupParameters(final String methodDescriptor, final StackDupMode mode) { + super.dupParameters( + methodDescriptor, isSuperCall ? StackDupMode.PREPEND_ARRAY_SUPER_CTOR : mode); + } + } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java index 503f8d2b131..39511a60a76 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java @@ -150,8 +150,8 @@ public static void dup(final MethodVisitor mv, final Type[] parameters, final St dup(mv, parameters); break; case PREPEND_ARRAY: - case PREPEND_ARRAY_ON_NEW_CTOR: - case PREPEND_ARRAY_ON_SUPER_CTOR: + case PREPEND_ARRAY_CTOR: + case PREPEND_ARRAY_SUPER_CTOR: case APPEND_ARRAY: dupN(mv, parameters, mode); break; @@ -280,7 +280,7 @@ private static void dupN( loadArray(mv, arraySize, parameters); mv.visitInsn(POP); break; - case PREPEND_ARRAY_ON_NEW_CTOR: + case PREPEND_ARRAY_CTOR: // move the array before the uninitialized entry created by NEW and DUP // stack start = [uninitialized, uninitialized, arg_0, ..., arg_n] // stack end = [array, uninitialized, uninitialized, arg_0, ..., arg_n] @@ -288,7 +288,7 @@ private static void dupN( loadArray(mv, arraySize, parameters); mv.visitInsn(POP); break; - case PREPEND_ARRAY_ON_SUPER_CTOR: + case PREPEND_ARRAY_SUPER_CTOR: // move the array before the uninitialized entry // stack start = [uninitialized, arg_0, ..., arg_n] // stack end = [array, uninitialized, arg_0, ..., arg_n] diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java index d1ada34f37f..0e43f3602e8 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java @@ -37,9 +37,6 @@ void invokeDynamic( Handle bootstrapMethodHandle, Object... bootstrapMethodArguments); - /** Duplicates all the ctor parameters in the stack just before the method is invoked. */ - void dupConstructor(String methodDescriptor); - /** Duplicates all the method parameters in the stack just before the method is invoked. */ void dupParameters(String methodDescriptor, StackDupMode mode); @@ -69,9 +66,9 @@ enum StackDupMode { /** Copies the parameters in an array and prepends it */ PREPEND_ARRAY, /** Copies the parameters in an array and adds it between NEW and DUP opcodes */ - PREPEND_ARRAY_ON_NEW_CTOR, + PREPEND_ARRAY_CTOR, /** Copies the parameters in an array and adds it before the uninitialized instance in a ctor */ - PREPEND_ARRAY_ON_SUPER_CTOR, + PREPEND_ARRAY_SUPER_CTOR, /** Copies the parameters in an array and appends it */ APPEND_ARRAY } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy index bc69370edd0..b56b94779ad 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy @@ -8,6 +8,8 @@ import net.bytebuddy.jar.asm.Type import java.util.concurrent.atomic.AtomicInteger +import static datadog.trace.agent.tooling.csi.CallSiteAdvice.StackDupMode.PREPEND_ARRAY_CTOR + class CallSiteInstrumentationTest extends BaseCallSiteTest { def 'test instrumentation adds type advice'() { @@ -73,7 +75,7 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest { final InvokeAdvice advice = new InvokeAdvice() { @Override void apply(CallSiteAdvice.MethodHandler handler, int opcode, String owner, String name, String descriptor, boolean isInterface) { - handler.dupConstructor(descriptor) + handler.dupParameters(descriptor, PREPEND_ARRAY_CTOR) handler.method(opcode, owner, name, descriptor, isInterface) handler.advice( Type.getType(SuperInCtorExampleAdvice).internalName, diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy index 29b98b7007f..6adb76d702c 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy @@ -246,7 +246,7 @@ class CallSiteUtilsTest extends DDSpecification { final visitor = mockMethodVisitor(stack) when: - CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_ON_SUPER_CTOR) + CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_SUPER_CTOR) then: 'the first element of the stack should be an array with the parameters' final arrayFromStack = stack.remove(0)