Skip to content

Context API beforeFinish Migration #9422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 36 commits into
base: master
Choose a base branch
from
+1,166 −779
Open

Context API beforeFinish Migration #9422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
36 commits
Select commit Hold shift + click to select a range
796b57e
Implement APIGW Inferred Proxy Spans (#8336)
jordan-wong May 7, 2025
e8c2baf
feat(serverless): Fix gateway inferred span design
PerfectSlayer May 12, 2025
737e783
init commit of instrumentations
zarirhamza Aug 26, 2025
c9a2e21
fix tests
zarirhamza Aug 26, 2025
a0ecac7
round 2 of context migration
zarirhamza Aug 26, 2025
2ef3dc7
netty changes
zarirhamza Aug 26, 2025
a14848f
tomcat changes
zarirhamza Aug 26, 2025
112ec20
liberty changes
zarirhamza Aug 26, 2025
909d174
store span
zarirhamza Aug 26, 2025
c883e8d
jetty
zarirhamza Aug 26, 2025
b847401
grizzly
zarirhamza Aug 26, 2025
bebe637
Merge branch 'master' into zarir/context-api-finish
zarirhamza Aug 26, 2025
f76415f
fix grizzly instrumentation
zarirhamza Aug 26, 2025
8f4a6ee
Merge branch 'master' into zarir/context-api-finish
zarirhamza Aug 27, 2025
762c233
remove serverless
zarirhamza Aug 27, 2025
c8c9ece
fixes part 1
zarirhamza Aug 27, 2025
f715d62
fix undertow tests
zarirhamza Aug 27, 2025
7fe9c70
remove span within request
zarirhamza Aug 27, 2025
0f0f228
fix tests
zarirhamza Aug 27, 2025
ecc5049
fix tests
zarirhamza Aug 27, 2025
4bd005d
implement fixes part 1
zarirhamza Aug 28, 2025
9b4d887
fix tests
zarirhamza Aug 28, 2025
a3bb790
merge main fixing merge conflicts
zarirhamza Oct 2, 2025
ee71b23
address jetty comments
zarirhamza Oct 4, 2025
dd4c4c0
Merge branch 'master' into zarir/context-api-finish
zarirhamza Oct 4, 2025
7fd52d8
Merge branch 'master' into zarir/context-api-finish
zarirhamza Oct 8, 2025
2c27431
use context scope
zarirhamza Oct 25, 2025
6e950a0
fix netty tests
zarirhamza Oct 25, 2025
5f20aee
Merge branch 'master' into zarir/context-api-finish
zarirhamza Oct 25, 2025
fd11e7a
fix baggage context usage
zarirhamza Oct 30, 2025
0c4ab54
Merge branch 'master' into zarir/context-api-finish
zarirhamza Oct 30, 2025
c13e8cd
Merge branch 'master' into zarir/context-api-finish
zarirhamza Nov 4, 2025
b9ad08c
fix context propagation
zarirhamza Nov 4, 2025
ce77411
Merge branch 'master' into zarir/context-api-finish
zarirhamza Nov 4, 2025
b33528c
fix netty timeout tests
zarirhamza Nov 5, 2025
c62d570
Merge branch 'master' into zarir/context-api-finish
zarirhamza Nov 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions ...tstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import static datadog.trace.api.cache.RadixTreeCache.UNSET_PORT;
import static datadog.trace.bootstrap.instrumentation.java.net.HostNameResolver.hostName;

import datadog.context.Context;
import datadog.context.ContextScope;
import datadog.trace.api.Config;
import datadog.trace.api.DDTags;
Expand Down Expand Up @@ -76,15 +77,19 @@ public AgentSpan afterStart(final AgentSpan span) {
return span;
}

public AgentScope beforeFinish(final AgentScope scope) {
beforeFinish(scope.span());
public ContextScope beforeFinish(final ContextScope scope) {
beforeFinish(scope.context());
return scope;
}

public AgentSpan beforeFinish(final AgentSpan span) {
return span;
}
Comment on lines 85 to 87
Copy link
Contributor

@PerfectSlayer PerfectSlayer Aug 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❔ question: ‏Can this one be removed then? It will prevent other decorator to override it and never being called.

Copy link
Contributor Author

@zarirhamza zarirhamza Aug 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cannot, WebSocketDecorator and its related instrumentation still call it. Needs to be a separate mini-migration

Copy link
Contributor

@PerfectSlayer PerfectSlayer Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is that the only one blocker? Because the WebSocketDecorator is calling the empty BaseDecorator.beforeFinish() and has no override... So technically, the call is doing nothing :/
So we can move it from BaseDecorator to WebSocketDecorator and migrate it in a second pass.

@amarziali What's the effort to migrate the WebSocket instrumentation to context rather than span?

Copy link
Contributor Author

@zarirhamza zarirhamza Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another place where this is called is within the Mule Instrumentation as well via the MuleDecorator in mule-4.

Moved the empty call to WebSocketDecorator and MuleDecorator thus allowing me to remove the function from BaseDecorator.

I assume the actual span -> context functionality will be taken care of later for both WebSockets and Mule.

Copy link
Contributor Author

@zarirhamza zarirhamza Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I take it back, I'm not sure how the tests aren't catching it but the ClientDecorator also uses beforeFinish(span) via the BaseDecorator. This seems to indicate that we need to invest time into looking through all decorators and confirming their reliance on beforeFinish. To save time, for now we can focus just on HttpServerDecorator related instrumentation

Copy link
Contributor

@ygree ygree Nov 5, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: It seems that there are no meaningful overrides of beforeFinish(final AgentSpan span), except for AxisMessageDecorator, which takes an extra parameter. To avoid confusion, let's go ahead and remove beforeFinish(final AgentSpan span).


public Context beforeFinish(final Context context) {
return context;
}

public AgentScope onError(final AgentScope scope, final Throwable throwable) {
onError(scope.span(), throwable);
return scope;
Expand Down
12 changes: 7 additions & 5 deletions .../src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
private static final Logger log = LoggerFactory.getLogger(HttpServerDecorator.class);
private static final int UNSET_PORT = 0;

public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
public static final String DD_CONTEXT_ATTRIBUTE = "datadog.context";
public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch";
public static final String DD_RUM_INJECTED = "datadog.rum.injected";
public static final String DD_FIN_DISP_LIST_SPAN_ATTRIBUTE =
Expand Down Expand Up @@ -537,12 +537,14 @@ private Flow<Void> callIGCallbackURI(
}

@Override
public AgentSpan beforeFinish(AgentSpan span) {
// TODO Migrate beforeFinish to Context API
Copy link
Contributor

@ygree ygree Nov 5, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: For cross reference, it's worth mentioning #9388 in this PR description.

public Context beforeFinish(Context context) {
AgentSpan span = AgentSpan.fromContext(context);
onRequestEndForInstrumentationGateway(span);

// Close Serverless Gateway Inferred Span if any
// finishInferredProxySpan(context);
Copy link
Contributor

@ygree ygree Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@zarirhamza This PR is more than just a refactoring. It now calls finishInferredProxySpan. Is there a test to verify that it actually makes a call to finishInferredProxySpan?

return super.beforeFinish(span);
finishInferredProxySpan(context);

return super.beforeFinish(context);
}

protected void finishInferredProxySpan(Context context) {
Expand Down
12 changes: 7 additions & 5 deletions ...10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import akka.http.scaladsl.model.HttpResponse;
import akka.http.scaladsl.util.FastFuture$;
import akka.stream.Materializer;
import datadog.context.Context;
import datadog.context.ContextScope;
import datadog.trace.api.gateway.Flow;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
Expand All @@ -29,7 +30,8 @@ public DatadogAsyncHandlerWrapper(
@Override
public Future<HttpResponse> apply(final HttpRequest request) {
final ContextScope scope = DatadogWrapperHelper.createSpan(request);
final AgentSpan span = fromContext(scope.context());
final Context context = scope.context();
final AgentSpan span = fromContext(context);
Future<HttpResponse> futureResponse;

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

try {
futureResponse = userHandler.apply(request);
} catch (final Throwable t) {
scope.close();
DatadogWrapperHelper.finishSpan(span, t);
DatadogWrapperHelper.finishSpan(context, t);
throw t;
}

Expand All @@ -67,14 +69,14 @@ public HttpResponse apply(HttpResponse response) {
response = newResponse;
}

DatadogWrapperHelper.finishSpan(span, response);
DatadogWrapperHelper.finishSpan(context, response);
return response;
}
},
new AbstractFunction1<Throwable, Throwable>() {
@Override
public Throwable apply(final Throwable t) {
DatadogWrapperHelper.finishSpan(span, t);
DatadogWrapperHelper.finishSpan(context, t);
return t;
}
},
Expand Down
6 changes: 3 additions & 3 deletions .../java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ public void onPush() throws Exception {
skipNextPull[0] = true;
requestContext.getTraceSegment().effectivelyBlocked();
emit(responseOutlet, response);
DatadogWrapperHelper.finishSpan(span, response);
DatadogWrapperHelper.finishSpan(scope.context(), response);
pull(requestInlet);
scope.close();
return;
Expand Down Expand Up @@ -142,7 +142,7 @@ public void onPush() throws Exception {
response.discardEntityBytes(materializer());
response = newResponse;
}
DatadogWrapperHelper.finishSpan(span, response);
DatadogWrapperHelper.finishSpan(scope.context(), response);
// Check if the active span matches the scope from when the request came in,
// and close it. If it's not, then it will be cleaned up actor message
// processing instrumentation that drives this state machine
Expand Down Expand Up @@ -172,7 +172,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception {
if (scope != null) {
// Mark the span as failed
AgentSpan span = fromContext(scope.context());
DatadogWrapperHelper.finishSpan(span, ex);
DatadogWrapperHelper.finishSpan(scope.context(), ex);
}
// We will not receive any more responses from the user code, so clean up any
// remaining spans
Expand Down
10 changes: 6 additions & 4 deletions ...-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,19 @@ public static ContextScope createSpan(final HttpRequest request) {
return context.attach();
}

public static void finishSpan(final AgentSpan span, final HttpResponse response) {
public static void finishSpan(final Context context, final HttpResponse response) {
final AgentSpan span = fromContext(context);
DECORATE.onResponse(span, response);
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(context);

span.finish();
}

public static void finishSpan(final AgentSpan span, final Throwable t) {
public static void finishSpan(final Context context, final Throwable t) {
final AgentSpan span = fromContext(context);
DECORATE.onError(span, t);
span.setHttpStatusCode(500);
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(context);

span.finish();
}
Expand Down
17 changes: 10 additions & 7 deletions ...ntation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,36 +1,39 @@
package datadog.trace.instrumentation.axway;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext;
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE;
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.SERVER_TRANSACTION_CLASS;

import datadog.context.Context;
import datadog.context.ContextScope;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import net.bytebuddy.asm.Advice;

public class HTTPPluginAdvice {

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

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void onExit(
@Advice.Enter final AgentScope scope,
@Advice.Enter final ContextScope scope,
@Advice.Argument(value = 2) final Object serverTransaction,
@Advice.Thrown final Throwable throwable) {
if (scope == null) {
return;
}
final AgentSpan span = scope.span();
final Context context = scope.context();
final AgentSpan span = fromContext(context);
try {
if (null != serverTransaction) {
// manual DECORATE.onResponse(span, serverTransaction):
Expand All @@ -45,7 +48,7 @@ public static void onExit(
if (throwable != null) {
DECORATE.onError(span, throwable);
}
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(context);
} finally {
scope.close();
span.finish();
Expand Down
4 changes: 3 additions & 1 deletion ...trumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.AXWAY_TRY_TRANSACTION;
import static datadog.trace.instrumentation.axway.AxwayHTTPPluginDecorator.DECORATE;

import datadog.context.Context;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import net.bytebuddy.asm.Advice;
Expand Down Expand Up @@ -34,11 +35,12 @@ public static void onExit(
return;
}
final AgentSpan span = scope.span();
final Context context = scope.context();
try {
if (throwable != null) {
DECORATE.onError(span, throwable);
}
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(context);
} finally {
scope.close();
span.finish();
Expand Down
2 changes: 1 addition & 1 deletion ...ain/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ public static void methodExit(
final AgentSpan span = fromContext(scope.context());
DECORATE.onError(span, throwable);
DECORATE.onResponse(span, response);
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(scope.context());
scope.close();
span.finish();
}
Expand Down
17 changes: 9 additions & 8 deletions ...ation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
package datadog.trace.instrumentation.cxf;

import com.google.auto.service.AutoService;
import datadog.context.Context;
import datadog.context.ContextScope;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers;
import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
import net.bytebuddy.asm.Advice;
Expand Down Expand Up @@ -55,21 +55,22 @@ public void methodAdvice(MethodTransformer transformer) {

public static class PropagateSpanAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) {
public static ContextScope beforeInvoke(@Advice.Argument(0) final Exchange exchange) {
if (exchange == null || exchange.getInMessage() == null || AgentTracer.activeSpan() != null) {
return null;
}
final Object span =
final Object contextObj =
ServletHelper.getServletRequestAttribute(
exchange.getInMessage().get("HTTP.REQUEST"), HttpServerDecorator.DD_SPAN_ATTRIBUTE);
if (span instanceof AgentSpan) {
return AgentTracer.activateSpan((AgentSpan) span);
exchange.getInMessage().get("HTTP.REQUEST"),
HttpServerDecorator.DD_CONTEXT_ATTRIBUTE);
if (contextObj instanceof Context) {
return ((Context) contextObj).attach();
}
return null;
}

@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
public static void afterInvoke(@Advice.Enter final AgentScope scope) {
public static void afterInvoke(@Advice.Enter final ContextScope scope) {
if (scope != null) {
scope.close();
}
Expand Down
15 changes: 8 additions & 7 deletions ...natra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR;
import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE;
import static datadog.trace.instrumentation.finatra.FinatraDecorator.FINATRA_CONTROLLER;
Expand All @@ -17,9 +18,9 @@
import com.twitter.finagle.http.Request;
import com.twitter.finagle.http.Response;
import com.twitter.util.Future;
import datadog.context.ContextScope;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.Tags;
import net.bytebuddy.asm.Advice;
Expand Down Expand Up @@ -61,7 +62,7 @@ public void methodAdvice(MethodTransformer transformer) {

public static class RouteAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope nameSpan(
public static ContextScope nameSpan(
@Advice.Argument(0) final Request request,
@Advice.FieldValue("path") final String path,
@Advice.FieldValue("clazz") final Class clazz) {
Expand All @@ -78,23 +79,23 @@ public static AgentScope nameSpan(
DECORATE.afterStart(span);
span.setResourceName(DECORATE.className(clazz));

return activateSpan(span);
return getCurrentContext().with(span).attach();
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void setupCallback(
@Advice.Enter final AgentScope scope,
@Advice.Enter final ContextScope scope,
@Advice.Thrown final Throwable throwable,
@Advice.Return final Some<Future<Response>> responseOption) {

if (scope == null) {
return;
}

final AgentSpan span = scope.span();
final AgentSpan span = spanFromContext(scope.context());
if (throwable != null) {
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
DECORATE.beforeFinish(scope.context());
span.finish();
scope.close();
return;
Expand Down
Loading