Skip to content

Commit 73c892e

Browse files
PerfectSlayerPerfectSlayer
committed
feat(context): Improve Context API usage
1 parent 2858dfb commit 73c892e

22 files changed

Lines changed: 95 additions & 89 deletions

File tree

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import java.util.function.Function;
4646
import java.util.function.Supplier;
4747
import javax.annotation.Nonnull;
48+
import javax.annotation.Nullable;
4849
import org.slf4j.Logger;
4950
import org.slf4j.LoggerFactory;
5051

@@ -181,7 +182,7 @@ public AgentSpan onRequest(
181182
final AgentSpan span,
182183
final CONNECTION connection,
183184
final REQUEST request,
184-
final AgentSpanContext.Extracted context) {
185+
@Nullable final AgentSpanContext.Extracted context) {
185186
Config config = Config.get();
186187
boolean clientIpResolverEnabled =
187188
config.isClientIpEnabled()

‎dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
package datadog.trace.instrumentation.akkahttp;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
4+
35
import akka.http.scaladsl.model.HttpRequest;
46
import akka.http.scaladsl.model.HttpResponse;
57
import akka.http.scaladsl.util.FastFuture$;
68
import akka.stream.Materializer;
7-
import datadog.context.Context;
89
import datadog.context.ContextScope;
910
import datadog.trace.api.gateway.Flow;
1011
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -28,8 +29,7 @@ public DatadogAsyncHandlerWrapper(
2829
@Override
2930
public Future<HttpResponse> apply(final HttpRequest request) {
3031
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
31-
Context context = scope.context();
32-
final AgentSpan span = AgentSpan.fromContext(context);
32+
final AgentSpan span = fromContext(scope.context());
3333
Future<HttpResponse> futureResponse;
3434

3535
// handle blocking in the beginning of the request

‎dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java‎

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

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
45

56
import akka.http.scaladsl.model.HttpRequest;
@@ -67,7 +68,7 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
6768
public void onPush() throws Exception {
6869
final HttpRequest request = grab(requestInlet);
6970
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
70-
final AgentSpan span = AgentSpan.fromContext(scope.context());
71+
final AgentSpan span = fromContext(scope.context());
7172
RequestContext requestContext = span.getRequestContext();
7273
if (requestContext != null) {
7374
HttpResponse response =
@@ -133,7 +134,7 @@ public void onPush() throws Exception {
133134
HttpResponse response = grab(responseInlet);
134135
final ContextScope scope = scopes.poll();
135136
if (scope != null) {
136-
AgentSpan span = AgentSpan.fromContext(scope.context());
137+
AgentSpan span = fromContext(scope.context());
137138
HttpResponse newResponse =
138139
BlockingResponseHelper.handleFinishForWaf(span, response);
139140
if (newResponse != response) {
@@ -159,7 +160,7 @@ public void onUpstreamFinish() throws Exception {
159160
// remaining spans
160161
ContextScope scope = scopes.poll();
161162
while (scope != null) {
162-
AgentSpan.fromContext(scope.context()).finish();
163+
fromContext(scope.context()).finish();
163164
scope = scopes.poll();
164165
}
165166
completeStage();
@@ -168,7 +169,7 @@ public void onUpstreamFinish() throws Exception {
168169
@Override
169170
public void onUpstreamFailure(final Throwable ex) throws Exception {
170171
ContextScope scope = scopes.poll();
171-
AgentSpan span = AgentSpan.fromContext(scope.context());
172+
AgentSpan span = fromContext(scope.context());
172173
if (scope != null) {
173174
// Mark the span as failed
174175
DatadogWrapperHelper.finishSpan(span, ex);

‎dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@
1010

1111
public class DatadogWrapperHelper {
1212
public static ContextScope createSpan(final HttpRequest request) {
13-
final Context extractedContext = DECORATE.extractContext(request);
14-
final AgentSpan span = DECORATE.startSpan(request, extractedContext);
13+
final Context context = DECORATE.extractContext(request);
14+
final AgentSpan span = DECORATE.startSpan(request, context);
1515
DECORATE.afterStart(span);
16-
DECORATE.onRequest(span, request, request, extractedContext);
16+
DECORATE.onRequest(span, request, request, context);
1717

18-
return extractedContext.with(span).attach();
18+
return context.with(span).attach();
1919
}
2020

2121
public static void finishSpan(final AgentSpan span, final HttpResponse response) {

‎dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,12 @@ public static ContextScope onEnter(
2626
return activateSpan((AgentSpan) existingSpan);
2727
}
2828

29-
final Context extractedContext = DECORATE.extractContext(req);
30-
span = DECORATE.startSpan(req, extractedContext);
31-
final ContextScope scope = extractedContext.with(span).attach();
29+
final Context context = DECORATE.extractContext(req);
30+
span = DECORATE.startSpan(req, context);
31+
final ContextScope scope = context.with(span).attach();
3232
span.setMeasured(true);
3333
DECORATE.afterStart(span);
34-
DECORATE.onRequest(span, req, req, extractedContext);
34+
DECORATE.onRequest(span, req, req, context);
3535

3636
req.setAttribute(DD_SPAN_ATTRIBUTE, span);
3737
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());

‎dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,12 @@ public static void onExit(
2727
}
2828
}
2929

30-
final Context extractedContext = JettyDecorator.DECORATE.extractContext(req);
31-
final AgentSpan span = JettyDecorator.DECORATE.startSpan(req, extractedContext);
32-
try (final ContextScope scope = extractedContext.with(span).attach()) {
30+
final Context context = JettyDecorator.DECORATE.extractContext(req);
31+
final AgentSpan span = JettyDecorator.DECORATE.startSpan(req, context);
32+
try (final ContextScope scope = context.with(span).attach()) {
3333
span.setMeasured(true);
3434
JettyDecorator.DECORATE.afterStart(span);
35-
JettyDecorator.DECORATE.onRequest(span, req, req, extractedContext);
35+
JettyDecorator.DECORATE.onRequest(span, req, req, context);
3636

3737
req.setAttribute(DD_SPAN_ATTRIBUTE, span);
3838
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());

‎dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,11 @@ public static ContextScope onEnter(
170170
return ((AgentSpan) existingSpan).attach();
171171
}
172172

173-
final Context extractedContext = DECORATE.extractContext(req);
174-
span = DECORATE.startSpan(req, extractedContext);
175-
final ContextScope scope = extractedContext.with(span).attach();
173+
final Context context = DECORATE.extractContext(req);
174+
span = DECORATE.startSpan(req, context);
175+
final ContextScope scope = context.with(span).attach();
176176
DECORATE.afterStart(span);
177-
DECORATE.onRequest(span, req, req, extractedContext);
177+
DECORATE.onRequest(span, req, req, context);
178178

179179
req.setAttribute(DD_SPAN_ATTRIBUTE, span);
180180
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,12 @@ public static ContextScope onEnter(
2424
return ((AgentSpan) existingSpan).attach();
2525
}
2626

27-
final Context extractedContext = DECORATE.extractContext(req);
28-
span = DECORATE.startSpan(req, extractedContext);
27+
final Context context = DECORATE.extractContext(req);
28+
span = DECORATE.startSpan(req, context);
2929
DECORATE.afterStart(span);
30-
DECORATE.onRequest(span, req, req, extractedContext);
30+
DECORATE.onRequest(span, req, req, context);
3131

32-
final ContextScope scope = extractedContext.with(span).attach();
32+
final ContextScope scope = context.with(span).attach();
3333
req.setAttribute(DD_SPAN_ATTRIBUTE, span);
3434
req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
3535
req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());

‎dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -105,10 +105,10 @@ public static class HandleRequestAdvice {
105105
} catch (NullPointerException e) {
106106
}
107107

108-
final Context extractedContext = DECORATE.extractContext(request);
109-
request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
110-
final AgentSpan span = DECORATE.startSpan(request, extractedContext);
111-
scope = extractedContext.with(span).attach();
108+
final Context context = DECORATE.extractContext(request);
109+
request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, context);
110+
final AgentSpan span = DECORATE.startSpan(request, context);
111+
scope = context.with(span).attach();
112112
if (Config.get().isJeeSplitByDeployment()) {
113113
final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext();
114114
if (dispatcherContext != null) {
@@ -122,7 +122,7 @@ public static class HandleRequestAdvice {
122122
}
123123
}
124124
DECORATE.afterStart(span);
125-
DECORATE.onRequest(span, request, request, extractedContext);
125+
DECORATE.onRequest(span, request, request, context);
126126
request.setAttribute(DD_SPAN_ATTRIBUTE, span);
127127
request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
128128
request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());

‎dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,10 @@ public static class HandleRequestAdvice {
107107
} catch (NullPointerException e) {
108108
}
109109

110-
final Context extractedContext = DECORATE.extractContext(request);
111-
request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
112-
final AgentSpan span = DECORATE.startSpan(request, extractedContext);
113-
scope = extractedContext.with(span).attach();
110+
final Context context = DECORATE.extractContext(request);
111+
request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, context);
112+
final AgentSpan span = DECORATE.startSpan(request, context);
113+
scope = context.with(span).attach();
114114
if (Config.get().isJeeSplitByDeployment()) {
115115
final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext();
116116
if (dispatcherContext != null) {
@@ -124,7 +124,7 @@ public static class HandleRequestAdvice {
124124
}
125125
}
126126
DECORATE.afterStart(span);
127-
DECORATE.onRequest(span, request, request, extractedContext);
127+
DECORATE.onRequest(span, request, request, context);
128128
request.setAttribute(DD_SPAN_ATTRIBUTE, span);
129129
request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
130130
request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());

0 commit comments

Comments
 (0)