Skip to content

Commit

Permalink
Use a new method visitor just for ctos
Browse files Browse the repository at this point in the history
  • Loading branch information
@manuel-alvarez-alvarez
manuel-alvarez-alvarez committed Nov 24, 2024
1 parent 5bc3947 commit 0468adc
Show file tree
Hide file tree
Showing 7 changed files with 105 additions and 82 deletions.
35 changes: 15 additions & 20 deletions ...strumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...mentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
advices(0) {
pointcut('java/net/URL', '<init>', '(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;");',
)
Expand Down Expand Up @@ -441,7 +441,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
advices(0) {
pointcut('java/lang/StringBuilder', '<init>', '(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");'
Expand Down
125 changes: 77 additions & 48 deletions ...-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,65 +118,35 @@ public MethodVisitor visitMethod(
final String[] exceptions) {
final MethodVisitor delegated =
super.visitMethod(access, name, descriptor, signature, exceptions);
return new CallSiteMethodVisitor(advices, delegated);
return "<init>".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<String> newInvocations = new LinkedList<>();
private StackDupMode ctorDupMode;

private CallSiteMethodVisitor(
@Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) {
super(ASM_API, delegated);
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,
final String owner,
final String name,
final String descriptor,
final boolean isInterface) {
try {
if (opcode == Opcodes.INVOKESPECIAL && "<init>".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);
}
}

Expand Down Expand Up @@ -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);
}

Expand All @@ -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);
Expand Down Expand Up @@ -319,4 +280,72 @@ private Type[] methodParamTypesWithThis(String owner, String methodDescriptor) {
return methodParameterTypesWithThis;
}
}

private static class CallSiteCtorMethodVisitor extends CallSiteMethodVisitor {

private final Deque<String> 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 && "<init>".equals(name)) {
if (owner.equals(newInvocations.peekLast())) {
newInvocations.removeLast();
isSuperCall = false;
} else {
// no new before call to <init>
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);
}
}
}
8 changes: 4 additions & 4 deletions .../agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -280,15 +280,15 @@ 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]
mv.visitInsn(DUP_X2);
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]
Expand Down
7 changes: 2 additions & 5 deletions ...ava-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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
}
Expand Down
4 changes: 3 additions & 1 deletion ...ooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'() {
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions ...nt/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ class CallSiteUtilsTest extends DDSpecification {
final visitor = mockMethodVisitor(stack)

when:
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_ON_NEW_CTOR)
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_CTOR)

then: 'the first element of the stack should be an array with the parameters'
final arrayFromStack = stack.remove(0)
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 0468adc

Please sign in to comment.