Skip to content

Commit 8d4a9ea

Browse files
PerfectSlayerPerfectSlayer
committed
wip(asm): Fix startSpan hook
1 parent 478bffa commit 8d4a9ea

File tree

2 files changed

+40
-20
lines changed

2 files changed

+40
-20
lines changed

‎dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty/HandleRequestVisitor.java

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.jetty;
22

3+
import datadog.context.Context;
34
import datadog.trace.api.gateway.Flow;
45
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
56
import java.util.List;
@@ -15,6 +16,7 @@
1516
* and replaces it with:
1617
*
1718
* <pre>
19+
* AgentSpan span = Java8BytecodeBridge(context);
1820
* if (span != null && span.getRequestBlockingAction() &&
1921
* JettyBlockingHelper.block(
2022
* this.getRequest(), this.getResponse(),
@@ -31,7 +33,7 @@ public class HandleRequestVisitor extends MethodVisitor {
3133
private static final Logger log = LoggerFactory.getLogger(HandleRequestVisitor.class);
3234

3335
private boolean lookForStore;
34-
private int agentSpanVar = -1;
36+
private int contextVar = -1;
3537
private final int classVersion;
3638
private final String connClassInternalName;
3739
private boolean success;
@@ -53,12 +55,12 @@ DelayLoadsMethodVisitor delayVisitorDelegate() {
5355
@Override
5456
public void visitMethodInsn(
5557
int opcode, String owner, String name, String descriptor, boolean isInterface) {
56-
if (agentSpanVar == -1) {
58+
if (contextVar == -1) {
5759
lookForStore =
5860
!lookForStore
5961
&& opcode == Opcodes.INVOKEVIRTUAL
6062
&& name.equals("startSpan")
61-
&& descriptor.endsWith("Ldatadog/trace/bootstrap/instrumentation/api/AgentSpan;");
63+
&& descriptor.endsWith("Ldatadog.context.Context;");
6264
} else if (opcode == Opcodes.INVOKEVIRTUAL
6365
&& owner.equals("org/eclipse/jetty/server/Server")
6466
&& name.equals("handle")
@@ -75,9 +77,17 @@ public void visitMethodInsn(
7577
Label beforeHandle = new Label();
7678
Label afterHandle = new Label();
7779

78-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
80+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
81+
super.visitMethodInsn(
82+
Opcodes.INVOKESTATIC,
83+
"datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge",
84+
"spanFromContext",
85+
"(" + Type.getDescriptor(Context.class) + ")" + Type.getDescriptor(AgentSpan.class),
86+
false);
87+
super.visitVarInsn(Opcodes.ASTORE, contextVar);
88+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
7989
super.visitJumpInsn(Opcodes.IFNULL, beforeHandle);
80-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
90+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
8191
super.visitMethodInsn(
8292
Opcodes.INVOKEINTERFACE,
8393
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
@@ -105,14 +115,14 @@ public void visitMethodInsn(
105115
"getResponse",
106116
"()Lorg/eclipse/jetty/server/Response;",
107117
false);
108-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
118+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
109119
super.visitMethodInsn(
110120
Opcodes.INVOKEINTERFACE,
111121
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
112122
"getRequestBlockingAction",
113123
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
114124
true);
115-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
125+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
116126
super.visitMethodInsn(
117127
Opcodes.INVOKESTATIC,
118128
Type.getInternalName(JettyBlockingHelper.class),
@@ -144,7 +154,7 @@ public void visitMethodInsn(
144154
@Override
145155
public void visitVarInsn(int opcode, int varIndex) {
146156
if (lookForStore && opcode == Opcodes.ASTORE) {
147-
agentSpanVar = varIndex;
157+
contextVar = varIndex;
148158
lookForStore = false;
149159
}
150160

‎dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty9/HandleVisitor.java

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.jetty9;
22

3+
import datadog.context.Context;
34
import datadog.trace.api.gateway.Flow;
45
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
56
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
@@ -25,6 +26,7 @@
2526
* </code> is replaced with: <code>
2627
* case REQUEST_DISPATCH:
2728
* // ...
29+
* AgentSpan span = Java8BytecodeBridge(context);
2830
* if (span != null && span.getBlockingAction() != null &&
2931
* JettyBlockingHelper.block(this.getRequest(), this.getResponse())) {
3032
* // nothing
@@ -82,7 +84,7 @@ public class HandleVisitor extends MethodVisitor {
8284
private static final Logger log = LoggerFactory.getLogger(HandleVisitor.class);
8385

8486
private boolean lookForStore;
85-
private int agentSpanVar = -1;
87+
private int contextVar = -1;
8688
private boolean success;
8789
private final String methodName;
8890

@@ -98,12 +100,12 @@ DelayCertainInsMethodVisitor delayVisitorDelegate() {
98100
@Override
99101
public void visitMethodInsn(
100102
int opcode, String owner, String name, String descriptor, boolean isInterface) {
101-
if (agentSpanVar == -1) {
103+
if (contextVar == -1) {
102104
lookForStore =
103105
!lookForStore
104106
&& opcode == Opcodes.INVOKEVIRTUAL
105107
&& name.equals("startSpan")
106-
&& descriptor.endsWith("Ldatadog/trace/bootstrap/instrumentation/api/AgentSpan;");
108+
&& descriptor.endsWith("Ldatadog.context.Context;");
107109
} else if (!success
108110
&& opcode == Opcodes.INVOKEVIRTUAL
109111
&& owner.equals("org/eclipse/jetty/server/Server")
@@ -128,9 +130,17 @@ public void visitMethodInsn(
128130
Label beforeHandle = new Label();
129131
Label afterHandle = new Label();
130132

131-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
133+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
134+
super.visitMethodInsn(
135+
Opcodes.INVOKESTATIC,
136+
"datadog/trace/bootstrap/instrumentation/api/Java8BytecodeBridge",
137+
"spanFromContext",
138+
"(" + Type.getDescriptor(Context.class) + ")" + Type.getDescriptor(AgentSpan.class),
139+
false);
140+
super.visitVarInsn(Opcodes.ASTORE, contextVar);
141+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
132142
super.visitJumpInsn(Opcodes.IFNULL, beforeHandle);
133-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
143+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
134144
super.visitMethodInsn(
135145
Opcodes.INVOKEINTERFACE,
136146
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
@@ -156,14 +166,14 @@ public void visitMethodInsn(
156166
"getResponse",
157167
"()Lorg/eclipse/jetty/server/Response;",
158168
false);
159-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
169+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
160170
super.visitMethodInsn(
161171
Opcodes.INVOKEINTERFACE,
162172
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
163173
"getRequestBlockingAction",
164174
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
165175
true);
166-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
176+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
167177
super.visitMethodInsn(
168178
Opcodes.INVOKESTATIC,
169179
Type.getInternalName(JettyBlockingHelper.class),
@@ -214,9 +224,9 @@ public void visitMethodInsn(
214224
Label beforeRegularDispatch = new Label();
215225
Label afterRegularDispatch = new Label();
216226

217-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
227+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
218228
super.visitJumpInsn(Opcodes.IFNULL, beforeRegularDispatch);
219-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
229+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
220230
super.visitMethodInsn(
221231
Opcodes.INVOKEINTERFACE,
222232
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
@@ -248,14 +258,14 @@ public void visitMethodInsn(
248258
"getResponse",
249259
"()Lorg/eclipse/jetty/server/Response;",
250260
false);
251-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
261+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
252262
super.visitMethodInsn(
253263
Opcodes.INVOKEINTERFACE,
254264
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
255265
"getRequestBlockingAction",
256266
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
257267
true);
258-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
268+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
259269

260270
// create the lambda
261271
super.visitInvokeDynamicInsn(
@@ -337,7 +347,7 @@ private boolean checkDispatchMethodState(final List<Function> savedVisitations)
337347
@Override
338348
public void visitVarInsn(int opcode, int varIndex) {
339349
if (lookForStore && opcode == Opcodes.ASTORE) {
340-
agentSpanVar = varIndex;
350+
contextVar = varIndex;
341351
lookForStore = false;
342352
}
343353

0 commit comments

Comments
 (0)