Skip to content

Commit 841b27c

Browse files
zarirhamzazarirhamza
committed
address all comments
1 parent c62d570 commit 841b27c

File tree

18 files changed

+129
-75
lines changed

18 files changed

+129
-75
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -539,10 +539,12 @@ private Flow<Void> callIGCallbackURI(
539539
@Override
540540
public Context beforeFinish(Context context) {
541541
AgentSpan span = AgentSpan.fromContext(context);
542-
onRequestEndForInstrumentationGateway(span);
542+
if (span != null) {
543+
onRequestEndForInstrumentationGateway(span);
544+
}
543545

544546
// Close Serverless Gateway Inferred Span if any
545-
finishInferredProxySpan(context);
547+
// finishInferredProxySpan(context);
546548

547549
return super.beforeFinish(context);
548550
}

‎dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ public void onAfterService(final Request request) {
2424
DECORATE.onResponse(span, request.getResponse());
2525
DECORATE.beforeFinish(context);
2626
span.finish();
27+
} else {
28+
DECORATE.beforeFinish(context);
2729
}
2830
}
2931
}

‎dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,8 @@ public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable
154154
DECORATE.onError(span, throwable);
155155
DECORATE.beforeFinish(context);
156156
span.finish();
157+
} else if (context != null) {
158+
DECORATE.beforeFinish(context);
157159
}
158160
ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE);
159161
ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE);

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -51,14 +51,26 @@ public class JettyCommitResponseHelper {
5151

5252
Request req = connection.getRequest();
5353

54-
Object contextObj;
55-
Context context;
56-
AgentSpan span;
57-
RequestContext requestContext;
58-
if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null
59-
|| !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context)
60-
|| (span = spanFromContext(context = (Context) contextObj)) == null
61-
|| (requestContext = span.getRequestContext()) == null) {
54+
if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) {
55+
state.partialResponse();
56+
return false;
57+
}
58+
59+
Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE);
60+
if (!(contextObj instanceof Context)) {
61+
state.partialResponse();
62+
return false;
63+
}
64+
65+
Context context = (Context) contextObj;
66+
AgentSpan span = spanFromContext(context);
67+
if (span == null) {
68+
state.partialResponse();
69+
return false;
70+
}
71+
72+
RequestContext requestContext = span.getRequestContext();
73+
if (requestContext == null) {
6274
state.partialResponse();
6375
return false;
6476
}

‎dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ public static void stopSpan(@Advice.This final HttpChannel channel) {
2626
JettyDecorator.OnResponse.onResponse(span, channel);
2727
DECORATE.beforeFinish(context);
2828
span.finish();
29+
} else {
30+
DECORATE.beforeFinish(context);
2931
}
3032
}
3133
}

‎dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
7878
NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE;
7979

8080
final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST);
81-
final Context spanContext = getCurrentContext().with(span);
81+
final Context context = getCurrentContext().with(span);
8282
try (final AgentScope scope = activateSpan(span)) {
8383
decorate.afterStart(span);
8484
decorate.onRequest(span, request);
@@ -93,7 +93,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
9393
DECORATE.injectContext(current(), request.headers(), SETTER);
9494
}
9595

96-
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext);
96+
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(context);
9797

9898
try {
9999
ctx.write(msg, prm);

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
7979
NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE;
8080

8181
final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST);
82-
final Context spanContext = getCurrentContext().with(span);
82+
final Context context = getCurrentContext().with(span);
8383
try (final AgentScope scope = activateSpan(span)) {
8484
decorate.afterStart(span);
8585
decorate.onRequest(span, request);
@@ -94,7 +94,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
9494
DECORATE.injectContext(current(), request.headers(), SETTER);
9595
}
9696

97-
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext);
97+
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(context);
9898

9999
try {
100100
ctx.write(msg, prm);

‎dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
5050
}
5151
} else {
5252
if (storedContext != null) {
53-
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(span));
53+
ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext);
5454
}
5555
}
5656
}

‎dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientDecorator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public final class SynapseClientDecorator extends HttpClientDecorator<HttpReques
1414
public static final CharSequence SYNAPSE_CLIENT = UTF8BytesString.create("synapse-client");
1515
public static final CharSequence SYNAPSE_REQUEST =
1616
UTF8BytesString.create(DECORATE.operationName());
17-
public static final String SYNAPSE_SPAN_KEY = "dd.trace.synapse.span";
17+
public static final String SYNAPSE_CONTEXT_KEY = "dd.trace.synapse.context";
1818
public static final String SYNAPSE_CONTINUATION_KEY = "dd.trace.synapse.continuation";
1919

2020
@Override

‎dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java‎

Lines changed: 33 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,21 @@
33
import static datadog.context.propagation.Propagators.defaultPropagator;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
6-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
76
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
87
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
8+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
99
import static datadog.trace.instrumentation.synapse3.SynapseClientDecorator.DECORATE;
10+
import static datadog.trace.instrumentation.synapse3.SynapseClientDecorator.SYNAPSE_CONTEXT_KEY;
1011
import static datadog.trace.instrumentation.synapse3.SynapseClientDecorator.SYNAPSE_REQUEST;
11-
import static datadog.trace.instrumentation.synapse3.SynapseClientDecorator.SYNAPSE_SPAN_KEY;
1212
import static datadog.trace.instrumentation.synapse3.TargetRequestInjectAdapter.SETTER;
1313
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1414
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1515

1616
import com.google.auto.service.AutoService;
17+
import datadog.context.Context;
18+
import datadog.context.ContextScope;
1719
import datadog.trace.agent.tooling.Instrumenter;
1820
import datadog.trace.agent.tooling.InstrumenterModule;
19-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2021
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2122
import net.bytebuddy.asm.Advice;
2223
import org.apache.axis2.context.MessageContext;
@@ -64,16 +65,16 @@ public void methodAdvice(final MethodTransformer transformer) {
6465

6566
public static final class ClientRequestAdvice {
6667
@Advice.OnMethodEnter(suppress = Throwable.class)
67-
public static AgentScope beginRequest(
68+
public static ContextScope beginRequest(
6869
@Advice.Argument(0) final NHttpClientConnection connection) {
6970

7071
// check for parent span propagated by SynapsePassthruInstrumentation
7172
AgentSpan parentSpan = null;
7273
MessageContext message = TargetContext.get(connection).getRequestMsgCtx();
7374
if (null != message) {
74-
parentSpan = (AgentSpan) message.getPropertyNonReplicable(SYNAPSE_SPAN_KEY);
75+
parentSpan = (AgentSpan) message.getPropertyNonReplicable(SYNAPSE_CONTEXT_KEY);
7576
if (null != parentSpan) {
76-
message.removePropertyNonReplicable(SYNAPSE_SPAN_KEY);
77+
message.removePropertyNonReplicable(SYNAPSE_CONTEXT_KEY);
7778
}
7879
}
7980

@@ -86,54 +87,58 @@ public static AgentScope beginRequest(
8687

8788
DECORATE.afterStart(span);
8889

90+
Context context = getCurrentContext().with(span);
91+
8992
// add trace id to client-side request before it gets submitted as an HttpRequest
90-
defaultPropagator()
91-
.inject(getCurrentContext().with(span), TargetContext.getRequest(connection), SETTER);
93+
defaultPropagator().inject(context, TargetContext.getRequest(connection), SETTER);
9294

93-
// capture span to be finished by one of the various client response advices
94-
connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span);
95+
// capture context to be finished by one of the various client response advices
96+
connection.getContext().setAttribute(SYNAPSE_CONTEXT_KEY, context);
9597

96-
return activateSpan(span);
98+
return context.attach();
9799
}
98100

99101
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
100102
public static void requestSubmitted(
101103
@Advice.Argument(0) final NHttpClientConnection connection,
102-
@Advice.Enter final AgentScope scope) {
104+
@Advice.Enter final ContextScope scope) {
103105
// populate span using details from the submitted HttpRequest (resolved URI, etc.)
104-
DECORATE.onRequest(scope.span(), TargetContext.getRequest(connection).getRequest());
106+
AgentSpan span = spanFromContext(scope.context());
107+
DECORATE.onRequest(span, TargetContext.getRequest(connection).getRequest());
105108
scope.close();
106109
}
107110
}
108111

109112
public static final class ClientResponseAdvice {
110113
@Advice.OnMethodEnter(suppress = Throwable.class)
111-
public static AgentScope beginResponse(
114+
public static ContextScope beginResponse(
112115
@Advice.Argument(0) final NHttpClientConnection connection) {
113-
// check and remove span from context so it won't be finished twice
114-
AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY);
115-
if (null != span) {
116-
return activateSpan(span);
116+
// check and remove context so it won't be finished twice
117+
Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY);
118+
if (null != context) {
119+
return context.attach();
117120
}
118121
return null;
119122
}
120123

121124
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
122125
public static void responseReceived(
123126
@Advice.Argument(0) final NHttpClientConnection connection,
124-
@Advice.Enter final AgentScope scope,
127+
@Advice.Enter final ContextScope scope,
125128
@Advice.Thrown final Throwable error) {
126129
if (null == scope) {
127130
return;
128131
}
129-
AgentSpan span = scope.span();
132+
AgentSpan span = spanFromContext(scope.context());
130133
DECORATE.onResponse(span, connection.getHttpResponse());
131134
if (null != error) {
132135
DECORATE.onError(span, error);
133136
}
134-
DECORATE.beforeFinish(span);
137+
DECORATE.beforeFinish(scope.context());
135138
scope.close();
136-
span.finish();
139+
if (span != null) {
140+
span.finish();
141+
}
137142
}
138143
}
139144

@@ -142,16 +147,19 @@ public static final class ClientErrorResponseAdvice {
142147
public static void errorResponse(
143148
@Advice.Argument(0) final NHttpClientConnection connection,
144149
@Advice.Argument(value = 1, optional = true) final Object error) {
145-
// check and remove span from context so it won't be finished twice
146-
AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY);
150+
// check and remove context so it won't be finished twice
151+
Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_CONTEXT_KEY);
152+
AgentSpan span = spanFromContext(context);
147153
if (null != span) {
148154
if (error instanceof Throwable) {
149155
DECORATE.onError(span, (Throwable) error);
150156
} else {
151157
span.setError(true);
152158
}
153-
DECORATE.beforeFinish(span);
159+
DECORATE.beforeFinish(context);
154160
span.finish();
161+
} else {
162+
DECORATE.beforeFinish(context);
155163
}
156164
}
157165
}

0 commit comments

Comments
 (0)