Skip to content

Commit 0468adc

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Use a new method visitor just for ctos
1 parent 5bc3947 commit 0468adc

7 files changed

Lines changed: 105 additions & 82 deletions

File tree

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java‎

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -286,31 +286,26 @@ private static void writeStackOperations(final AdviceSpecification advice, final
286286
body.addStatement(dupMethod);
287287
} else {
288288
final MethodCallExpr dupMethod = new MethodCallExpr().setScope(new NameExpr("handler"));
289-
if (advice.isConstructor() && advice instanceof AfterSpecification) {
290-
dupMethod.setName("dupConstructor");
289+
if (advice.includeThis()) {
290+
dupMethod.setName("dupInvoke");
291+
dupMethod.addArgument(new NameExpr("owner"));
291292
dupMethod.addArgument(new NameExpr("descriptor"));
292293
} else {
293-
if (advice.includeThis()) {
294-
dupMethod.setName("dupInvoke");
295-
dupMethod.addArgument(new NameExpr("owner"));
296-
dupMethod.addArgument(new NameExpr("descriptor"));
294+
dupMethod.setName("dupParameters");
295+
dupMethod.addArgument(new NameExpr("descriptor"));
296+
}
297+
String mode = "COPY";
298+
if (allArgsSpec != null) {
299+
if (advice instanceof AfterSpecification) {
300+
mode = advice.isConstructor() ? "PREPEND_ARRAY_CTOR" : "PREPEND_ARRAY";
297301
} else {
298-
dupMethod.setName("dupParameters");
299-
dupMethod.addArgument(new NameExpr("descriptor"));
300-
}
301-
String mode = "COPY";
302-
if (allArgsSpec != null) {
303-
if (advice instanceof AfterSpecification) {
304-
mode = "PREPEND_ARRAY";
305-
} else {
306-
mode = "APPEND_ARRAY";
307-
}
302+
mode = "APPEND_ARRAY";
308303
}
309-
dupMethod.addArgument(
310-
new FieldAccessExpr()
311-
.setScope(new TypeExpr(new ClassOrInterfaceType().setName(STACK_DUP_MODE_CLASS)))
312-
.setName(mode));
313304
}
305+
dupMethod.addArgument(
306+
new FieldAccessExpr()
307+
.setScope(new TypeExpr(new ClassOrInterfaceType().setName(STACK_DUP_MODE_CLASS)))
308+
.setName(mode));
314309
body.addStatement(dupMethod);
315310
}
316311
}

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
140140
advices(0) {
141141
pointcut('java/net/URL', '<init>', '(Ljava/lang/String;)V')
142142
statements(
143-
'handler.dupConstructor(descriptor);',
143+
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);',
144144
'handler.method(opcode, owner, name, descriptor, isInterface);',
145145
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;");',
146146
)
@@ -441,7 +441,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
441441
advices(0) {
442442
pointcut('java/lang/StringBuilder', '<init>', '(Ljava/lang/String;)V')
443443
statements(
444-
'handler.dupConstructor(descriptor);',
444+
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);',
445445
'handler.method(opcode, owner, name, descriptor, isInterface);',
446446
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");',
447447
'handler.instruction(Opcodes.CHECKCAST, "java/lang/StringBuilder");'

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java‎

Lines changed: 77 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -118,65 +118,35 @@ public MethodVisitor visitMethod(
118118
final String[] exceptions) {
119119
final MethodVisitor delegated =
120120
super.visitMethod(access, name, descriptor, signature, exceptions);
121-
return new CallSiteMethodVisitor(advices, delegated);
121+
return "<init>".equals(name)
122+
? new CallSiteCtorMethodVisitor(advices, delegated)
123+
: new CallSiteMethodVisitor(advices, delegated);
122124
}
123125
}
124126

125127
private static class CallSiteMethodVisitor extends MethodVisitor
126128
implements CallSiteAdvice.MethodHandler {
127129
private final Advices advices;
128-
private final Deque<String> newInvocations = new LinkedList<>();
129-
private StackDupMode ctorDupMode;
130130

131131
private CallSiteMethodVisitor(
132132
@Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) {
133133
super(ASM_API, delegated);
134134
this.advices = advices;
135135
}
136136

137-
@Override
138-
public void visitEnd() {
139-
super.visitEnd();
140-
if (!newInvocations.isEmpty()) {
141-
LOGGER.debug(
142-
SEND_TELEMETRY,
143-
"There is an issue handling NEW bytecodes, remaining types {}",
144-
newInvocations);
145-
}
146-
}
147-
148-
@Override
149-
public void visitTypeInsn(final int opcode, final String type) {
150-
if (opcode == Opcodes.NEW) {
151-
newInvocations.addLast(type);
152-
}
153-
super.visitTypeInsn(opcode, type);
154-
}
155-
156137
@Override
157138
public void visitMethodInsn(
158139
final int opcode,
159140
final String owner,
160141
final String name,
161142
final String descriptor,
162143
final boolean isInterface) {
163-
try {
164-
if (opcode == Opcodes.INVOKESPECIAL && "<init>".equals(name)) {
165-
if (owner.equals(newInvocations.peekLast())) {
166-
newInvocations.removeLast();
167-
ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_NEW_CTOR;
168-
} else {
169-
ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR;
170-
}
171-
}
172-
CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor);
173-
if (advice instanceof InvokeAdvice) {
174-
((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface);
175-
} else {
176-
mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
177-
}
178-
} finally {
179-
ctorDupMode = null;
144+
145+
CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor);
146+
if (advice instanceof InvokeAdvice) {
147+
((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface);
148+
} else {
149+
mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
180150
}
181151
}
182152

@@ -239,10 +209,6 @@ public void method(
239209

240210
@Override
241211
public void advice(String owner, String name, String descriptor) {
242-
if (ctorDupMode == StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR) {
243-
// append this to the stack after super call
244-
mv.visitIntInsn(Opcodes.ALOAD, 0);
245-
}
246212
mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false);
247213
}
248214

@@ -255,11 +221,6 @@ public void invokeDynamic(
255221
mv.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
256222
}
257223

258-
@Override
259-
public void dupConstructor(final String methodDescriptor) {
260-
dupParameters(methodDescriptor, ctorDupMode);
261-
}
262-
263224
@Override
264225
public void dupParameters(final String methodDescriptor, final StackDupMode mode) {
265226
final Type method = Type.getMethodType(methodDescriptor);
@@ -319,4 +280,72 @@ private Type[] methodParamTypesWithThis(String owner, String methodDescriptor) {
319280
return methodParameterTypesWithThis;
320281
}
321282
}
283+
284+
private static class CallSiteCtorMethodVisitor extends CallSiteMethodVisitor {
285+
286+
private final Deque<String> newInvocations = new LinkedList<>();
287+
private boolean isSuperCall = false;
288+
289+
private CallSiteCtorMethodVisitor(
290+
@Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) {
291+
super(advices, delegated);
292+
}
293+
294+
@Override
295+
public void visitEnd() {
296+
super.visitEnd();
297+
if (!newInvocations.isEmpty()) {
298+
LOGGER.debug(
299+
SEND_TELEMETRY,
300+
"There is an issue handling NEW bytecodes, remaining types {}",
301+
newInvocations);
302+
}
303+
}
304+
305+
@Override
306+
public void visitTypeInsn(final int opcode, final String type) {
307+
if (opcode == Opcodes.NEW) {
308+
newInvocations.addLast(type);
309+
}
310+
super.visitTypeInsn(opcode, type);
311+
}
312+
313+
@Override
314+
public void visitMethodInsn(
315+
final int opcode,
316+
final String owner,
317+
final String name,
318+
final String descriptor,
319+
final boolean isInterface) {
320+
try {
321+
if (opcode == Opcodes.INVOKESPECIAL && "<init>".equals(name)) {
322+
if (owner.equals(newInvocations.peekLast())) {
323+
newInvocations.removeLast();
324+
isSuperCall = false;
325+
} else {
326+
// no new before call to <init>
327+
isSuperCall = true;
328+
}
329+
}
330+
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
331+
} finally {
332+
isSuperCall = false;
333+
}
334+
}
335+
336+
@Override
337+
public void advice(final String owner, final String name, final String descriptor) {
338+
if (isSuperCall) {
339+
// append this to the stack after super call
340+
mv.visitIntInsn(Opcodes.ALOAD, 0);
341+
}
342+
mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false);
343+
}
344+
345+
@Override
346+
public void dupParameters(final String methodDescriptor, final StackDupMode mode) {
347+
super.dupParameters(
348+
methodDescriptor, isSuperCall ? StackDupMode.PREPEND_ARRAY_SUPER_CTOR : mode);
349+
}
350+
}
322351
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,8 @@ public static void dup(final MethodVisitor mv, final Type[] parameters, final St
150150
dup(mv, parameters);
151151
break;
152152
case PREPEND_ARRAY:
153-
case PREPEND_ARRAY_ON_NEW_CTOR:
154-
case PREPEND_ARRAY_ON_SUPER_CTOR:
153+
case PREPEND_ARRAY_CTOR:
154+
case PREPEND_ARRAY_SUPER_CTOR:
155155
case APPEND_ARRAY:
156156
dupN(mv, parameters, mode);
157157
break;
@@ -280,15 +280,15 @@ private static void dupN(
280280
loadArray(mv, arraySize, parameters);
281281
mv.visitInsn(POP);
282282
break;
283-
case PREPEND_ARRAY_ON_NEW_CTOR:
283+
case PREPEND_ARRAY_CTOR:
284284
// move the array before the uninitialized entry created by NEW and DUP
285285
// stack start = [uninitialized, uninitialized, arg_0, ..., arg_n]
286286
// stack end = [array, uninitialized, uninitialized, arg_0, ..., arg_n]
287287
mv.visitInsn(DUP_X2);
288288
loadArray(mv, arraySize, parameters);
289289
mv.visitInsn(POP);
290290
break;
291-
case PREPEND_ARRAY_ON_SUPER_CTOR:
291+
case PREPEND_ARRAY_SUPER_CTOR:
292292
// move the array before the uninitialized entry
293293
// stack start = [uninitialized, arg_0, ..., arg_n]
294294
// stack end = [array, uninitialized, arg_0, ..., arg_n]

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,6 @@ void invokeDynamic(
3737
Handle bootstrapMethodHandle,
3838
Object... bootstrapMethodArguments);
3939

40-
/** Duplicates all the ctor parameters in the stack just before the method is invoked. */
41-
void dupConstructor(String methodDescriptor);
42-
4340
/** Duplicates all the method parameters in the stack just before the method is invoked. */
4441
void dupParameters(String methodDescriptor, StackDupMode mode);
4542

@@ -69,9 +66,9 @@ enum StackDupMode {
6966
/** Copies the parameters in an array and prepends it */
7067
PREPEND_ARRAY,
7168
/** Copies the parameters in an array and adds it between NEW and DUP opcodes */
72-
PREPEND_ARRAY_ON_NEW_CTOR,
69+
PREPEND_ARRAY_CTOR,
7370
/** Copies the parameters in an array and adds it before the uninitialized instance in a ctor */
74-
PREPEND_ARRAY_ON_SUPER_CTOR,
71+
PREPEND_ARRAY_SUPER_CTOR,
7572
/** Copies the parameters in an array and appends it */
7673
APPEND_ARRAY
7774
}

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import net.bytebuddy.jar.asm.Type
88

99
import java.util.concurrent.atomic.AtomicInteger
1010

11+
import static datadog.trace.agent.tooling.csi.CallSiteAdvice.StackDupMode.PREPEND_ARRAY_CTOR
12+
1113
class CallSiteInstrumentationTest extends BaseCallSiteTest {
1214

1315
def 'test instrumentation adds type advice'() {
@@ -73,7 +75,7 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
7375
final InvokeAdvice advice = new InvokeAdvice() {
7476
@Override
7577
void apply(CallSiteAdvice.MethodHandler handler, int opcode, String owner, String name, String descriptor, boolean isInterface) {
76-
handler.dupConstructor(descriptor)
78+
handler.dupParameters(descriptor, PREPEND_ARRAY_CTOR)
7779
handler.method(opcode, owner, name, descriptor, isInterface)
7880
handler.advice(
7981
Type.getType(SuperInCtorExampleAdvice).internalName,

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteUtilsTest.groovy‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ class CallSiteUtilsTest extends DDSpecification {
207207
final visitor = mockMethodVisitor(stack)
208208

209209
when:
210-
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_ON_NEW_CTOR)
210+
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_CTOR)
211211

212212
then: 'the first element of the stack should be an array with the parameters'
213213
final arrayFromStack = stack.remove(0)
@@ -246,7 +246,7 @@ class CallSiteUtilsTest extends DDSpecification {
246246
final visitor = mockMethodVisitor(stack)
247247

248248
when:
249-
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_ON_SUPER_CTOR)
249+
CallSiteUtils.dup(visitor, expected*.type as Type[], PREPEND_ARRAY_SUPER_CTOR)
250250

251251
then: 'the first element of the stack should be an array with the parameters'
252252
final arrayFromStack = stack.remove(0)

0 commit comments

Comments
 (0)