diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java index bb575d92080..9cde0e0e244 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/BaseDecorator.java @@ -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; @@ -76,8 +77,8 @@ 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; } @@ -85,6 +86,10 @@ public AgentSpan beforeFinish(final AgentSpan span) { return span; } + public Context beforeFinish(final Context context) { + return context; + } + public AgentScope onError(final AgentScope scope, final Throwable throwable) { onError(scope.span(), throwable); return scope; diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 76163e9bb6f..e2a53ee7e8c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -56,7 +56,7 @@ public abstract class HttpServerDecorator callIGCallbackURI( } @Override - public AgentSpan beforeFinish(AgentSpan span) { - // TODO Migrate beforeFinish to Context API + public Context beforeFinish(Context context) { + AgentSpan span = AgentSpan.fromContext(context); onRequestEndForInstrumentationGateway(span); + // Close Serverless Gateway Inferred Span if any - // finishInferredProxySpan(context); - return super.beforeFinish(span); + finishInferredProxySpan(context); + + return super.beforeFinish(context); } protected void finishInferredProxySpan(Context context) { diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java index 195af7e84c9..91850d3621f 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogAsyncHandlerWrapper.java @@ -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; @@ -29,7 +30,8 @@ public DatadogAsyncHandlerWrapper( @Override public Future 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 futureResponse; // handle blocking in the beginning of the request @@ -38,7 +40,7 @@ public Future 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$.successful().apply(response); } @@ -46,7 +48,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(context, t); throw t; } @@ -67,14 +69,14 @@ public HttpResponse apply(HttpResponse response) { response = newResponse; } - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(context, response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(context, t); return t; } }, diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java index 4becaf8fe2c..29f0cb497d6 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java @@ -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; @@ -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 @@ -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 diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index 4c45b352c83..2acf49c3ad6 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java @@ -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(); } diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java index b698fbdedb9..e887ffec7f1 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/HTTPPluginAdvice.java @@ -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): @@ -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(); diff --git a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java index c9cc01ba869..9c63a5d31d9 100644 --- a/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java +++ b/dd-java-agent/instrumentation/axway-api/src/main/java/datadog/trace/instrumentation/axway/StateAdvice.java @@ -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; @@ -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(); diff --git a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java index 472f45bbc77..5edb0f4030e 100644 --- a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java +++ b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java @@ -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(); } diff --git a/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java b/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java index 4a5ecbc715d..775fe5fced3 100644 --- a/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java +++ b/dd-java-agent/instrumentation/cxf-2.1/src/main/java/datadog/trace/instrumentation/cxf/InvokerInstrumentation.java @@ -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; @@ -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(); } diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java index cb3a307001b..e93d0858dc6 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/FinatraInstrumentation.java @@ -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; @@ -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; @@ -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) { @@ -78,12 +79,12 @@ 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> responseOption) { @@ -91,10 +92,10 @@ public static void setupCallback( 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; diff --git a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java index 677ede1884f..6d948f86532 100644 --- a/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java +++ b/dd-java-agent/instrumentation/finatra-2.9/src/main/java/datadog/trace/instrumentation/finatra/Listener.java @@ -1,36 +1,40 @@ package datadog.trace.instrumentation.finatra; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.finatra.FinatraDecorator.DECORATE; import com.twitter.finagle.http.Response; import com.twitter.util.FutureEventListener; +import datadog.context.ContextScope; import datadog.trace.api.Config; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; public class Listener implements FutureEventListener { - private final AgentScope scope; + private final ContextScope scope; - public Listener(final AgentScope scope) { + public Listener(final ContextScope scope) { this.scope = scope; } @Override public void onSuccess(final Response response) { + final AgentSpan span = spanFromContext(scope.context()); // Don't use DECORATE.onResponse because this is the controller span if (Config.get().getHttpServerErrorStatuses().get(DECORATE.status(response))) { - scope.span().setError(true); + span.setError(true); } - DECORATE.beforeFinish(scope.span()); - scope.span().finish(); + DECORATE.beforeFinish(scope.context()); + span.finish(); scope.close(); } @Override public void onFailure(final Throwable cause) { - DECORATE.onError(scope.span(), cause); - DECORATE.beforeFinish(scope.span()); - scope.span().finish(); + final AgentSpan span = spanFromContext(scope.context()); + DECORATE.onError(span, cause); + DECORATE.beforeFinish(scope.context()); + span.finish(); scope.close(); } } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java index 0be7ae6ee01..076d3ab66ae 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyBlockingHelper.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.grizzly; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.blocking.BlockingActionHelper; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -35,14 +37,14 @@ public class GrizzlyBlockingHelper { private GrizzlyBlockingHelper() {} public static boolean block( - Request request, Response response, Flow.Action.RequestBlockingAction rba, AgentSpan span) { + Request request, Response response, Flow.Action.RequestBlockingAction rba, Context context) { return block( request, response, rba.getStatusCode(), rba.getBlockingContentType(), rba.getExtraHeaders(), - span); + context); } public static boolean block( @@ -51,11 +53,12 @@ public static boolean block( int statusCode, BlockingContentType bct, Map extraHeaders, - AgentSpan span) { + Context context) { if (GET_OUTPUT_STREAM == null) { return false; } + AgentSpan span = spanFromContext(context); try { OutputStream os = (OutputStream) GET_OUTPUT_STREAM.invoke(response); response.setStatus(BlockingActionHelper.getHttpCode(statusCode)); @@ -76,13 +79,17 @@ public static boolean block( os.close(); response.finish(); - span.getRequestContext().getTraceSegment().effectivelyBlocked(); + if (span != null) { + span.getRequestContext().getTraceSegment().effectivelyBlocked(); + } SpanClosingListener.LISTENER.onAfterService(request); } catch (Throwable e) { log.info("Error committing blocking response", e); - DECORATE.onError(span, e); - DECORATE.beforeFinish(span); - span.finish(); + if (span != null) { + DECORATE.onError(span, e); + DECORATE.beforeFinish(context); + span.finish(); + } } return true; diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java index 833f4043332..d5c228aa293 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java @@ -2,7 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -42,7 +42,7 @@ public static class HandleAdvice { @Advice.Local("contextScope") ContextScope scope, @Advice.Argument(0) final Request request, @Advice.Argument(1) final Response response) { - if (request.getAttribute(DD_SPAN_ATTRIBUTE) != null) { + if (request.getAttribute(DD_CONTEXT_ATTRIBUTE) != null) { return false; } @@ -54,14 +54,14 @@ public static class HandleAdvice { scope = context.attach(); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); if (rba != null) { - boolean success = GrizzlyBlockingHelper.block(request, response, rba, span); + boolean success = GrizzlyBlockingHelper.block(request, response, rba, context); if (success) { return true; /* skip body */ } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java index c5c243e3369..f6ae98af71c 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-2.0/src/main/java/datadog/trace/instrumentation/grizzly/SpanClosingListener.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.grizzly; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.glassfish.grizzly.http.server.AfterServiceListener; import org.glassfish.grizzly.http.server.Request; @@ -12,13 +14,17 @@ public class SpanClosingListener implements AfterServiceListener { @Override public void onAfterService(final Request request) { - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanAttr instanceof AgentSpan) { - request.removeAttribute(DD_SPAN_ATTRIBUTE); - final AgentSpan span = (AgentSpan) spanAttr; - DECORATE.onResponse(span, request.getResponse()); - DECORATE.beforeFinish(span); - span.finish(); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + request.removeAttribute(DD_CONTEXT_ATTRIBUTE); + + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, request.getResponse()); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java index bab2a4218ed..84efdb6e24b 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/DefaultFilterChainInstrumentation.java @@ -1,15 +1,15 @@ package datadog.trace.instrumentation.grizzlyhttp232; 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.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPrivate; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.glassfish.grizzly.filterchain.FilterChainContext; @@ -40,22 +40,21 @@ public void methodAdvice(MethodTransformer transformer) { public static class PropagateServerSpanAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(2) final FilterChainContext ctx) { + public static ContextScope onEnter(@Advice.Argument(2) final FilterChainContext ctx) { final AgentSpan active = activeSpan(); // don't activate a span if already one is active if (active != null) { return null; } - final Object span = ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (span instanceof AgentSpan) { - // activate the http server span when nothing is already active - return activateSpan((AgentSpan) span); + final Object contextObj = ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + return ((Context) contextObj).attach(); } return null; } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void onExit(@Advice.Enter final AgentScope scope) { + public static void onExit(@Advice.Enter final ContextScope scope) { if (scope != null) { scope.close(); } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java index d7a20c33e99..9e39339d732 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/FilterAdvice.java @@ -1,27 +1,30 @@ package datadog.trace.instrumentation.grizzlyhttp232; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.context.Context; +import datadog.context.ContextScope; import net.bytebuddy.asm.Advice; import org.glassfish.grizzly.filterchain.FilterChainContext; public class FilterAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) { - Object span = ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (span == null || activeSpan() != null) { + public static ContextScope onEnter(@Advice.Argument(0) final FilterChainContext ctx) { + if (getCurrentContext() != getRootContext()) { return null; } - return activateSpan((AgentSpan) span); + Object contextObj = ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + return ((Context) contextObj).attach(); + } + return null; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void onExit(@Advice.Enter final AgentScope scope) { + public static void onExit(@Advice.Enter final ContextScope scope) { if (scope != null) { scope.close(); } diff --git a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index 978a0d65824..9c75322a1aa 100644 --- a/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.grizzlyhttp232; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import datadog.appsec.api.blocking.BlockingContentType; import datadog.context.Context; @@ -87,20 +88,22 @@ protected CharSequence component() { public static void onHttpServerFilterPrepareResponseEnter( FilterChainContext ctx, HttpResponsePacket responsePacket) { - AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (null != span) { + Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { DECORATE.onResponse(span, responsePacket); } } public static void onHttpServerFilterPrepareResponseExit( FilterChainContext ctx, HttpResponsePacket responsePacket) { - AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (null != span) { - DECORATE.beforeFinish(span); + Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { + DECORATE.beforeFinish(context); span.finish(); } - ctx.getAttributes().removeAttribute(DD_SPAN_ATTRIBUTE); + ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); } @@ -108,7 +111,7 @@ public static NextAction onHttpCodecFilterExit( FilterChainContext ctx, HttpHeader httpHeader, HttpCodecFilter thiz, NextAction nextAction) { // only create a span if there isn't another one attached to the current ctx // and if the httpHeader has been parsed into a HttpRequestPacket - if (ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE) != null + if (ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE) != null || !(httpHeader instanceof HttpRequestPacket)) { return nextAction; } @@ -119,8 +122,8 @@ public static NextAction onHttpCodecFilterExit( ContextScope scope = context.attach(); AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); - ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span); ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse); + ctx.getAttributes().setAttribute(DD_CONTEXT_ATTRIBUTE, context); DECORATE.onRequest(span, httpRequest, httpRequest, parentContext); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -145,13 +148,14 @@ public static NextAction onHttpCodecFilterExit( } public static void onFilterChainFail(FilterChainContext ctx, Throwable throwable) { - AgentSpan span = (AgentSpan) ctx.getAttributes().getAttribute(DD_SPAN_ATTRIBUTE); - if (null != span) { + Context context = (Context) ctx.getAttributes().getAttribute(DD_CONTEXT_ATTRIBUTE); + AgentSpan span; + if (context != null && (span = spanFromContext(context)) != null) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span).finish(); + DECORATE.beforeFinish(context); span.finish(); } - ctx.getAttributes().removeAttribute(DD_SPAN_ATTRIBUTE); + ctx.getAttributes().removeAttribute(DD_CONTEXT_ATTRIBUTE); ctx.getAttributes().removeAttribute(DD_RESPONSE_ATTRIBUTE); } diff --git a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java index fcba2189ce5..63f1e9a1de5 100644 --- a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java +++ b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java @@ -6,7 +6,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jaxrs.v1.InjectAdapter.SETTER; import static datadog.trace.instrumentation.jaxrs.v1.JaxRsClientV1Decorator.DECORATE; import static datadog.trace.instrumentation.jaxrs.v1.JaxRsClientV1Decorator.JAX_RS_CLIENT_CALL; @@ -67,12 +67,12 @@ public static AgentScope onEnter( @Advice.This final ClientHandler thisObj) { // WARNING: this might be a chain...so we only have to trace the first in the chain. - final boolean isRootClientHandler = null == request.getProperties().get(DD_SPAN_ATTRIBUTE); + final boolean isRootClientHandler = null == request.getProperties().get(DD_CONTEXT_ATTRIBUTE); if (isRootClientHandler) { final AgentSpan span = startSpan(JAX_RS_CLIENT_CALL); DECORATE.afterStart(span); DECORATE.onRequest(span, request); - request.getProperties().put(DD_SPAN_ATTRIBUTE, span); + request.getProperties().put(DD_CONTEXT_ATTRIBUTE, span); DECORATE.injectContext(getCurrentContext().with(span), request.getHeaders(), SETTER); return activateSpan(span); } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/HandleAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/HandleAdvice.java index ece1c7720d2..6c18ce73c2e 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/HandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/HandleAdvice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty10; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; import datadog.context.Context; @@ -19,9 +19,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -31,7 +31,7 @@ public static ContextScope onEnter( DECORATE.onRequest(span, req, req, parentContext); final ContextScope scope = context.attach(); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java index 009965e4337..4dd2494355a 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyCommitResponseHelper.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.jetty10; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -49,19 +51,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - state.partialResponse(); - return false; - } - - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { - state.partialResponse(); - return false; - } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { state.partialResponse(); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java index 8e7e7c68cd2..f398135f14f 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ResetAdvice.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -16,12 +18,15 @@ public class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - JettyDecorator.OnResponse.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + JettyDecorator.OnResponse.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java index 4c759e94ea0..b121af3f8eb 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/ServerHandleAdvice.java @@ -1,10 +1,11 @@ package datadog.trace.instrumentation.jetty10; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -12,12 +13,15 @@ class ServerHandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpChannel channel, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = channel.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // same logic as in Servlet3Advice. We need to activate/finish the dispatch span here // because we don't know if a servlet is going to be called and therefore whether // Servlet3Advice will have an opportunity to run. @@ -30,11 +34,18 @@ static AgentScope onEnter( // this is not great, but we leave the attribute. This is because // Set{Servlet,Context}PathAdvice // looks for this attribute, and we need a way to tell Servlet3Advice not to activate - // the root span, stored in DD_SPAN_ATTRIBUTE. + // the root span, stored in DD_CONTEXT_ATTRIBUTE. // req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -42,7 +53,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -55,7 +66,8 @@ public static void onExit( } if (!req.isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java index 0fcf638e53a..6c0540843d9 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetContextPathAdvice.java @@ -1,10 +1,12 @@ package datadog.trace.instrumentation.jetty10; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.Request; @@ -18,11 +20,15 @@ public class SetContextPathAdvice { public static void updateContextPath( @Advice.This final Request req, @Advice.Argument(0) final String contextPath) { if (contextPath != null) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_CONTEXT, contextPath); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_CONTEXT, contextPath); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java index 571266d94b0..af307de6144 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/SetServletPathAdvice.java @@ -1,10 +1,12 @@ package datadog.trace.instrumentation.jetty10; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import javax.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; @@ -21,11 +23,15 @@ public class SetServletPathAdvice { public static void updateServletPath( @Advice.This final Request req, @Advice.Argument(0) final String servletPath) { if (servletPath != null && !servletPath.isEmpty()) { // bypass cleanup - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_PATH, servletPath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_PATH, servletPath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java index 65dbaffbe4e..fa7c46b710c 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java @@ -1,8 +1,7 @@ package datadog.trace.instrumentation.jetty11; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE; import datadog.context.Context; @@ -21,9 +20,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -34,7 +33,7 @@ public static ContextScope onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, req, req, parentContext); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -58,12 +57,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java index 109d6c73467..e8007c38d23 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/SetContextPathAdvice.java @@ -2,11 +2,13 @@ import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.Request; @@ -22,20 +24,24 @@ public static void updateContextPath( @Advice.This final Request req, @Advice.Argument(0) final ContextHandler.Context context, @Advice.Argument(1) final String pathInContext) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - if (context != null && context.getContextPath() != null) { - final String servletContext = context.getContextPath(); - ((AgentSpan) span).setTag(SERVLET_CONTEXT, servletContext); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, servletContext); - if (pathInContext != null) { - final String relativePath = - pathInContext.startsWith(servletContext) - ? pathInContext.substring(servletContext.length()) - : pathInContext; - ((AgentSpan) span).setTag(SERVLET_PATH, relativePath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, relativePath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context ctx = (Context) contextObj; + AgentSpan span = spanFromContext(ctx); + if (span != null) { + if (context != null && context.getContextPath() != null) { + final String servletContext = context.getContextPath(); + span.setTag(SERVLET_CONTEXT, servletContext); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, servletContext); + if (pathInContext != null) { + final String relativePath = + pathInContext.startsWith(servletContext) + ? pathInContext.substring(servletContext.length()) + : pathInContext; + span.setTag(SERVLET_PATH, relativePath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, relativePath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java index ed8c193b334..c60fbe0d575 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty12; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty12.JettyDecorator.DECORATE; import datadog.context.Context; @@ -20,9 +20,9 @@ public static void onExit( @Advice.This final HttpChannelState channel, @Advice.Return(readOnly = false) Runnable ret) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - try (final ContextScope scope = ((AgentSpan) existingSpan).attach()) { + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + try (final ContextScope scope = ((Context) existingContext).attach()) { ret = JettyRunnableWrapper.wrapIfNeeded(ret); return; } @@ -36,7 +36,7 @@ public static void onExit( DECORATE.afterStart(span); DECORATE.onRequest(span, req, req, parentContext); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); ret = JettyRunnableWrapper.wrapIfNeeded(ret); @@ -52,12 +52,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannelState channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java index d954023d95d..6b1eb97e6b5 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/SetContextPathAdvice.java @@ -2,11 +2,13 @@ import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty12.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.jetty12.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.http.HttpURI; @@ -21,33 +23,37 @@ public class SetContextPathAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void updateContextPath( @Advice.This final ContextHandler contextHandler, @Advice.Argument(0) final Request req) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - if (contextHandler != null && contextHandler.getContextPath() != null) { - final String servletContext = contextHandler.getContextPath(); - final String pathInContext; - final HttpURI uri = req.getHttpURI(); - if (contextHandler.getContext() != null && uri != null && uri.getDecodedPath() != null) { - pathInContext = contextHandler.getContext().getPathInContext(uri.getDecodedPath()); - } else { - pathInContext = null; - } - ((AgentSpan) span).setTag(SERVLET_CONTEXT, servletContext); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, servletContext); - if (pathInContext != null) { - // the following can be cached however than can be issues for application having - // dynamically generated URL - // since a bounded cache might collide - String relativePath = - pathInContext.startsWith(servletContext) - ? pathInContext.substring(servletContext.length()) - : pathInContext; - if (relativePath.isEmpty() || relativePath.charAt(0) != '/') { - relativePath = "/" + relativePath; + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + if (contextHandler != null && contextHandler.getContextPath() != null) { + final String servletContext = contextHandler.getContextPath(); + final String pathInContext; + final HttpURI uri = req.getHttpURI(); + if (contextHandler.getContext() != null && uri != null && uri.getDecodedPath() != null) { + pathInContext = contextHandler.getContext().getPathInContext(uri.getDecodedPath()); + } else { + pathInContext = null; + } + span.setTag(SERVLET_CONTEXT, servletContext); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, servletContext); + if (pathInContext != null) { + // the following can be cached however than can be issues for application having + // dynamically generated URL + // since a bounded cache might collide + String relativePath = + pathInContext.startsWith(servletContext) + ? pathInContext.substring(servletContext.length()) + : pathInContext; + if (relativePath.isEmpty() || relativePath.charAt(0) != '/') { + relativePath = "/" + relativePath; + } + span.setTag(SERVLET_PATH, relativePath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, relativePath); } - ((AgentSpan) span).setTag(SERVLET_PATH, relativePath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, relativePath); } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java index 4728c725ca0..f5311a224ec 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyCommitResponseInstrumentation.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.jetty70; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.gateway.BlockResponseFunction; @@ -70,13 +72,14 @@ static class CommitResponseAdvice { return false; } - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (!(contextObj instanceof Context)) { return false; } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + RequestContext requestContext; + if (span == null || (requestContext = span.getRequestContext()) == null) { return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java index af7f8a4c7d6..6953877d52c 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java @@ -2,7 +2,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isConstructor; @@ -149,10 +149,10 @@ public static ContextScope onEnter( @Advice.This final HttpConnection connection, @Advice.Local("agentSpan") AgentSpan span) { Request req = connection.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // Request already gone through initial processing, so just activate the span. - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + // Request already gone through initial processing, so just activate the context. + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -162,7 +162,7 @@ public static ContextScope onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, req, req, parentContext); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -183,12 +183,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpConnection channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java index 9cfd9c992b1..a79b766affd 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java @@ -4,14 +4,16 @@ import static datadog.trace.api.gateway.Events.EVENTS; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.advice.ActiveRequestContext; import datadog.trace.advice.RequiresRequestContext; import datadog.trace.agent.tooling.Instrumenter; @@ -62,11 +64,15 @@ public static class SetContextPathAdvice { public static void updateContextPath( @Advice.This final Request req, @Advice.Argument(0) final String contextPath) { if (contextPath != null) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_CONTEXT, contextPath); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_CONTEXT, contextPath); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + } } } } @@ -81,11 +87,15 @@ public static class SetServletPathAdvice { public static void updateServletPath( @Advice.This final Request req, @Advice.Argument(0) final String servletPath) { if (servletPath != null && !servletPath.isEmpty()) { // bypass cleanup - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_PATH, servletPath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_PATH, servletPath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java index 58ddd318ab4..0acc3649f3f 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/ServerHandleInstrumentation.java @@ -1,16 +1,17 @@ package datadog.trace.instrumentation.jetty70; 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.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty70.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; 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.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpConnection; @@ -53,12 +54,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpConnection connection, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = connection.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // see comments in HandleRequestAdvice for jetty-9 Object dispatchSpan; synchronized (req) { @@ -66,8 +70,15 @@ static AgentScope onEnter( } if (dispatchSpan instanceof AgentSpan) { span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -75,7 +86,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -88,7 +99,8 @@ public static void onExit( } if (!req.getAsyncContinuation().isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java index 380bf5a30d1..7f6dbac9c40 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyCommitResponseInstrumentation.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.jetty76; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.Reference; @@ -81,13 +83,14 @@ static class CommitResponseAdvice { return false; } - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (!(contextObj instanceof Context)) { return false; } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + RequestContext requestContext; + if (span == null || (requestContext = span.getRequestContext()) == null) { return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java index ec9ef262fb1..9b498dc79ac 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java @@ -2,8 +2,8 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isConstructor; @@ -150,10 +150,10 @@ public static ContextScope onEnter( @Advice.Local("newSpan") AgentSpan span) { Request req = connection.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // Request already gone through initial processing, so just activate the span. - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + // Request already gone through initial processing, so just activate the context. + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -163,7 +163,7 @@ public static ContextScope onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, req, req, parentContext); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -184,12 +184,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final AbstractHttpConnection connection) { Request req = connection.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, connection); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, connection); + DECORATE.beforeFinish(context); + span.finish(); + } } // Jetty doesn't always call async listeners diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java index be74a281258..346ec1b0e84 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java @@ -4,14 +4,16 @@ import static datadog.trace.api.gateway.Events.EVENTS; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.advice.ActiveRequestContext; import datadog.trace.advice.RequiresRequestContext; import datadog.trace.agent.tooling.Instrumenter; @@ -62,11 +64,15 @@ public static class SetContextPathAdvice { public static void updateContextPath( @Advice.This final Request req, @Advice.Argument(0) final String contextPath) { if (contextPath != null) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_CONTEXT, contextPath); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_CONTEXT, contextPath); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + } } } } @@ -81,11 +87,15 @@ public static class SetServletPathAdvice { public static void updateServletPath( @Advice.This final Request req, @Advice.Argument(0) final String servletPath) { if (servletPath != null && !servletPath.isEmpty()) { // bypass cleanup - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_PATH, servletPath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_PATH, servletPath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java index 86f239b9e69..ae8515aec21 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/ServerHandleInstrumentation.java @@ -1,16 +1,17 @@ package datadog.trace.instrumentation.jetty76; 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.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty76.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; 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.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.AbstractHttpConnection; @@ -53,12 +54,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) AbstractHttpConnection connection, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span) { req = connection.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // see comments in HandleRequestAdvice for jetty-9 Object dispatchSpan; synchronized (req) { @@ -66,8 +70,15 @@ static AgentScope onEnter( } if (dispatchSpan instanceof AgentSpan) { span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -75,7 +86,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("agentSpan") AgentSpan span, @Advice.Local("request") Request req, @Advice.Thrown Throwable t) { @@ -88,7 +99,8 @@ public static void onExit( } if (!req.getAsyncContinuation().isAsyncStarted()) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0.4/src/main/java/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0.4/src/main/java/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java index 7fb3bec4fcc..22548b59291 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0.4/src/main/java/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0.4/src/main/java/datadog/trace/instrumentation/jetty904/JettyCommitResponseHelper.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.jetty904; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -48,19 +50,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { - _committed.set(false); - return false; - } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java index 8e478cad61a..c5aebd187ef 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyCommitResponseInstrumentation.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.jetty9; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.Reference; @@ -98,19 +100,14 @@ static class CommitResponseAdvice { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { - _committed.set(false); - return false; - } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java index b85a639fb2b..c263f6ae4d3 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java @@ -4,7 +4,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy; @@ -165,9 +165,9 @@ public static ContextScope onEnter( @Advice.This final HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span) { Request req = channel.getRequest(); - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - return ((AgentSpan) existingSpan).attach(); + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final Context parentContext = DECORATE.extract(req); @@ -177,7 +177,7 @@ public static ContextScope onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, req, req, parentContext); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -197,12 +197,15 @@ public static class ResetAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void stopSpan(@Advice.This final HttpChannel channel) { Request req = channel.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, channel); - DECORATE.beforeFinish(span); - span.finish(); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); + if (span != null) { + DECORATE.onResponse(span, channel); + DECORATE.beforeFinish(context); + span.finish(); + } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java index ce06bf0ef21..996596c4008 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java @@ -4,14 +4,16 @@ import static datadog.trace.api.gateway.Events.EVENTS; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DD_SERVLET_PATH_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.advice.ActiveRequestContext; import datadog.trace.advice.RequiresRequestContext; import datadog.trace.agent.tooling.Instrumenter; @@ -69,11 +71,15 @@ public static class SetContextPathAdvice { public static void updateContextPath( @Advice.This final Request req, @Advice.Argument(0) final String contextPath) { if (contextPath != null) { - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_CONTEXT, contextPath); - req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_CONTEXT, contextPath); + req.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath); + } } } } @@ -88,11 +94,15 @@ public static class SetServletPathAdvice { public static void updateServletPath( @Advice.This final Request req, @Advice.Argument(0) final String servletPath) { if (servletPath != null && !servletPath.isEmpty()) { // bypass cleanup - Object span = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); // Don't want to update while being dispatched to new servlet - if (span instanceof AgentSpan && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { - ((AgentSpan) span).setTag(SERVLET_PATH, servletPath); - req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + if (contextObj instanceof Context && req.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE) == null) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + span.setTag(SERVLET_PATH, servletPath); + req.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath); + } } } } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java index 477e60d80b5..590390d8404 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/ServerHandleInstrumentation.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.jetty9; 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.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; @@ -9,9 +9,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; 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.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.eclipse.jetty.server.HttpChannel; @@ -69,12 +70,15 @@ public void methodAdvice(MethodTransformer transformer) { static class HandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentScope onEnter( + static ContextScope onEnter( @Advice.Argument(0) HttpChannel channel, @Advice.Local("agentSpan") AgentSpan span, @Advice.Local("request") Request req) { req = channel.getRequest(); + // First check if there's an existing context in the request (from main server span) + Object existingContext = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + // same logic as in Servlet3Advice. We need to activate/finish the dispatch span here // because we don't know if a servlet is going to be called and therefore whether // Servlet3Advice will have an opportunity to run. @@ -92,11 +96,18 @@ static AgentScope onEnter( // this is not great, but we leave the attribute. This is because // Set{Servlet,Context}PathAdvice // looks for this attribute, and we need a way to tell Servlet3Advice not to activate - // the root span, stored in DD_SPAN_ATTRIBUTE. + // the root span, stored in DD_CONTEXT_ATTRIBUTE. // req.removeAttribute(DD_DISPATCH_SPAN_ATTRIBUTE); span = (AgentSpan) dispatchSpan; - AgentScope scope = activateSpan(span); - return scope; + + // If we have an existing context, create a new context with the dispatch span + // Otherwise just attach the dispatch span + if (existingContext instanceof Context) { + Context contextWithDispatchSpan = ((Context) existingContext).with(span); + return contextWithDispatchSpan.attach(); + } else { + return span.attach(); + } } return null; @@ -104,7 +115,7 @@ static AgentScope onEnter( @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) public static void onExit( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Local("request") Request req, @Advice.Local("agentSpan") AgentSpan span, @Advice.Thrown Throwable t) { @@ -131,7 +142,8 @@ public static void onExit( } if (!(req.isAsyncStarted() && registeredFinishListener)) { // finish will be handled by the async listener - DECORATE.beforeFinish(span); + // Use the full context from the scope for beforeFinish + DECORATE.beforeFinish(scope.context()); span.finish(); } scope.close(); diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.3/src/main/java/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.3/src/main/java/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java index 3e9afaebc7d..5e4dc5fa078 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.3/src/main/java/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.3/src/main/java/datadog/trace/instrumentation/jetty93/JettyCommitResponseHelper.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.jetty93; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -48,19 +50,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - _committed.set(false); - return false; - } - - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { - _committed.set(false); - return false; - } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { _committed.set(false); return false; } diff --git a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.4.21/src/main/java/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.4.21/src/main/java/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java index 2fd93b65141..ba34bd6412d 100644 --- a/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.4.21/src/main/java/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java +++ b/dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.4.21/src/main/java/datadog/trace/instrumentation/jetty9421/JettyCommitResponseHelper.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.jetty9421; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty9.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -51,19 +53,14 @@ public class JettyCommitResponseHelper { Request req = connection.getRequest(); - if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null) { - state.partialResponse(); - return false; - } - - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (!(existingSpan instanceof AgentSpan)) { - state.partialResponse(); - return false; - } - AgentSpan span = (AgentSpan) existingSpan; - RequestContext requestContext = span.getRequestContext(); - if (requestContext == null) { + Object contextObj; + Context context; + AgentSpan span; + RequestContext requestContext; + if (req.getAttribute(DD_IGNORE_COMMIT_ATTRIBUTE) != null + || !((contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE)) instanceof Context) + || (span = spanFromContext(context = (Context) contextObj)) == null + || (requestContext = span.getRequestContext()) == null) { state.partialResponse(); return false; } diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java index 1cb460dfbcb..cc984c4d071 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -96,9 +96,9 @@ public static class HandleRequestAdvice { // if we try to get an attribute that doesn't exist open liberty might complain with an // exception try { - Object existingSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - scope = ((AgentSpan) existingSpan).attach(); + Object existingContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + scope = ((Context) existingContext).attach(); return false; } } catch (NullPointerException e) { @@ -123,7 +123,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java index fa50b98f53b..7daf4045919 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/RequestFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty20; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -9,6 +10,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletRequest; import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.servlet.IExtendedResponse; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -57,19 +59,22 @@ public static void stopSpan(@Advice.This SRTServletRequest req) { // this should be a servlet response if (resp instanceof SRTServletResponse) { SRTServletResponse httpResp = (SRTServletResponse) resp; - Object spanObj = null; + Object contextObj = null; try { - spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - if (spanObj instanceof AgentSpan) { - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, httpResp); - DECORATE.beforeFinish(span); - span.finish(); + if (contextObj instanceof Context) { + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, httpResp); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java index 3bd85a7f2d1..5041e9ee6a6 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/ResponseFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty20; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty20.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -11,6 +12,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.WebContainerRequestState; import com.ibm.wsspi.webcontainer.servlet.IExtendedRequest; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,33 +82,39 @@ static void after(@Advice.Argument(0) boolean releaseChannel) { @SuppressFBWarnings("DCN_NULLPOINTER_EXCEPTION") public static class ResponseFinishAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentSpan onEnter(@Advice.This SRTServletResponse resp) { + static Context onEnter(@Advice.This SRTServletResponse resp) { // this is the last opportunity to have any meaningful // interaction with the response - AgentSpan span = null; + Context context = null; IExtendedRequest req = resp.getRequest(); try { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - span = (AgentSpan) spanObj; - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - DECORATE.onResponse(span, resp); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + context = (Context) contextObj; + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, resp); + } } } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - return span; + return context; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.This SRTServletResponse resp, @Advice.Enter AgentSpan span) { - if (span == null) { + @Advice.This SRTServletResponse resp, @Advice.Enter Context context) { + if (context == null) { return; } - DECORATE.beforeFinish(span); - span.finish(); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java index 2fb526942d7..644644bee89 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java @@ -2,9 +2,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.HttpInboundServiceContextImplInstrumentation.REQUEST_MSG_TYPE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -98,9 +98,9 @@ public static class HandleRequestAdvice { // if we try to get an attribute that doesn't exist open liberty might complain with an // exception try { - Object existingSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - scope = ((AgentSpan) existingSpan).attach(); + Object existingContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingContext instanceof Context) { + scope = ((Context) existingContext).attach(); return false; } } catch (NullPointerException e) { @@ -125,7 +125,7 @@ public static class HandleRequestAdvice { } DECORATE.afterStart(span); DECORATE.onRequest(span, request, request, parentContext); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, context); request.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java index a2b2172dde4..a6e633bcfba 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/RequestFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty23; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -9,6 +10,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletRequest; import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.servlet.IExtendedResponse; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -57,19 +59,22 @@ public static void stopSpan(@Advice.This SRTServletRequest req) { // this should be a servlet response if (resp instanceof SRTServletResponse) { SRTServletResponse httpResp = (SRTServletResponse) resp; - Object spanObj = null; + Object contextObj = null; try { - spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - if (spanObj instanceof AgentSpan) { - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; - DECORATE.onResponse(span, httpResp); - DECORATE.beforeFinish(span); - span.finish(); + if (contextObj instanceof Context) { + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + final Context context = (Context) contextObj; + final AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, httpResp); + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java index 96ebcc7252e..d8b4c2a013e 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/ResponseFinishInstrumentation.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.liberty23; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.liberty23.LibertyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -11,6 +12,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.wsspi.webcontainer.WebContainerRequestState; import com.ibm.wsspi.webcontainer.servlet.IExtendedRequest; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -80,33 +82,39 @@ static void after(@Advice.Argument(0) boolean releaseChannel) { @SuppressFBWarnings("DCN_NULLPOINTER_EXCEPTION") public static class ResponseFinishAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - static AgentSpan onEnter(@Advice.This SRTServletResponse resp) { + static Context onEnter(@Advice.This SRTServletResponse resp) { // this is the last opportunity to have any meaningful // interaction with the response - AgentSpan span = null; + Context context = null; IExtendedRequest req = resp.getRequest(); try { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - span = (AgentSpan) spanObj; - req.setAttribute(DD_SPAN_ATTRIBUTE, null); - DECORATE.onResponse(span, resp); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + context = (Context) contextObj; + req.setAttribute(DD_CONTEXT_ATTRIBUTE, null); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.onResponse(span, resp); + } } } catch (NullPointerException e) { // OpenLiberty will throw NPE on getAttribute if the response has already been closed. } - return span; + return context; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.This SRTServletResponse resp, @Advice.Enter AgentSpan span) { - if (span == null) { + @Advice.This SRTServletResponse resp, @Advice.Enter Context context) { + if (context == null) { return; } - DECORATE.beforeFinish(span); - span.finish(); + AgentSpan span = fromContext(context); + if (span != null) { + DECORATE.beforeFinish(context); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java index 778b35ddd43..47802519796 100644 --- a/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java +++ b/dd-java-agent/instrumentation/micronaut/src/main/java/datadog/trace/instrumentation/micronaut/WriteFinalNettyResponseAdvice.java @@ -1,10 +1,10 @@ package datadog.trace.instrumentation.micronaut; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.DECORATE; import static datadog.trace.instrumentation.micronaut.MicronautDecorator.SPAN_ATTRIBUTE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.micronaut.http.HttpRequest; import io.micronaut.http.MutableHttpResponse; @@ -20,9 +20,9 @@ public static void beginRequest( return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = getCurrentContext().with(span).attach()) { DECORATE.onResponse(span, message); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 4a211023143..34281d33a46 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -2,8 +2,8 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; 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.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.ContextStore; @@ -109,9 +110,9 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { DECORATE.onError(errorSpan, cause); - DECORATE.beforeFinish(errorSpan); + DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java index 862dc01cb2d..ae7c9a463dc 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java @@ -70,7 +70,7 @@ public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent ctx.sendUpstream(msg); } catch (final Throwable throwable) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false throw throwable; } diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index a32891cb820..984af432114 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -1,10 +1,10 @@ package datadog.trace.instrumentation.netty38.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.ContextScope; import datadog.trace.bootstrap.ContextStore; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import datadog.trace.instrumentation.netty38.ChannelTraceContext; @@ -36,7 +36,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = getCurrentContext().with(span).attach()) { final HttpResponse response = (HttpResponse) msg.getMessage(); try { @@ -58,7 +58,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m && (response.getStatus() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java index e08f930e495..8f8ab0d5a8c 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.netty40; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.api.GenericClassValue; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -16,7 +17,8 @@ public final class AttributeKeys { private static final ClassValue>> MAPS = GenericClassValue.constructing(ConcurrentHashMap.class); - public static final AttributeKey SPAN_ATTRIBUTE_KEY = attributeKey(DD_SPAN_ATTRIBUTE); + public static final AttributeKey CONTEXT_ATTRIBUTE_KEY = + attributeKey(DD_CONTEXT_ATTRIBUTE); public static final AttributeKey CLIENT_PARENT_ATTRIBUTE_KEY = attributeKey("datadog.client.parent.span"); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java index e0b426f7697..60834870a8c 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/ChannelFutureListenerInstrumentation.java @@ -2,8 +2,8 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; 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.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -94,9 +95,9 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); - NettyHttpServerDecorator.DECORATE.beforeFinish(errorSpan); + NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java index 7fe652004ff..7cc2884896d 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelHandlerContextInstrumentation.java @@ -6,13 +6,15 @@ 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.noopScope; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.NettyChannelPipelineInstrumentation.ADDITIONAL_INSTRUMENTATION_NAMES; import static datadog.trace.instrumentation.netty40.NettyChannelPipelineInstrumentation.INSTRUMENTATION_NAME; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -69,7 +71,8 @@ public void methodAdvice(MethodTransformer transformer) { public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { - final AgentSpan channelSpan = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan channelSpan = spanFromContext(storedContext); if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index fbbf2309d32..763e5e5c8e5 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -4,14 +4,16 @@ 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.instrumentation.netty40.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE_SECURE; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty40.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -76,6 +78,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); + final Context spanContext = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -90,7 +93,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java index b0cc383f5e1..75ec585739f 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientResponseTracingHandler.java @@ -2,10 +2,12 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -24,7 +26,13 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } final boolean finishSpan = msg instanceof HttpResponse; @@ -47,7 +55,14 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } + if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -68,7 +83,14 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } + if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java index 58bb49cb647..d627282b43b 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java @@ -3,8 +3,8 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; @@ -27,7 +27,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { Channel channel = ctx.channel(); if (!(msg instanceof HttpRequest)) { - final AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); if (span == null) { ctx.fireChannelRead(msg); // superclass does not throw } else { @@ -51,7 +52,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); - channel.attr(SPAN_ATTRIBUTE_KEY).set(span); + channel.attr(CONTEXT_ATTRIBUTE_KEY).set(context); channel.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).set(request.headers()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -67,9 +68,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { ctx.fireChannelRead(msg); } catch (final Throwable throwable) { DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } } @@ -81,7 +82,8 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { super.channelInactive(ctx); } finally { try { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndRemove(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); + final AgentSpan span = spanFromContext(storedContext); if (span != null && span.phasedFinish()) { // at this point we can just publish this span to avoid loosing the rest of the trace span.publish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java index 4c97c1a52bb..6b5dd488e76 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerResponseTracingHandler.java @@ -1,12 +1,13 @@ package datadog.trace.instrumentation.netty40.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.CHANNEL_ID; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import datadog.trace.util.RandomUtils; @@ -24,13 +25,14 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = storedContext.attach()) { final HttpResponse response = (HttpResponse) msg; try { @@ -39,7 +41,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.onError(span, throwable); span.setHttpStatusCode(500); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } final boolean isWebsocketUpgrade = @@ -58,8 +60,8 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann && (response.getStatus() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + DECORATE.beforeFinish(scope.context()); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java index 6e3ec443840..c5abbe0c807 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/MaybeBlockResponseHandler.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.netty40.server; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty40.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty40.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; import static io.netty.handler.codec.http.HttpHeaders.setContentLength; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; @@ -56,7 +58,8 @@ private static void markBlockedResponse(Channel ch) { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) throws Exception { Channel channel = ctx.channel(); - AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan span = spanFromContext(storedContext); RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java index 23211173f4a..f0d04a533d2 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/ChannelFutureListenerInstrumentation.java @@ -2,8 +2,8 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; 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.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.NETTY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.NETTY_CONNECT; @@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -94,9 +95,9 @@ public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture f final AgentSpan errorSpan = startSpan(NETTY_CONNECT).setTag(Tags.COMPONENT, "netty"); errorSpan.context().setIntegrationName(NETTY); - try (final AgentScope scope = activateSpan(errorSpan)) { + try (final ContextScope scope = getCurrentContext().with(errorSpan).attach()) { NettyHttpServerDecorator.DECORATE.onError(errorSpan, cause); - NettyHttpServerDecorator.DECORATE.beforeFinish(errorSpan); + NettyHttpServerDecorator.DECORATE.beforeFinish(scope.context()); errorSpan.finish(); } diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java index d532a5e1606..261c05662c0 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2MultiplexHandlerStreamChannelInstrumentation.java @@ -37,10 +37,10 @@ public static class PropagateContextAdvice { @Advice.OnMethodExit(suppress = Throwable.class) public static void afterCreate(@Advice.This Channel self) { if (self.parent() != null - && self.parent().hasAttr(AttributeKeys.SPAN_ATTRIBUTE_KEY) - && !self.hasAttr(AttributeKeys.SPAN_ATTRIBUTE_KEY)) { - self.attr(AttributeKeys.SPAN_ATTRIBUTE_KEY) - .set(self.parent().attr(AttributeKeys.SPAN_ATTRIBUTE_KEY).getAndRemove()); + && self.parent().hasAttr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY) + && !self.hasAttr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY)) { + self.attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY) + .set(self.parent().attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).getAndRemove()); } } } diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java index 66f89dc58d9..62363fda0cc 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelHandlerContextInstrumentation.java @@ -6,13 +6,15 @@ 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.noopScope; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.NettyChannelPipelineInstrumentation.ADDITIONAL_INSTRUMENTATION_NAMES; import static datadog.trace.instrumentation.netty41.NettyChannelPipelineInstrumentation.INSTRUMENTATION_NAME; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -69,7 +71,8 @@ public void methodAdvice(MethodTransformer transformer) { public static class FireAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) { - final AgentSpan channelSpan = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan channelSpan = spanFromContext(storedContext); if (channelSpan == null || channelSpan == activeSpan()) { // don't modify the scope return noopScope(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index c944965fd34..688e30f2dc7 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -4,14 +4,16 @@ 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.instrumentation.netty41.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE_SECURE; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty41.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -77,6 +79,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann NettyHttpClientDecorator decorate = isSecure ? DECORATE_SECURE : DECORATE; final AgentSpan span = startSpan("netty", NETTY_CLIENT_REQUEST); + final Context spanContext = getCurrentContext().with(span); try (final AgentScope scope = activateSpan(span)) { decorate.afterStart(span); decorate.onRequest(span, request); @@ -91,7 +94,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.injectContext(current(), request.headers(), SETTER); } - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(spanContext); try { ctx.write(msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java index 408c3530e20..ce764cf7bb9 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientResponseTracingHandler.java @@ -2,10 +2,12 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.CLIENT_PARENT_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.netty.channel.ChannelHandler; @@ -26,7 +28,13 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } if (span != null) { final boolean finishSpan = @@ -41,7 +49,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { span.finish(); } } else { - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(span)); + } } } @@ -56,7 +66,14 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } + if (span != null) { // If an exception is passed to this point, it likely means it was unhandled and the // client span won't be finished with a proper response, so we should finish the span here. @@ -77,7 +94,14 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { final Attribute parentAttr = ctx.channel().attr(CLIENT_PARENT_ATTRIBUTE_KEY); parentAttr.setIfAbsent(noopSpan()); final AgentSpan parent = parentAttr.get(); - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndSet(parent); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); + + // Set parent context back to maintain the same functionality as getAndSet(parent) + if (storedContext != null) { + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).set(storedContext.with(parent)); + } + if (span != null && span != parent) { try (final AgentScope scope = activateSpan(span)) { DECORATE.beforeFinish(span); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java index 30197668254..58892004518 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java @@ -3,8 +3,8 @@ import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; import datadog.context.Context; @@ -26,11 +26,11 @@ public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapte public void channelRead(final ChannelHandlerContext ctx, final Object msg) { Channel channel = ctx.channel(); if (!(msg instanceof HttpRequest)) { - final AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); - if (span == null) { + final Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + if (storedContext == null) { ctx.fireChannelRead(msg); // superclass does not throw } else { - try (final ContextScope scope = span.attach()) { + try (final ContextScope scope = storedContext.attach()) { ctx.fireChannelRead(msg); // superclass does not throw } } @@ -50,7 +50,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); - channel.attr(SPAN_ATTRIBUTE_KEY).set(span); + channel.attr(CONTEXT_ATTRIBUTE_KEY).set(context); channel.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).set(request.headers()); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); @@ -72,9 +72,9 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { The handler chain failed with exception - need to finish the span here */ DECORATE.onError(span, throwable); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(ignored.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } } @@ -86,7 +86,8 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception { super.channelInactive(ctx); } finally { try { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).getAndRemove(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove(); + final AgentSpan span = spanFromContext(storedContext); if (span != null && span.phasedFinish()) { // at this point we can just publish this span to avoid loosing the rest of the trace span.publish(); diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java index 07235193fc8..1cf31f91eae 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerResponseTracingHandler.java @@ -1,11 +1,12 @@ package datadog.trace.instrumentation.netty41.server; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext; import io.netty.channel.ChannelHandler; @@ -22,14 +23,15 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap @Override public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) { - final AgentSpan span = ctx.channel().attr(SPAN_ATTRIBUTE_KEY).get(); + final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get(); + final AgentSpan span = spanFromContext(storedContext); if (span == null || !(msg instanceof HttpResponse)) { ctx.write(msg, prm); return; } - try (final AgentScope scope = activateSpan(span)) { + try (final ContextScope scope = storedContext.attach()) { final HttpResponse response = (HttpResponse) msg; try { @@ -38,7 +40,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann DECORATE.onError(span, throwable); span.setHttpStatusCode(500); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); throw throwable; } final boolean isWebsocketUpgrade = @@ -52,9 +54,9 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann if (response.status() != HttpResponseStatus.CONTINUE && (response.status() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) { DECORATE.onResponse(span, response); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false - ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove(); + ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove(); } } } diff --git a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java index 9139e318703..7503e871d65 100644 --- a/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/MaybeBlockResponseHandler.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.netty41.server; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY; +import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY; -import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; import static io.netty.handler.codec.http.HttpHeaders.setContentLength; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; import datadog.trace.bootstrap.blocking.BlockingActionHelper; @@ -55,7 +57,8 @@ private static void markBlockedResponse(Channel ch) { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) throws Exception { Channel channel = ctx.channel(); - AgentSpan span = channel.attr(SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan span = spanFromContext(storedContext); RequestContext requestContext; if (span == null || (requestContext = span.getRequestContext()) == null) { super.write(ctx, msg, prm); diff --git a/dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java b/dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java index 41ab0146d3d..b6af132f173 100644 --- a/dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java +++ b/dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java @@ -1,7 +1,8 @@ package datadog.trace.instrumentation.netty41; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.api.GenericClassValue; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -16,7 +17,8 @@ public final class AttributeKeys { private static final ClassValue>> MAPS = GenericClassValue.constructing(ConcurrentHashMap.class); - public static final AttributeKey SPAN_ATTRIBUTE_KEY = attributeKey(DD_SPAN_ATTRIBUTE); + public static final AttributeKey CONTEXT_ATTRIBUTE_KEY = + attributeKey(DD_CONTEXT_ATTRIBUTE); public static final AttributeKey CLIENT_PARENT_ATTRIBUTE_KEY = attributeKey("datadog.client.parent.span"); diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java index 8842e6e9be0..0578eda8aba 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogAsyncHandlerWrapper.java @@ -32,7 +32,7 @@ public Future apply(final HttpRequest request) { futureResponse = userHandler.apply(request); } catch (final Throwable t) { scope.close(); - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); throw t; } final Future wrapped = @@ -40,14 +40,14 @@ public Future apply(final HttpRequest request) { new AbstractFunction1() { @Override public HttpResponse apply(final HttpResponse response) { - DatadogWrapperHelper.finishSpan(span, response); + DatadogWrapperHelper.finishSpan(scope.context(), response); return response; } }, new AbstractFunction1() { @Override public Throwable apply(final Throwable t) { - DatadogWrapperHelper.finishSpan(span, t); + DatadogWrapperHelper.finishSpan(scope.context(), t); return t; } }, diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java index 533a77e1ec4..00c1057b3f0 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java @@ -114,12 +114,12 @@ public void onPush() throws Exception { final HttpResponse response = grab(responseInlet); final ContextScope scope = scopes.poll(); if (scope != null) { - AgentSpan span = fromContext(scope.context()); - 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 AgentSpan activeSpan = activeSpan(); + AgentSpan span = fromContext(scope.context()); if (activeSpan == span) { scope.close(); } @@ -144,8 +144,7 @@ public void onUpstreamFailure(final Throwable ex) throws Exception { ContextScope scope = scopes.poll(); 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 diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index 5c15a75fa46..d4565c68bfc 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -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(); } diff --git a/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java b/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java index b3014547ce7..fdbc677b0c5 100644 --- a/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java @@ -55,14 +55,14 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); if (REPORT_HTTP_STATUS) { playControllerSpan.setHttpStatusCode(500); } - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java index 68b56de4ddb..6a51c2cf294 100644 --- a/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play/play-2.3/src/main/java/datadog/trace/instrumentation/play23/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play23; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -12,10 +14,12 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1, Object> { private static final Logger log = LoggerFactory.getLogger(RequestCompleteCallback.class); + private final ContextScope scope; private final AgentSpan span; - public RequestCompleteCallback(final AgentSpan span) { - this.span = span; + public RequestCompleteCallback(final ContextScope scope) { + this.scope = scope; + this.span = spanFromContext(scope.context()); } @Override @@ -28,7 +32,7 @@ public Object apply(final Try result) { DECORATE.onResponse(span, result.get()); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java b/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java index ff617daa07a..ce5713e5afa 100644 --- a/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java @@ -70,14 +70,14 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); if (REPORT_HTTP_STATUS) { playControllerSpan.setHttpStatusCode(500); } - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java index 7841c6c4977..07fd7efa2d6 100644 --- a/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play/play-2.4/src/main/java/datadog/trace/instrumentation/play24/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play24; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -12,9 +14,11 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1, Object> { private static final Logger log = LoggerFactory.getLogger(RequestCompleteCallback.class); private final AgentSpan span; + private final ContextScope scope; - public RequestCompleteCallback(final AgentSpan span) { - this.span = span; + public RequestCompleteCallback(final ContextScope scope) { + this.span = spanFromContext(scope.context()); + this.scope = scope; } @Override @@ -27,7 +31,7 @@ public Object apply(final Try result) { DECORATE.onResponse(span, result.get()); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java b/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java index 83d6a74cf76..89103f8e2b6 100644 --- a/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java @@ -72,11 +72,11 @@ public static void stopTraceOnResponse( if (throwable == null) { responseFuture.onComplete( - new RequestCompleteCallback(playControllerSpan), + new RequestCompleteCallback(playControllerScope), ((Action) thisAction).executionContext()); } else { DECORATE.onError(playControllerSpan, throwable); - DECORATE.beforeFinish(playControllerSpan); + DECORATE.beforeFinish(playControllerScope.context()); playControllerSpan.finish(); } playControllerScope.close(); diff --git a/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java b/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java index 14c96f8a3bd..97e84394109 100644 --- a/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java +++ b/dd-java-agent/instrumentation/play/play-2.6/src/main/java/datadog/trace/instrumentation/play26/RequestCompleteCallback.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.play26; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,9 +16,11 @@ public class RequestCompleteCallback extends scala.runtime.AbstractFunction1 result) { DECORATE.updateOn404Only(span, response); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } catch (final Throwable t) { log.debug("error in play instrumentation", t); } finally { diff --git a/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java b/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java index ad00fa0041d..1043cb5b20b 100644 --- a/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java +++ b/dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/TracingHandler.java @@ -1,9 +1,10 @@ package datadog.trace.instrumentation.ratpack; import static datadog.context.Context.root; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.ratpack.RatpackServerDecorator.DECORATE; import com.google.common.reflect.TypeToken; @@ -24,16 +25,17 @@ public final class TracingHandler implements Handler { Types.token(Flow.Action.RequestBlockingAction.class); /** This constant must stay in sync with datadog.trace.instrumentation.netty41.AttributeKeys. */ - public static final AttributeKey SERVER_ATTRIBUTE_KEY = - AttributeKey.valueOf(DD_SPAN_ATTRIBUTE); + public static final AttributeKey SERVER_CONTEXT_ATTRIBUTE_KEY = + AttributeKey.valueOf(DD_CONTEXT_ATTRIBUTE); @Override public void handle(final Context ctx) { final Request request = ctx.getRequest(); - final Attribute spanAttribute = - ctx.getDirectChannelAccess().getChannel().attr(SERVER_ATTRIBUTE_KEY); - final AgentSpan nettySpan = spanAttribute.get(); + final Attribute contextAttribute = + ctx.getDirectChannelAccess().getChannel().attr(SERVER_CONTEXT_ATTRIBUTE_KEY); + final datadog.context.Context nettyContext = contextAttribute.get(); + final AgentSpan nettySpan = nettyContext != null ? fromContext(nettyContext) : null; // Relying on executor instrumentation to assume the netty span is in context as the parent. final AgentSpan ratpackSpan = startSpan(DECORATE.spanName()).setMeasured(true); diff --git a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java index 6c6b16a938c..5a498a65829 100644 --- a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java +++ b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java @@ -81,7 +81,7 @@ public static void finishRequest( DECORATE.onError(span, error); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java index 6b82e220519..f9c1e49c30d 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java @@ -3,14 +3,16 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.ClassloaderConfigurationOverrides; @@ -97,12 +99,15 @@ public static AgentSpan before( } } - Object span = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (span instanceof AgentSpan + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context && CallDepthThreadLocalMap.incrementCallDepth(HttpServletRequest.class) == 0) { - final AgentSpan agentSpan = (AgentSpan) span; - ClassloaderConfigurationOverrides.maybeEnrichSpan(agentSpan); - return agentSpan; + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + return span; + } } return null; } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-2.2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-2.2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index 120fbb039a0..ae67ab27c27 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-2.2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-2.2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java @@ -1,7 +1,7 @@ package datadog.trace.instrumentation.servlet2; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE; import datadog.context.Context; @@ -37,13 +37,16 @@ public static boolean onEnter( } final HttpServletRequest httpServletRequest = (HttpServletRequest) request; - Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - final boolean hasServletTrace = spanAttr instanceof AgentSpan; - if (hasServletTrace) { - final AgentSpan span = (AgentSpan) spanAttr; - ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - // Tracing might already be applied by the FilterChain or a parent request (forward/include). - return false; + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context existingContext = (Context) contextAttr; + final AgentSpan span = spanFromContext(existingContext); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + // Tracing might already be applied by the FilterChain or a parent request + // (forward/include). + return false; + } } if (response instanceof HttpServletResponse) { @@ -58,7 +61,7 @@ public static boolean onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, httpServletRequest, httpServletRequest, parentContext); - httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); + httpServletRequest.setAttribute(DD_CONTEXT_ATTRIBUTE, context); httpServletRequest.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); httpServletRequest.setAttribute( @@ -85,13 +88,17 @@ public static void stopSpan( @Advice.Local("contextScope") final ContextScope scope, @Advice.Thrown final Throwable throwable) { // Set user.principal regardless of who created this span. - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); if (Config.get().isServletPrincipalEnabled() - && spanAttr instanceof AgentSpan + && contextAttr instanceof Context && request instanceof HttpServletRequest) { - final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); - if (principal != null) { - ((AgentSpan) spanAttr).setTag(DDTags.USER_NAME, principal.getName()); + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); + if (principal != null) { + span.setTag(DDTags.USER_NAME, principal.getName()); + } } } @@ -116,7 +123,7 @@ public static void stopSpan( } DECORATE.onError(span, throwable); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java index 39ddb3734aa..eeec7e5bf93 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java @@ -6,8 +6,9 @@ import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_DISPATCH; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.AsyncDispatcherDecorator.DECORATE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DD_CONTEXT_PATH_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DD_SERVLET_PATH_ATTRIBUTE; @@ -15,6 +16,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.CallDepthThreadLocalMap; @@ -84,33 +86,36 @@ public static boolean enter( } final ServletRequest request = context.getRequest(); - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanAttr instanceof AgentSpan) { - final AgentSpan parent = (AgentSpan) spanAttr; - - final AgentSpan span = startSpan(SERVLET_DISPATCH, parent.context()); - // This span should get finished by Servlet3Advice - // However, when using Jetty without servlets (directly org.eclipse.jetty.server.Handler), - // that's not the case (see jetty's HandleAdvice) - DECORATE.afterStart(span); - - // These are pulled from attributes because jetty clears them from the request too early. - span.setTag(SERVLET_CONTEXT, request.getAttribute(DD_CONTEXT_PATH_ATTRIBUTE)); - span.setTag(SERVLET_PATH, request.getAttribute(DD_SERVLET_PATH_ATTRIBUTE)); - - synchronized (request) { - request.setAttribute(DD_DISPATCH_SPAN_ATTRIBUTE, span); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context ctxt = (Context) contextAttr; + final AgentSpan parent = spanFromContext(ctxt); + if (parent != null) { + + final AgentSpan span = startSpan(SERVLET_DISPATCH, parent.context()); + // This span should get finished by Servlet3Advice + // However, when using Jetty without servlets (directly org.eclipse.jetty.server.Handler), + // that's not the case (see jetty's HandleAdvice) + DECORATE.afterStart(span); + + // These are pulled from attributes because jetty clears them from the request too early. + span.setTag(SERVLET_CONTEXT, request.getAttribute(DD_CONTEXT_PATH_ATTRIBUTE)); + span.setTag(SERVLET_PATH, request.getAttribute(DD_SERVLET_PATH_ATTRIBUTE)); + + synchronized (request) { + request.setAttribute(DD_DISPATCH_SPAN_ATTRIBUTE, span); + } + + if (args.length == 1 && args[0] instanceof String) { + span.setResourceName((String) args[0]); + } else if (args.length == 2 && args[1] instanceof String) { + span.setResourceName((String) args[1]); + } + + // We can't register FinishAsyncDispatchListener here. + // The dispatch may happen on an onTimeout/onError, and adding listeners + // when listeners are being iterated on causes a ConcurrentModificationException on jetty } - - if (args.length == 1 && args[0] instanceof String) { - span.setResourceName((String) args[0]); - } else if (args.length == 2 && args[1] instanceof String) { - span.setResourceName((String) args[1]); - } - - // We can't register FinishAsyncDispatchListener here. - // The dispatch may happen on an onTimeout/onError, and adding listeners - // when listeners are being iterated on causes a ConcurrentModificationException on jetty } return true; } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java index 3c70e9be54c..8d7080bba95 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/FinishAsyncDispatchListener.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.servlet3; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.TIMEOUT; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE; +import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; @@ -26,14 +28,16 @@ public class FinishAsyncDispatchListener implements AsyncListener, Runnable { private final AtomicBoolean activated; private final AgentSpan span; private final boolean doOnResponse; + private final ContextScope scope; - public FinishAsyncDispatchListener(final AgentSpan span, boolean doOnResponse) { - this(span, new AtomicBoolean(), doOnResponse); + public FinishAsyncDispatchListener(final ContextScope scope, boolean doOnResponse) { + this(scope, new AtomicBoolean(), doOnResponse); } public FinishAsyncDispatchListener( - final AgentSpan span, AtomicBoolean activated, boolean doOnResponse) { - this.span = span; + final ContextScope scope, AtomicBoolean activated, boolean doOnResponse) { + this.scope = scope; + this.span = spanFromContext(scope.context()); this.activated = activated; this.doOnResponse = doOnResponse; } @@ -54,7 +58,7 @@ public void onComplete(final AsyncEvent event) throws IOException { DECORATE.onError(span, (Throwable) error); } } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -62,7 +66,7 @@ public void onComplete(final AsyncEvent event) throws IOException { @Override public void run() { if (activated.compareAndSet(false, true)) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -74,7 +78,7 @@ public void onTimeout(final AsyncEvent event) throws IOException { span.setError(true); } span.setTag(TIMEOUT, event.getAsyncContext().getTimeout()); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } @@ -83,7 +87,7 @@ public void onTimeout(final AsyncEvent event) throws IOException { public void onError(final AsyncEvent event) throws IOException { if (event.getThrowable() != null && activated.compareAndSet(false, true)) { DECORATE.onError(span, event.getThrowable()); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); maybeFinishSpan(); } } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java index c123c212515..e8030aa367f 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java @@ -2,10 +2,10 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE; import datadog.context.Context; @@ -77,14 +77,16 @@ public static boolean onEnter( finishSpan = true; - Object spanAttrValue = request.getAttribute(DD_SPAN_ATTRIBUTE); - final boolean hasServletTrace = spanAttrValue instanceof AgentSpan; - if (hasServletTrace) { - final AgentSpan span = (AgentSpan) spanAttrValue; - ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - // Tracing might already be applied by other instrumentation, - // the FilterChain or a parent request (forward/include). - return false; + Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextAttr instanceof Context) { + final Context existingContext = (Context) contextAttr; + final AgentSpan span = spanFromContext(existingContext); + if (span != null) { + ClassloaderConfigurationOverrides.maybeEnrichSpan(span); + // Tracing might already be applied by other instrumentation, + // the FilterChain or a parent request (forward/include). + return false; + } } final Context parentContext = DECORATE.extract(httpServletRequest); @@ -95,7 +97,7 @@ public static boolean onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, httpServletRequest, httpServletRequest, parentContext); - httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); + httpServletRequest.setAttribute(DD_CONTEXT_ATTRIBUTE, context); httpServletRequest.setAttribute( CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); httpServletRequest.setAttribute( @@ -125,13 +127,17 @@ public static void stopSpan( rumServletWrapper.commit(); } // Set user.principal regardless of who created this span. - final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); if (Config.get().isServletPrincipalEnabled() - && spanAttr instanceof AgentSpan + && contextAttr instanceof Context && request instanceof HttpServletRequest) { - final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); - if (principal != null) { - ((AgentSpan) spanAttr).setTag(DDTags.USER_NAME, principal.getName()); + final Context context = (Context) contextAttr; + final AgentSpan span = spanFromContext(context); + if (span != null) { + final Principal principal = ((HttpServletRequest) request).getUserPrincipal(); + if (principal != null) { + span.setTag(DDTags.USER_NAME, principal.getName()); + } } } @@ -147,7 +153,7 @@ public static void stopSpan( if (request.isAsyncStarted()) { AtomicBoolean activated = new AtomicBoolean(); FinishAsyncDispatchListener asyncListener = - new FinishAsyncDispatchListener(span, activated, !isDispatch); + new FinishAsyncDispatchListener(scope, activated, !isDispatch); // Jetty doesn't always call the listener, if the request ends before request.setAttribute(DD_FIN_DISP_LIST_SPAN_ATTRIBUTE, asyncListener); try { @@ -161,7 +167,7 @@ public static void stopSpan( DECORATE.onResponse(span, resp); } if (finishSpan) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } @@ -184,7 +190,7 @@ public static void stopSpan( DECORATE.onError(span, throwable); } if (finishSpan) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); // Finish the span manually since finishSpanOnClose was false } } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy index 11bf3eb20e7..312caadca48 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy @@ -9,6 +9,10 @@ import datadog.trace.instrumentation.servlet3.AsyncDispatcherDecorator import datadog.trace.instrumentation.servlet3.HtmlRumServlet import datadog.trace.instrumentation.servlet3.TestServlet3 import datadog.trace.instrumentation.servlet3.XmlRumServlet +import datadog.context.Context + +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE import groovy.servlet.AbstractHttpServlet import org.eclipse.jetty.server.Request import org.eclipse.jetty.server.Server @@ -526,7 +530,8 @@ class ServeFromOnAsyncTimeout extends AbstractHttpServlet { @Override void onTimeout(AsyncEvent event) throws IOException { - AgentSpan span = event.getSuppliedRequest().getAttribute('datadog.span') + Context ddContext = (Context) event.getSuppliedRequest().getAttribute(DD_CONTEXT_ATTRIBUTE) + AgentSpan span = spanFromContext(ddContext) activateSpan(span).withCloseable { try { delegateServlet.service(req, resp) diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java index 28676b5b8df..60835aa76b0 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java @@ -3,12 +3,13 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf; -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.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +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.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.servlet.ServletRequestSetter.SETTER; import static datadog.trace.instrumentation.servlet.SpanNameCache.SERVLET_PREFIX; import static datadog.trace.instrumentation.servlet.SpanNameCache.SPAN_NAME_CACHE; @@ -21,10 +22,11 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; 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.bootstrap.InstrumentationContext; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; @@ -82,16 +84,16 @@ public void methodAdvice(MethodTransformer transformer) { public static class RequestDispatcherAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope start( + public static ContextScope start( @Advice.Origin("#m") final String method, @Advice.This final RequestDispatcher dispatcher, - @Advice.Local("_requestSpan") Object requestSpan, + @Advice.Local("_requestContext") Object requestContext, @Advice.Argument(0) final ServletRequest request) { final AgentSpan parentSpan = activeSpan(); - final Object servletSpanObject = request.getAttribute(DD_SPAN_ATTRIBUTE); + final Object contextAttr = request.getAttribute(DD_CONTEXT_ATTRIBUTE); final AgentSpan servletSpan = - servletSpanObject instanceof AgentSpan ? (AgentSpan) servletSpanObject : null; + contextAttr instanceof Context ? spanFromContext((Context) contextAttr) : null; if (parentSpan == null && servletSpan == null) { // Don't want to generate a new top-level span @@ -122,16 +124,19 @@ public static AgentScope start( DECORATE.injectContext(span, request, SETTER); // temporarily replace from request to avoid spring resource name bubbling up: - requestSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - request.setAttribute(DD_SPAN_ATTRIBUTE, span); + requestContext = request.getAttribute(DD_CONTEXT_ATTRIBUTE); - return activateSpan(span); + final ContextScope scope = getCurrentContext().with(span).attach(); + // Set the context after activation so we have the proper Context object + request.setAttribute(DD_CONTEXT_ATTRIBUTE, scope.context()); + + return scope; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stop( - @Advice.Enter final AgentScope scope, - @Advice.Local("_requestSpan") final Object requestSpan, + @Advice.Enter final ContextScope scope, + @Advice.Local("_requestContext") final Object requestContext, @Advice.Argument(0) final ServletRequest request, @Advice.Argument(1) final ServletResponse response, @Advice.Thrown final Throwable throwable) { @@ -139,15 +144,15 @@ public static void stop( return; } - if (requestSpan != null) { - // now add it back... - request.setAttribute(DD_SPAN_ATTRIBUTE, requestSpan); + if (requestContext != null) { + request.setAttribute(DD_CONTEXT_ATTRIBUTE, requestContext); } - DECORATE.onError(scope, throwable); - DECORATE.beforeFinish(scope); + final AgentSpan span = spanFromContext(scope.context()); + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/test/groovy/RequestDispatcherTest.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/test/groovy/RequestDispatcherTest.groovy index 305896a5bcb..3193f8dc3af 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/test/groovy/RequestDispatcherTest.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-common/src/test/groovy/RequestDispatcherTest.groovy @@ -1,6 +1,6 @@ import datadog.trace.agent.test.InstrumentationSpecification import datadog.trace.api.DDSpanTypes -import datadog.trace.core.DDSpan +import datadog.context.Context import javax.servlet.ServletException import javax.servlet.http.HttpServletRequest @@ -8,7 +8,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE import static datadog.trace.core.propagation.DatadogHttpCodec.SAMPLING_PRIORITY_KEY import static datadog.trace.core.propagation.DatadogHttpCodec.SPAN_ID_KEY import static datadog.trace.core.propagation.DatadogHttpCodec.TRACE_ID_KEY @@ -17,7 +17,7 @@ class RequestDispatcherTest extends InstrumentationSpecification { def request = Mock(HttpServletRequest) def response = Mock(HttpServletResponse) - def mockSpan = Stub(DDSpan) + def mockContext = Mock(Context) def dispatcher = new RequestDispatcherUtils(request, response) def "test dispatch no-parent"() { @@ -26,7 +26,7 @@ class RequestDispatcherTest extends InstrumentationSpecification { dispatcher.include("") then: - 2 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 2 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(2) { trace(1) { basicSpan(it, "forward-child") @@ -44,7 +44,7 @@ class RequestDispatcherTest extends InstrumentationSpecification { } then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(1) { trace(3) { basicSpan(it, "parent") @@ -67,11 +67,13 @@ class RequestDispatcherTest extends InstrumentationSpecification { 1 * request.setAttribute(SPAN_ID_KEY, _) 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) >> mockContext then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName.toString() == "servlet.$operation" }) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, { Context ctx -> + ctx != null && ctx != mockContext // Verify it's a new context + }) then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, mockContext) where: operation | method @@ -97,7 +99,7 @@ class RequestDispatcherTest extends InstrumentationSpecification { def th = thrown(ServletException) th == ex - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) assertTraces(1) { trace(3) { basicSpan(it, "parent", null, ex) @@ -122,11 +124,13 @@ class RequestDispatcherTest extends InstrumentationSpecification { 1 * request.setAttribute(SPAN_ID_KEY, _) 1 * request.setAttribute(SAMPLING_PRIORITY_KEY, _) then: - 1 * request.getAttribute(DD_SPAN_ATTRIBUTE) >> mockSpan + 1 * request.getAttribute(DD_CONTEXT_ATTRIBUTE) >> mockContext then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, { it.spanName.toString() == "servlet.$operation" }) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, { Context ctx -> + ctx != null && ctx != mockContext // Verify it's a new context + }) then: - 1 * request.setAttribute(DD_SPAN_ATTRIBUTE, mockSpan) + 1 * request.setAttribute(DD_CONTEXT_ATTRIBUTE, mockContext) where: operation | method diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala index 29436ce4bd4..f1c994b33d2 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHelper.scala @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.spray import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; @@ -14,7 +15,8 @@ object SprayHelper { def wrapRequestContext( ctx: RequestContext, span: AgentSpan, - parentContext: Context + parentContext: Context, + scope: ContextScope ): RequestContext = { ctx.withRouteResponseMapped(message => { DECORATE.onRequest(span, ctx, ctx.request, parentContext) @@ -23,7 +25,8 @@ object SprayHelper { case throwable: Throwable => DECORATE.onError(span, throwable) case x => } - DECORATE.beforeFinish(span) + DECORATE.beforeFinish(scope.context()) + scope.close() span.finish() message }) diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java index 9a44db3476d..7bb8e735eb8 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java @@ -36,7 +36,7 @@ public static ContextScope enter( ContextScope scope = context.attach(); DECORATE.afterStart(span); - ctx = SprayHelper.wrapRequestContext(ctx, span, parentContext); + ctx = SprayHelper.wrapRequestContext(ctx, span, parentContext, scope); return scope; } diff --git a/dd-java-agent/instrumentation/spring/spring-cloud-zuul-2.0/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-cloud-zuul-2.0/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java index 77714dabea7..f167301e029 100644 --- a/dd-java-agent/instrumentation/spring/spring-cloud-zuul-2.0/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-cloud-zuul-2.0/src/main/java/datadog/trace/instrumentation/springcloudzuul2/ZuulSendForwardFilterInstrumentation.java @@ -1,13 +1,15 @@ package datadog.trace.instrumentation.springcloudzuul2; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; import com.netflix.zuul.context.RequestContext; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -48,9 +50,10 @@ public static void onEnter( request = ctx.getRequest(); if (request != null) { // Capture the span from the request before forwarding. - Object span = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (span instanceof AgentSpan) { - parentSpan = (AgentSpan) span; + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + parentSpan = spanFromContext(context); } } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy index a271f69a743..c4a179c1b44 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/latestDepTest/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy @@ -15,7 +15,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE @WebAppConfiguration // TODO this doesn't exist in 3.1.0 so will need rework if we want to test that version // https://mvnrepository.com/artifact/org.springframework/spring-test-mvc?repo=springio-plugins-release @@ -39,12 +39,12 @@ class HandlerMappingResourceNameFilterForkedTest extends InstrumentationSpecific when: runUnderTrace("test-servlet", { - request.setAttribute(DD_SPAN_ATTRIBUTE, activeSpan()) + request.setAttribute(DD_CONTEXT_ATTRIBUTE, activeSpan().context()) filter.doFilterInternal(request, response, filterChain) }) then: - assert request.getAttributeNames().toList() == [DD_SPAN_ATTRIBUTE] + assert request.getAttributeNames().toList() == [DD_CONTEXT_ATTRIBUTE] 0 * response._ where: diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java index 82a310fe133..8fd5cb59219 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/DispatcherServletInstrumentation.java @@ -2,9 +2,10 @@ 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.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE_RENDER; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.RESPONSE_RENDER; @@ -14,9 +15,9 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +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.ErrorPriorities; import java.util.List; @@ -95,20 +96,21 @@ public static void afterRefresh( public static class RenderAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { + public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(RESPONSE_RENDER); DECORATE_RENDER.afterStart(span); DECORATE_RENDER.onRender(span, mv); - return activateSpan(span); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + final AgentSpan span = spanFromContext(scope.context()); DECORATE_RENDER.onError(scope, throwable); - DECORATE_RENDER.beforeFinish(scope); + DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java index 9304644ea12..6558878f8ad 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerAdapterInstrumentation.java @@ -3,11 +3,12 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; 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.getRootContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; @@ -17,9 +18,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; 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.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import javax.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; @@ -66,16 +68,20 @@ public void methodAdvice(MethodTransformer transformer) { public static class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope nameResourceAndStartSpan( + public static ContextScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, @Advice.Argument(2) final Object handler, @Advice.Local("handlerSpanKey") String handlerSpanKey) { handlerSpanKey = ""; // Name the parent span based on the matching pattern - Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - DECORATE.onRequest((AgentSpan) parentSpan, request, request, getRootContext()); + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + DECORATE.onRequest(parentSpan, request, request, getRootContext()); + } } if (activeSpan() == null) { @@ -91,22 +97,25 @@ public static AgentScope nameResourceAndStartSpan( handlerKey = handler.getClass().getName(); } handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; - final Object existingSpan = request.getAttribute(handlerSpanKey); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + + // If the context already exists, return it + final Object existingContext = request.getAttribute(handlerSpanKey); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); DECORATE.onHandle(span, handler); + request.setAttribute(handlerSpanKey, span); - return activateSpan(span); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( @Advice.Argument(0) final HttpServletRequest request, - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable, @Advice.Local("handlerSpanKey") String handlerSpanKey) { if (scope == null) { @@ -115,14 +124,14 @@ public static void stopSpan( boolean finish = !Boolean.TRUE.equals( request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); - final AgentSpan span = scope.span(); + final AgentSpan span = spanFromContext(scope.context()); scope.close(); if (throwable != null) { DECORATE.onError(span, throwable); finish = true; } if (finish) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java index f78268e61fc..073bce859f1 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilter.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.springweb; import static datadog.context.Context.root; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb.SpringWebHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; import java.util.List; @@ -33,18 +35,22 @@ protected void doFilterInternal( final FilterChain filterChain) throws ServletException, IOException { - final Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - PathMatchingHttpServletRequestWrapper wrappedRequest = - new PathMatchingHttpServletRequestWrapper(request); - try { - if (findMapping(wrappedRequest)) { - // Name the parent span based on the matching pattern - // Let the parent span resource name be set with the attribute set in findMapping. - DECORATE.onRequest((AgentSpan) parentSpan, wrappedRequest, wrappedRequest, root()); + final Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + PathMatchingHttpServletRequestWrapper wrappedRequest = + new PathMatchingHttpServletRequestWrapper(request); + try { + if (findMapping(wrappedRequest)) { + // Name the parent span based on the matching pattern + // Let the parent span resource name be set with the attribute set in findMapping. + DECORATE.onRequest(parentSpan, wrappedRequest, wrappedRequest, root()); + } + } catch (final Exception ignored) { + // mapping.getHandler() threw exception. Ignore } - } catch (final Exception ignored) { - // mapping.getHandler() threw exception. Ignore } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy index a271f69a743..c4a179c1b44 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/datadog/trace/instrumentation/springweb/HandlerMappingResourceNameFilterForkedTest.groovy @@ -15,7 +15,7 @@ import javax.servlet.http.HttpServletResponse import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE @WebAppConfiguration // TODO this doesn't exist in 3.1.0 so will need rework if we want to test that version // https://mvnrepository.com/artifact/org.springframework/spring-test-mvc?repo=springio-plugins-release @@ -39,12 +39,12 @@ class HandlerMappingResourceNameFilterForkedTest extends InstrumentationSpecific when: runUnderTrace("test-servlet", { - request.setAttribute(DD_SPAN_ATTRIBUTE, activeSpan()) + request.setAttribute(DD_CONTEXT_ATTRIBUTE, activeSpan().context()) filter.doFilterInternal(request, response, filterChain) }) then: - assert request.getAttributeNames().toList() == [DD_SPAN_ATTRIBUTE] + assert request.getAttributeNames().toList() == [DD_CONTEXT_ATTRIBUTE] 0 * response._ where: diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java index 670b4f2ac45..bb95ff99f8f 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/ControllerAdvice.java @@ -1,15 +1,17 @@ package datadog.trace.instrumentation.springweb6; import static datadog.context.Context.root; -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.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +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.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_CONTINUE_SUFFIX; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DD_HANDLER_SPAN_PREFIX_KEY; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; -import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.http.HttpServletRequest; import net.bytebuddy.asm.Advice; @@ -18,15 +20,19 @@ public class ControllerAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope nameResourceAndStartSpan( + public static ContextScope nameResourceAndStartSpan( @Advice.Argument(0) final HttpServletRequest request, @Advice.Argument(2) final Object handler, @Advice.Local("handlerSpanKey") String handlerSpanKey) { handlerSpanKey = ""; // Name the parent span based on the matching pattern - Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - DECORATE.onRequest((AgentSpan) parentSpan, request, request, root()); + Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + DECORATE.onRequest(parentSpan, request, request, root()); + } } if (activeSpan() == null) { @@ -42,21 +48,24 @@ public static AgentScope nameResourceAndStartSpan( handlerKey = handler.getClass().getName(); } handlerSpanKey = DD_HANDLER_SPAN_PREFIX_KEY + handlerKey; - final Object existingSpan = request.getAttribute(handlerSpanKey); - if (existingSpan instanceof AgentSpan) { - return activateSpan((AgentSpan) existingSpan); + + // If the context already exists, return it + final Object existingContext = request.getAttribute(handlerSpanKey); + if (existingContext instanceof Context) { + return ((Context) existingContext).attach(); } final AgentSpan span = startSpan(DECORATE.spanName()).setMeasured(true); DECORATE.afterStart(span); DECORATE.onHandle(span, handler); + request.setAttribute(handlerSpanKey, span); - return activateSpan(span); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.Argument(0) final HttpServletRequest request, @Advice.Thrown final Throwable throwable, @Advice.Local("handlerSpanKey") String handlerSpanKey) { @@ -66,14 +75,14 @@ public static void stopSpan( boolean finish = !Boolean.TRUE.equals( request.getAttribute(handlerSpanKey + DD_HANDLER_SPAN_CONTINUE_SUFFIX)); - final AgentSpan span = scope.span(); + final AgentSpan span = spanFromContext(scope.context()); scope.close(); if (throwable != null) { DECORATE.onError(span, throwable); finish = true; } if (finish) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); span.finish(); } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java index 0b5008f5e00..eac11f69ca2 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/HandlerMappingResourceNameFilter.java @@ -1,9 +1,11 @@ package datadog.trace.instrumentation.springweb6; import static datadog.context.Context.root; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.springweb6.SpringWebHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import jakarta.servlet.FilterChain; import jakarta.servlet.ServletException; @@ -33,18 +35,22 @@ protected void doFilterInternal( final FilterChain filterChain) throws ServletException, IOException { - final Object parentSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); - if (parentSpan instanceof AgentSpan) { - PathMatchingHttpServletRequestWrapper wrappedRequest = - new PathMatchingHttpServletRequestWrapper(request); - try { - if (findMapping(wrappedRequest)) { - // Name the parent span based on the matching pattern - // Let the parent span resource name be set with the attribute set in findMapping. - DECORATE.onRequest((AgentSpan) parentSpan, wrappedRequest, wrappedRequest, root()); + final Object contextObj = request.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan parentSpan = spanFromContext(context); + if (parentSpan != null) { + PathMatchingHttpServletRequestWrapper wrappedRequest = + new PathMatchingHttpServletRequestWrapper(request); + try { + if (findMapping(wrappedRequest)) { + // Name the parent span based on the matching pattern + // Let the parent span resource name be set with the attribute set in findMapping. + DECORATE.onRequest(parentSpan, wrappedRequest, wrappedRequest, root()); + } + } catch (final Exception ignored) { + // mapping.getHandler() threw exception. Ignore } - } catch (final Exception ignored) { - // mapping.getHandler() threw exception. Ignore } } diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java index 44e88d34688..542af56afda 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java +++ b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/RenderAdvice.java @@ -1,9 +1,10 @@ package datadog.trace.instrumentation.springweb6; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; 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 datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.context.ContextScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import net.bytebuddy.asm.Advice; import org.springframework.web.servlet.ModelAndView; @@ -11,19 +12,20 @@ public class RenderAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope onEnter(@Advice.Argument(0) final ModelAndView mv) { + public static ContextScope onEnter(@Advice.Argument(0) final ModelAndView mv) { final AgentSpan span = startSpan(SpringWebHttpServerDecorator.RESPONSE_RENDER); SpringWebHttpServerDecorator.DECORATE_RENDER.afterStart(span); SpringWebHttpServerDecorator.DECORATE_RENDER.onRender(span, mv); - return activateSpan(span); + return getCurrentContext().with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) { + @Advice.Enter final ContextScope scope, @Advice.Thrown final Throwable throwable) { + final AgentSpan span = spanFromContext(scope.context()); SpringWebHttpServerDecorator.DECORATE_RENDER.onError(scope, throwable); - SpringWebHttpServerDecorator.DECORATE_RENDER.beforeFinish(scope); + SpringWebHttpServerDecorator.DECORATE_RENDER.beforeFinish(scope.context()); scope.close(); - scope.span().finish(); + span.finish(); } } diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java index a73af873ffb..ab9bd59e1c5 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java @@ -74,8 +74,9 @@ public static ContextScope beginRequest( DECORATE.afterStart(span); DECORATE.onRequest(span, connection, request, parentContext); - // capture span to be finished by one of the various server response advices - connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span); + // capture context (which contains span) to be finished by one of the various server response + // advices + connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, context); return scope; } @@ -90,10 +91,10 @@ public static final class ServerResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope beginResponse( @Advice.Argument(0) final NHttpServerConnection connection) { - // check and remove span from context so it won't be finished twice - AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); - if (null != span) { - return span.attach(); + // check and remove context so it won't be finished twice + Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); + if (null != context) { + return context.attach(); } return null; } @@ -111,7 +112,7 @@ public static void responseReady( if (null != error) { DECORATE.onError(span, error); } - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } @@ -122,16 +123,19 @@ public static final class ServerErrorResponseAdvice { public static void errorResponse( @Advice.Argument(0) final NHttpServerConnection connection, @Advice.Argument(value = 1, optional = true) final Object error) { - // check and remove span from context so it won't be finished twice - AgentSpan span = (AgentSpan) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); - if (null != span) { - if (error instanceof Throwable) { - DECORATE.onError(span, (Throwable) error); - } else { - span.setError(true); + // check and remove context so it won't be finished twice + Context context = (Context) connection.getContext().removeAttribute(SYNAPSE_SPAN_KEY); + if (null != context) { + AgentSpan span = spanFromContext(context); + if (null != span) { + if (error instanceof Throwable) { + DECORATE.onError(span, (Throwable) error); + } else { + span.setError(true); + } + DECORATE.beforeFinish(context); + span.finish(); } - DECORATE.beforeFinish(span); - span.finish(); } } } diff --git a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java index 6bfa5a5db2e..22720024845 100644 --- a/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3.0/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerWorkerInstrumentation.java @@ -3,6 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.DECORATE; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.SYNAPSE_CONTINUATION_KEY; import static datadog.trace.instrumentation.synapse3.SynapseServerDecorator.SYNAPSE_SPAN_KEY; @@ -12,6 +13,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.ContextScope; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -66,26 +68,32 @@ public static void createWorker(@Advice.Argument(0) final SourceRequest request) public static final class ServerWorkerResponseAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope beginResponse( + public static ContextScope beginResponse( @Advice.FieldValue("request") final SourceRequest request) { AgentScope.Continuation continuation = (AgentScope.Continuation) request.getConnection().getContext().removeAttribute(SYNAPSE_CONTINUATION_KEY); if (null != continuation) { - return continuation.activate(); + // Activate the continuation to get the span, then attach it to current context + AgentScope agentScope = continuation.activate(); + try { + return agentScope.span().attach(); + } finally { + agentScope.close(); + } } return null; } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void responseReady( - @Advice.Enter final AgentScope scope, + @Advice.Enter final ContextScope scope, @Advice.FieldValue("request") final SourceRequest request, @Advice.Thrown final Throwable error) { if (null == scope) { return; } - AgentSpan span = scope.span(); + AgentSpan span = spanFromContext(scope.context()); HttpResponse httpResponse = request.getConnection().getHttpResponse(); if (null != httpResponse) { DECORATE.onResponse(span, httpResponse); @@ -97,7 +105,7 @@ public static void responseReady( // (if there's an ACK response or error we might not get a separate response event) if ((null != httpResponse || null != error) && null != request.getConnection().getContext().removeAttribute(SYNAPSE_SPAN_KEY)) { - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); scope.close(); span.finish(); } else { diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java index 9c292ae4e1d..8bd72786a64 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/RequestInstrumentation.java @@ -1,12 +1,14 @@ package datadog.trace.instrumentation.tomcat; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -79,20 +81,21 @@ public static class RecycleAdvice { public static void stopSpan(@Advice.This final Request req) { Response resp = req.getResponse(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { + if (contextObj instanceof Context) { /** - * This advice will be called for both Request and Response. The span is removed from the + * This advice will be called for both Request and Response. The context is removed from the * request so the advice only applies the first invocation. (So it doesn't matter which is * recycled first.) */ - // value set on the coyote request, so we must remove directly from there. - req.getCoyoteRequest().setAttribute(DD_SPAN_ATTRIBUTE, null); + // values set on the coyote request, so we must remove directly from there. + req.getCoyoteRequest().setAttribute(DD_CONTEXT_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); DECORATE.onResponse(span, resp); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(context); span.finish(); } } diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java index e17ad93f3c0..17234fad51d 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/ResponseInstrumentation.java @@ -1,11 +1,13 @@ package datadog.trace.instrumentation.tomcat; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -56,20 +58,21 @@ public static class RecycleAdvice { public static void stopSpan(@Advice.This final Response resp) { Request req = resp.getRequest(); - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { + if (contextObj instanceof Context) { /** - * This advice will be called for both Request and Response. The span is removed from the + * This advice will be called for both Request and Response. The context is removed from the * request so the advice only applies the first invocation. (So it doesn't matter which is * recycled first.) */ - // value set on the coyote request, so we must remove directly from there. - req.getCoyoteRequest().setAttribute(DD_SPAN_ATTRIBUTE, null); + // values set on the coyote request, so we must remove directly from there. + req.getCoyoteRequest().setAttribute(DD_CONTEXT_ATTRIBUTE, null); - final AgentSpan span = (AgentSpan) spanObj; + final Context context = (Context) contextObj; + final AgentSpan span = spanFromContext(context); DECORATE.onResponse(span, resp); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(context); span.finish(); } } diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index 4534a35b8e5..ba3e5ad848f 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -3,10 +3,9 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.agent.tooling.muzzle.Reference.EXPECTS_NON_STATIC; import static datadog.trace.agent.tooling.muzzle.Reference.EXPECTS_PUBLIC; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DD_PARENT_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.tomcat.TomcatDecorator.DECORATE; @@ -117,10 +116,10 @@ public static class ServiceAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Request req) { - Object existingSpan = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (existingSpan instanceof AgentSpan) { - // Request already gone through initial processing, so just activate the span. - return activateSpan((AgentSpan) existingSpan); + Object existingCtx = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (existingCtx instanceof Context) { + // Request already gone through initial processing, so just attach the context. + return ((Context) existingCtx).attach(); } final Context parentContext = DECORATE.extract(req); @@ -132,7 +131,7 @@ public static ContextScope onService(@Advice.Argument(0) org.apache.coyote.Reque final AgentSpan span = spanFromContext(context); DECORATE.afterStart(span); - req.setAttribute(DD_SPAN_ATTRIBUTE, span); + req.setAttribute(DD_CONTEXT_ATTRIBUTE, context); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId()); return scope; @@ -162,19 +161,23 @@ public static void afterParse( @Advice.Argument(1) Request req, @Advice.Argument(3) Response resp, @Advice.Return(readOnly = false) Boolean ret) { - Object spanObj = req.getAttribute(DD_SPAN_ATTRIBUTE); - if (spanObj instanceof AgentSpan) { - AgentSpan span = (AgentSpan) spanObj; - req.setAttribute(CorrelationIdentifier.getTraceIdKey(), AgentTracer.get().getTraceId(span)); - req.setAttribute(CorrelationIdentifier.getSpanIdKey(), AgentTracer.get().getSpanId(span)); - Object ctxObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE); - Context parentContext = ctxObj instanceof Context ? (Context) ctxObj : getRootContext(); - DECORATE.onRequest(span, req, req, parentContext); - Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); - if (rba != null) { - TomcatBlockingHelper.commitBlockingResponse( - span.getRequestContext().getTraceSegment(), req, resp, rba); - ret = false; // skip pipeline + Object contextObj = req.getAttribute(DD_CONTEXT_ATTRIBUTE); + if (contextObj instanceof Context) { + Context context = (Context) contextObj; + AgentSpan span = spanFromContext(context); + if (span != null) { + req.setAttribute( + CorrelationIdentifier.getTraceIdKey(), AgentTracer.get().getTraceId(span)); + req.setAttribute(CorrelationIdentifier.getSpanIdKey(), AgentTracer.get().getSpanId(span)); + Object ctxObj = req.getAttribute(DD_PARENT_CONTEXT_ATTRIBUTE); + Context parentContext = ctxObj instanceof Context ? (Context) ctxObj : getRootContext(); + DECORATE.onRequest(span, req, req, parentContext); + Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); + if (rba != null) { + TomcatBlockingHelper.commitBlockingResponse( + span.getRequestContext().getTraceSegment(), req, resp, rba); + ret = false; // skip pipeline + } } } } diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java index 670c877ffad..8583b92156c 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-appsec/tomcat-appsec-7.0/src/main/java/datadog/trace/instrumentation/tomcat7/CommitActionInstrumentation.java @@ -1,11 +1,13 @@ package datadog.trace.instrumentation.tomcat7; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.gateway.BlockResponseFunction; @@ -92,7 +94,12 @@ static class ProcessCommitActionAdvice { request.setAttribute(HttpServerDecorator.DD_IGNORE_COMMIT_ATTRIBUTE, Boolean.TRUE); // on async requests AgentTracer.activeSpan() may return null - AgentSpan agentSpan = (AgentSpan) request.getAttribute(HttpServerDecorator.DD_SPAN_ATTRIBUTE); + Object contextObj = request.getAttribute(HttpServerDecorator.DD_CONTEXT_ATTRIBUTE); + if (!(contextObj instanceof Context)) { + return false; + } + Context context = (Context) contextObj; + AgentSpan agentSpan = spanFromContext(context); if (agentSpan == null) { return false; } diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java index 76cb6392f4c..2b57a5b14ae 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ExchangeEndSpanListener.java @@ -1,8 +1,10 @@ package datadog.trace.instrumentation.undertow; +import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromContext; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import io.undertow.server.DefaultResponseListener; @@ -21,7 +23,8 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener return; } - AgentSpan span = continuation.span(); + Context context = continuation.context(); + AgentSpan span = fromContext(context); Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION); if (throwable != null) { @@ -29,7 +32,7 @@ public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener } DECORATE.onResponse(span, exchange); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(context); continuation.cancel(); span.finish(); nextListener.proceed(); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java index 5977c2af5b4..13f126deca2 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java @@ -85,7 +85,7 @@ public static void afterRequestParse( DECORATE.onError(span, throwable); // because we know that a http 400 will be thrown DECORATE.onResponseStatus(span, 400); - DECORATE.beforeFinish(span); + DECORATE.beforeFinish(scope.context()); } finally { if (span != null) { span.finish(); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java index 9a63d2c947c..b97a0a4c3b8 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/ServletInstrumentation.java @@ -3,7 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.SERVLET_REQUEST; @@ -68,7 +68,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_SPAN_ATTRIBUTE, undertowSpan); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, continuation.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java index 58fb7553e76..ab5b64fc3b5 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.2/src/main/java/datadog/trace/instrumentation/undertow/JakartaServletInstrumentation.java @@ -3,7 +3,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH; -import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_CONTEXT_ATTRIBUTE; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DD_UNDERTOW_CONTINUATION; import static datadog.trace.instrumentation.undertow.UndertowDecorator.SERVLET_REQUEST; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -60,7 +60,7 @@ public static void enter( if (continuation != null) { AgentSpan undertowSpan = continuation.span(); ServletRequest request = servletRequestContext.getServletRequest(); - request.setAttribute(DD_SPAN_ATTRIBUTE, undertowSpan); + request.setAttribute(DD_CONTEXT_ATTRIBUTE, continuation.context()); undertowSpan.setSpanName(SERVLET_REQUEST); undertowSpan.setTag(SERVLET_CONTEXT, request.getServletContext().getContextPath()); diff --git a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java index b2ffddc5bd1..5c617ce0275 100644 --- a/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx/vertx-web/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/client/HttpClientRequestBaseInstrumentation.java @@ -2,6 +2,7 @@ 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.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.httpurlconnection.HttpUrlConnectionDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPackagePrivate; @@ -9,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -57,8 +59,9 @@ public static void onExit( @Advice.FieldValue("stream") final HttpClientStream stream, @Advice.Return boolean result) { if (result) { - AgentSpan nettySpan = - stream.connection().channel().attr(AttributeKeys.SPAN_ATTRIBUTE_KEY).get(); + Context storedContext = + stream.connection().channel().attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).get(); + AgentSpan nettySpan = spanFromContext(storedContext); if (nettySpan != null) { try (final AgentScope scope = activateSpan(nettySpan)) { DECORATE.onError(scope, cause);