Merge pull request #5546 from DataDog/rgs/jetty-scope-activation

richardstartin · web-flow · commit a82bdedfd5a8 · 2023-07-07T09:41:00.000+01:00
activate spans before calling decorators
diff --git a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java
@@ -27,12 +27,12 @@ public static AgentScope onEnter(
 
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(req);
       span = DECORATE.startSpan(req, extractedContext);
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
       span.setMeasured(true);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, req, req, extractedContext);
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
       req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java
@@ -154,11 +154,11 @@ public static AgentScope onEnter(
 
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(req);
       span = DECORATE.startSpan(req, extractedContext);
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, req, req, extractedContext);
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
       req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java
@@ -155,11 +155,11 @@ public static AgentScope onEnter(
 
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(req);
       span = DECORATE.startSpan(req, extractedContext);
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, req, req, extractedContext);
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
       req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java
@@ -162,11 +162,11 @@ public static AgentScope onEnter(
 
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(req);
       span = DECORATE.startSpan(req, extractedContext);
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, req, req, extractedContext);
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
       req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
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
@@ -89,11 +89,11 @@ public static class HandleRequestAdvice {
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(request);
       request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
       final AgentSpan span = DECORATE.startSpan(request, extractedContext);
+      scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
 
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, request, request, extractedContext);
-      scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);
       request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
diff --git a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java
@@ -31,10 +31,10 @@ public static AgentScope onEnter(@Advice.Argument(0) final Request req) {
       span = startSpan(PLAY_REQUEST);
       span.setMeasured(true);
     }
-    DECORATE.afterStart(span);
 
     final AgentScope scope = activateSpan(span);
     scope.setAsyncPropagation(true);
+    DECORATE.afterStart(span);
     return scope;
   }
 
diff --git a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java
@@ -37,10 +37,10 @@ public static AgentScope onEnter(@Advice.Argument(value = 0, readOnly = false) R
       span = startSpan(PLAY_REQUEST);
       span.setMeasured(true);
     }
-    DECORATE.afterStart(span);
 
     final AgentScope scope = activateSpan(span);
     scope.setAsyncPropagation(true);
+    DECORATE.afterStart(span);
 
     req = RequestHelper.withTag(req, "_dd_HasPlayRequestSpan", "true");
 
diff --git a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java
@@ -35,11 +35,10 @@ public static AgentScope onEnter(@Advice.Argument(value = 0, readOnly = false) R
       // Do not extract the context.
       span = startSpan(PLAY_REQUEST);
     }
-    span.setMeasured(true);
-    DECORATE.afterStart(span);
-
     final AgentScope scope = activateSpan(span);
     scope.setAsyncPropagation(true);
+    span.setMeasured(true);
+    DECORATE.afterStart(span);
 
     req = req.addAttr(HasPlayRequestSpan.KEY, HasPlayRequestSpan.INSTANCE);
 
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
@@ -54,11 +54,12 @@ public static class RestletHandleAdvice {
     public static AgentScope beginRequest(@Advice.Argument(0) final HttpExchange exchange) {
       AgentSpan.Context.Extracted context = DECORATE.extract(exchange);
       AgentSpan span = DECORATE.startSpan(exchange, context);
+      AgentScope scope = activateSpan(span);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, exchange, exchange, context);
       DECORATE.onPeerConnection(span, exchange.getRemoteAddress());
 
-      return activateSpan(span);
+      return scope;
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java
@@ -50,11 +50,10 @@ public static boolean onEnter(
 
     final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(httpServletRequest);
     final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedContext);
-    DECORATE.afterStart(span);
-    DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedContext);
-
     scope = activateSpan(span);
     scope.setAsyncPropagation(true);
+    DECORATE.afterStart(span);
+    DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedContext);
 
     httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
     httpServletRequest.setAttribute(
diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java
@@ -67,13 +67,12 @@ public static boolean onEnter(
 
     final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(httpServletRequest);
     final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedContext);
+    scope = activateSpan(span);
+    scope.setAsyncPropagation(true);
 
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedContext);
 
-    scope = activateSpan(span);
-    scope.setAsyncPropagation(true);
-
     httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
     httpServletRequest.setAttribute(
         CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
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
@@ -26,11 +26,11 @@ public static AgentScope enter(@Advice.Argument(value = 1, readOnly = false) Req
       extractedContext = null;
       span = startSpan(DECORATE.spanName());
     }
+    final AgentScope scope = activateSpan(span);
+    scope.setAsyncPropagation(true);
 
     DECORATE.afterStart(span);
 
-    final AgentScope scope = activateSpan(span);
-    scope.setAsyncPropagation(true);
     ctx = SprayHelper.wrapRequestContext(ctx, scope.span(), extractedContext);
     return scope;
   }
diff --git a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java
@@ -75,13 +75,14 @@ public static AgentScope beginRequest(
         span = startSpan(DECORATE.spanName());
         span.setMeasured(true);
       }
+      AgentScope scope = activateSpan(span);
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, connection, request, extractedContext);
 
       // capture span to be finished by one of the various server response advices
       connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span);
 
-      return activateSpan(span);
+      return scope;
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java
@@ -124,11 +124,11 @@ public static AgentScope onService(@Advice.Argument(0) org.apache.coyote.Request
       req.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext);
 
       final AgentSpan span = DECORATE.startSpan(req, extractedContext);
+      final AgentScope scope = activateSpan(span);
+      scope.setAsyncPropagation(true);
       // This span is finished when Request.recycle() is called by RequestInstrumentation.
       DECORATE.afterStart(span);
 
-      final AgentScope scope = activateSpan(span);
-      scope.setAsyncPropagation(true);
       req.setAttribute(DD_SPAN_ATTRIBUTE, span);
       req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
       req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java
@@ -90,11 +90,10 @@ public static void onEnter(
 
       final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(exchange);
       final AgentSpan span = DECORATE.startSpan(exchange, extractedContext).setMeasured(true);
-      DECORATE.afterStart(span);
-      DECORATE.onRequest(span, exchange, exchange, extractedContext);
-
       scope = activateSpan(span);
       scope.setAsyncPropagation(true);
+      DECORATE.afterStart(span);
+      DECORATE.onRequest(span, exchange, exchange, extractedContext);
 
       exchange.putAttachment(DD_UNDERTOW_CONTINUATION, scope.capture());
 
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
@@ -76,9 +76,9 @@ public static void afterRequestParse(
         } else {
           final AgentSpan.Context.Extracted extractedContext = DECORATE.extract(exchange);
           span = DECORATE.startSpan(exchange, extractedContext).setMeasured(true);
+          scope = activateSpan(span);
           DECORATE.afterStart(span);
           DECORATE.onRequest(span, exchange, exchange, extractedContext);
-          scope = activateSpan(span);
         }
         DECORATE.onError(span, throwable);
         // because we know that a http 400 will be thrown

Original file line number	Diff line number	Diff line change
`@@ -31,10 +31,10 @@ public static AgentScope onEnter(@Advice.Argument(0) final Request req) {`
`31`	`31`	`span = startSpan(PLAY_REQUEST);`
`32`	`32`	`span.setMeasured(true);`
`33`	`33`	`}`
`34`		`- DECORATE.afterStart(span);`
`35`	`34`
`36`	`35`	`final AgentScope scope = activateSpan(span);`
`37`	`36`	`scope.setAsyncPropagation(true);`
	`37`	`+ DECORATE.afterStart(span);`
`38`	`38`	`return scope;`
`39`	`39`	`}`
`40`	`40`
Original file line number	Diff line number	Diff line change
`@@ -37,10 +37,10 @@ public static AgentScope onEnter(@Advice.Argument(value = 0, readOnly = false) R`
`37`	`37`	`span = startSpan(PLAY_REQUEST);`
`38`	`38`	`span.setMeasured(true);`
`39`	`39`	`}`
`40`		`- DECORATE.afterStart(span);`
`41`	`40`
`42`	`41`	`final AgentScope scope = activateSpan(span);`
`43`	`42`	`scope.setAsyncPropagation(true);`
	`43`	`+ DECORATE.afterStart(span);`
`44`	`44`
`45`	`45`	`req = RequestHelper.withTag(req, "_dd_HasPlayRequestSpan", "true");`
`46`	`46`