Skip to content

Commit b99070a

Browse files
zarirhamzajordan-wongPerfectSlayer
authored
Context API beforeFinish Migration (#9422)
Completes the context API migration for HTTP server spans, specifically within the beforeFinish lifecycle. --------- Co-authored-by: Zarir Hamza <zarir.hamza@datadoghq.com> Co-authored-by: Jordan Wong <61242306+jordan-wong@users.noreply.github.com> Co-authored-by: Bruce Bujon <bruce.bujon@datadoghq.com>
1 parent 83fdea5 commit b99070a

File tree

115 files changed

+1265
-810
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

115 files changed

+1265
-810
lines changed

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT;
55
import static datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver.hostName;
66

7+
import datadog.context.Context;
78
import datadog.context.ContextScope;
89
import datadog.trace.api.Config;
910
import datadog.trace.api.DDTags;
@@ -76,15 +77,19 @@ public AgentSpan afterStart(final AgentSpan span) {
7677
return span;
7778
}
7879

79-
public AgentScope beforeFinish(final AgentScope scope) {
80-
beforeFinish(scope.span());
80+
public ContextScope beforeFinish(final ContextScope scope) {
81+
beforeFinish(scope.context());
8182
return scope;
8283
}
8384

8485
public AgentSpan beforeFinish(final AgentSpan span) {
8586
return span;
8687
}
8788

89+
public Context beforeFinish(final Context context) {
90+
return context;
91+
}
92+
8893
public AgentScope onError(final AgentScope scope, final Throwable throwable) {
8994
onError(scope.span(), throwable);
9095
return scope;

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
5656
private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class);
5757
private static final int UNSET_PORT = 0;
5858

59-
public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
59+
public static final String DD_CONTEXT_ATTRIBUTE = "datadog.context";
6060
public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch";
6161
public static final String DD_RUM_INJECTED = "datadog.rum.injected";
6262
public static final String DD_FIN_DISP_LIST_SPAN_ATTRIBUTE =
@@ -537,12 +537,16 @@ private Flow<Void> callIGCallbackURI(
537537
}
538538

539539
@Override
540-
public AgentSpan beforeFinish(AgentSpan span) {
541-
// TODO Migrate beforeFinish to Context API
542-
onRequestEndForInstrumentationGateway(span);
540+
public Context beforeFinish(Context context) {
541+
AgentSpan span = AgentSpan.fromContext(context);
542+
if (span != null) {
543+
onRequestEndForInstrumentationGateway(span);
544+
}
545+
543546
// Close Serverless Gateway Inferred Span if any
544547
// finishInferredProxySpan(context);
545-
return super.beforeFinish(span);
548+
549+
return super.beforeFinish(context);
546550
}
547551

548552
protected void finishInferredProxySpan(Context context) {

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import akka.http.scaladsl.model.HttpResponse;
77
import akka.http.scaladsl.util.FastFuture$;
88
import akka.stream.Materializer;
9+
import datadog.context.Context;
910
import datadog.context.ContextScope;
1011
import datadog.trace.api.gateway.Flow;
1112
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -29,7 +30,8 @@ public DatadogAsyncHandlerWrapper(
2930
@Override
3031
public Future<HttpResponse> apply(final HttpRequest request) {
3132
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
32-
final AgentSpan span = fromContext(scope.context());
33+
final Context context = scope.context();
34+
final AgentSpan span = fromContext(context);
3335
Future<HttpResponse> futureResponse;
3436

3537
// handle blocking in the beginning of the request
@@ -38,15 +40,15 @@ public Future<HttpResponse> apply(final HttpRequest request) {
3840
request.discardEntityBytes(materializer);
3941
HttpResponse response = BlockingResponseHelper.maybeCreateBlockingResponse(rba, request);
4042
span.getRequestContext().getTraceSegment().effectivelyBlocked();
41-
DatadogWrapperHelper.finishSpan(span, response);
43+
DatadogWrapperHelper.finishSpan(context, response);
4244
return FastFuture$.MODULE$.<HttpResponse>successful().apply(response);
4345
}
4446

4547
try {
4648
futureResponse = userHandler.apply(request);
4749
} catch (final Throwable t) {
4850
scope.close();
49-
DatadogWrapperHelper.finishSpan(span, t);
51+
DatadogWrapperHelper.finishSpan(context, t);
5052
throw t;
5153
}
5254

@@ -67,14 +69,14 @@ public HttpResponse apply(HttpResponse response) {
6769
response = newResponse;
6870
}
6971

70-
DatadogWrapperHelper.finishSpan(span, response);
72+
DatadogWrapperHelper.finishSpan(context, response);
7173
return response;
7274
}
7375
},
7476
new AbstractFunction1<Throwable, Throwable>() {
7577
@Override
7678
public Throwable apply(final Throwable t) {
77-
DatadogWrapperHelper.finishSpan(span, t);
79+
DatadogWrapperHelper.finishSpan(context, t);
7880
return t;
7981
}
8082
},

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public void onPush() throws Exception {
7878
skipNextPull[0] = true;
7979
requestContext.getTraceSegment().effectivelyBlocked();
8080
emit(responseOutlet, response);
81-
DatadogWrapperHelper.finishSpan(span, response);
81+
DatadogWrapperHelper.finishSpan(scope.context(), response);
8282
pull(requestInlet);
8383
scope.close();
8484
return;
@@ -142,7 +142,7 @@ public void onPush() throws Exception {
142142
response.discardEntityBytes(materializer());
143143
response = newResponse;
144144
}
145-
DatadogWrapperHelper.finishSpan(span, response);
145+
DatadogWrapperHelper.finishSpan(scope.context(), response);
146146
// Check if the active span matches the scope from when the request came in,
147147
// and close it. If it's not, then it will be cleaned up actor message
148148
// processing instrumentation that drives this state machine
@@ -172,7 +172,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception {
172172
if (scope != null) {
173173
// Mark the span as failed
174174
AgentSpan span = fromContext(scope.context());
175-
DatadogWrapperHelper.finishSpan(span, ex);
175+
DatadogWrapperHelper.finishSpan(scope.context(), ex);
176176
}
177177
// We will not receive any more responses from the user code, so clean up any
178178
// remaining spans

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) {
2020
return context.attach();
2121
}
2222

23-
public static void finishSpan(final AgentSpan span, final HttpResponse response) {
23+
public static void finishSpan(final Context context, final HttpResponse response) {
24+
final AgentSpan span = fromContext(context);
2425
DECORATE.onResponse(span, response);
25-
DECORATE.beforeFinish(span);
26+
DECORATE.beforeFinish(context);
2627

2728
span.finish();
2829
}
2930

30-
public static void finishSpan(final AgentSpan span, final Throwable t) {
31+
public static void finishSpan(final Context context, final Throwable t) {
32+
final AgentSpan span = fromContext(context);
3133
DECORATE.onError(span, t);
3234
span.setHttpStatusCode(500);
33-
DECORATE.beforeFinish(span);
35+
DECORATE.beforeFinish(context);
3436

3537
span.finish();
3638
}

‎dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java‎

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

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
5+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
56
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
67
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE;
78
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.SERVER_TRANSACTION_CLASS;
89

10+
import datadog.context.Context;
11+
import datadog.context.ContextScope;
912
import datadog.trace.bootstrap.InstrumentationContext;
10-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1113
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1214
import net.bytebuddy.asm.Advice;
1315

1416
public class HTTPPluginAdvice {
1517

1618
@Advice.OnMethodEnter(suppress = Throwable.class)
17-
public static AgentScope onEnter(@Advice.Argument(value = 2) final Object serverTransaction) {
19+
public static ContextScope onEnter(@Advice.Argument(value = 2) final Object serverTransaction) {
1820
final AgentSpan span = startSpan("axway-api", DECORATE.spanName()).setMeasured(true);
1921
DECORATE.afterStart(span);
2022
// serverTransaction is like request + connection in one object:
2123
DECORATE.onRequest(span, serverTransaction, serverTransaction, getRootContext());
22-
return activateSpan(span);
24+
return getCurrentContext().with(span).attach();
2325
}
2426

2527
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
2628
public static void onExit(
27-
@Advice.Enter final AgentScope scope,
29+
@Advice.Enter final ContextScope scope,
2830
@Advice.Argument(value = 2) final Object serverTransaction,
2931
@Advice.Thrown final Throwable throwable) {
3032
if (scope == null) {
3133
return;
3234
}
33-
final AgentSpan span = scope.span();
35+
final Context context = scope.context();
36+
final AgentSpan span = fromContext(context);
3437
try {
3538
if (null != serverTransaction) {
3639
// manual DECORATE.onResponse(span, serverTransaction):
@@ -45,7 +48,7 @@ public static void onExit(
4548
if (throwable != null) {
4649
DECORATE.onError(span, throwable);
4750
}
48-
DECORATE.beforeFinish(span);
51+
DECORATE.beforeFinish(context);
4952
} finally {
5053
scope.close();
5154
span.finish();

‎dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.AXWAY_TRY_TRANSACTION;
66
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE;
77

8+
import datadog.context.Context;
89
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
910
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1011
import net.bytebuddy.asm.Advice;
@@ -34,11 +35,12 @@ public static void onExit(
3435
return;
3536
}
3637
final AgentSpan span = scope.span();
38+
final Context context = scope.context();
3739
try {
3840
if (throwable != null) {
3941
DECORATE.onError(span, throwable);
4042
}
41-
DECORATE.beforeFinish(span);
43+
DECORATE.beforeFinish(context);
4244
} finally {
4345
scope.close();
4446
span.finish();

‎dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ public static void methodExit(
8585
final AgentSpan span = fromContext(scope.context());
8686
DECORATE.onError(span, throwable);
8787
DECORATE.onResponse(span, response);
88-
DECORATE.beforeFinish(span);
88+
DECORATE.beforeFinish(scope.context());
8989
scope.close();
9090
span.finish();
9191
}

‎dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
package datadog.trace.instrumentation.cxf;
22

33
import com.google.auto.service.AutoService;
4+
import datadog.context.Context;
5+
import datadog.context.ContextScope;
46
import datadog.trace.agent.tooling.Instrumenter;
57
import datadog.trace.agent.tooling.InstrumenterModule;
68
import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
79
import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
810
import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers;
9-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
10-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1111
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
1212
import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
1313
import net.bytebuddy.asm.Advice;
@@ -55,21 +55,22 @@ public void methodAdvice(MethodTransformer transformer) {
5555

5656
public static class PropagateSpanAdvice {
5757
@Advice.OnMethodEnter(suppress = Throwable.class)
58-
public static AgentScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) {
58+
public static ContextScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) {
5959
if (exchange == null || exchange.getInMessage() == null || AgentTracer.activeSpan() != null) {
6060
return null;
6161
}
62-
final Object span =
62+
final Object contextObj =
6363
ServletHelper.getServletRequestAttribute(
64-
exchange.getInMessage().get("HTTP.REQUEST"), HttpServerDecorator.DD_SPAN_ATTRIBUTE);
65-
if (span instanceof AgentSpan) {
66-
return AgentTracer.activateSpan((AgentSpan) span);
64+
exchange.getInMessage().get("HTTP.REQUEST"),
65+
HttpServerDecorator.DD_CONTEXT_ATTRIBUTE);
66+
if (contextObj instanceof Context) {
67+
return ((Context) contextObj).attach();
6768
}
6869
return null;
6970
}
7071

7172
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
72-
public static void afterInvoke(@Advice.Enter final AgentScope scope) {
73+
public static void afterInvoke(@Advice.Enter final ContextScope scope) {
7374
if (scope != null) {
7475
scope.close();
7576
}

‎dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,10 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
55
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
6-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
76
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
87
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
8+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
9+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
910
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
1011
import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE;
1112
import static datadog.trace.instrumentation.finatra.FinatraDecorator.FINATRA_CONTROLLER;
@@ -17,9 +18,9 @@
1718
import com.twitter.finagle.http.Request;
1819
import com.twitter.finagle.http.Response;
1920
import com.twitter.util.Future;
21+
import datadog.context.ContextScope;
2022
import datadog.trace.agent.tooling.Instrumenter;
2123
import datadog.trace.agent.tooling.InstrumenterModule;
22-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
2324
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
2425
import datadog.trace.bootstrap.instrumentation.api.Tags;
2526
import net.bytebuddy.asm.Advice;
@@ -61,7 +62,7 @@ public void methodAdvice(MethodTransformer transformer) {
6162

6263
public static class RouteAdvice {
6364
@Advice.OnMethodEnter(suppress = Throwable.class)
64-
public static AgentScope nameSpan(
65+
public static ContextScope nameSpan(
6566
@Advice.Argument(0) final Request request,
6667
@Advice.FieldValue("path") final String path,
6768
@Advice.FieldValue("clazz") final Class clazz) {
@@ -78,23 +79,23 @@ public static AgentScope nameSpan(
7879
DECORATE.afterStart(span);
7980
span.setResourceName(DECORATE.className(clazz));
8081

81-
return activateSpan(span);
82+
return getCurrentContext().with(span).attach();
8283
}
8384

8485
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
8586
public static void setupCallback(
86-
@Advice.Enter final AgentScope scope,
87+
@Advice.Enter final ContextScope scope,
8788
@Advice.Thrown final Throwable throwable,
8889
@Advice.Return final Some<Future<Response>> responseOption) {
8990

9091
if (scope == null) {
9192
return;
9293
}
9394

94-
final AgentSpan span = scope.span();
95+
final AgentSpan span = spanFromContext(scope.context());
9596
if (throwable != null) {
9697
DECORATE.onError(span, throwable);
97-
DECORATE.beforeFinish(span);
98+
DECORATE.beforeFinish(scope.context());
9899
span.finish();
99100
scope.close();
100101
return;

0 commit comments

Comments
 (0)