HttpServerDecorator.onRequest expects to be passed the parent context containing the extracted trace data (#9513)

Author: mcculls

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

@@ -138,16 +138,17 @@ public Context extract(REQUEST_CARRIER carrier) {
    * Starts a span.
    *
    * @param carrier The request carrier.
-   * @param context The parent context of the span to create.
+   * @param parentContext The parent context of the span to create.
    * @return A new context bundling the span, child of the given parent context.
    */
-  public Context startSpan(REQUEST_CARRIER carrier, Context context) {
+  public Context startSpan(REQUEST_CARRIER carrier, Context parentContext) {
     String[] instrumentationNames = instrumentationNames();
     String instrumentationName =
         instrumentationNames != null && instrumentationNames.length > 0
             ? instrumentationNames[0]
             : DEFAULT_INSTRUMENTATION_NAME;
-    AgentSpanContext.Extracted extracted = callIGCallbackStart(getExtractedSpanContext(context));
+    AgentSpanContext.Extracted extracted =
+        callIGCallbackStart(getExtractedSpanContext(parentContext));
     AgentSpan span =
         tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
     Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
@@ -158,14 +159,14 @@ public Context startSpan(REQUEST_CARRIER carrier, Context context) {
     if (null != carrier && null != getter) {
       tracer().getDataStreamsMonitoring().setCheckpoint(span, forHttpServer());
     }
-    return context.with(span);
+    return parentContext.with(span);
   }
 
   public AgentSpan onRequest(
       final AgentSpan span,
       final CONNECTION connection,
       final REQUEST request,
-      final Context context) {
+      final Context parentContext) {
     Config config = Config.get();
 
     if (APPSEC_ACTIVE) {
@@ -183,7 +184,7 @@ public AgentSpan onRequest(
       }
     }
 
-    AgentSpanContext.Extracted extracted = getExtractedSpanContext(context);
+    AgentSpanContext.Extracted extracted = getExtractedSpanContext(parentContext);
     boolean clientIpResolverEnabled =
         config.isClientIpEnabled() || traceClientIpResolverEnabled && APPSEC_ACTIVE;
     if (extracted != null) {
@@ -316,9 +317,17 @@ public AgentSpan onRequest(
     return span;
   }
 
-  protected static AgentSpanContext.Extracted getExtractedSpanContext(Context context) {
-    AgentSpan extractedSpan = AgentSpan.fromContext(context);
-    return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
+  protected static AgentSpanContext.Extracted getExtractedSpanContext(Context parentContext) {
+    AgentSpan extractedSpan = AgentSpan.fromContext(parentContext);
+    if (extractedSpan != null) {
+      AgentSpanContext extractedSpanContext = extractedSpan.context();
+      if (extractedSpanContext instanceof AgentSpanContext.Extracted) {
+        return (AgentSpanContext.Extracted) extractedSpanContext;
+      } else {
+        log.warn("Expected AgentSpanContext.Extracted but found {}", extractedSpanContext);
+      }
+    }
+    return null;
   }
 
   protected BlockResponseFunction createBlockResponseFunction(

dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayHttpServerDecorator.java

@@ -87,8 +87,8 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final Request<?> connection,
       final Request<?> request,
-      final Context context) {
-    super.onRequest(span, connection, request, context);
+      final Context parentContext) {
+    super.onRequest(span, connection, request, parentContext);
     if (request != null) {
       // more about routes here:
       // https://github.com/playframework/playframework/blob/master/documentation/manual/releases/release26/migration26/Migration26.md#router-tags-are-now-attributes

dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayHttpServerDecorator.java

@@ -87,8 +87,8 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final Request<?> connection,
       final Request<?> request,
-      final Context context) {
-    super.onRequest(span, connection, request, context);
+      final Context parentContext) {
+    super.onRequest(span, connection, request, parentContext);
     if (request != null) {
       // more about routes here:
       // https://github.com/playframework/playframework/blob/master/documentation/manual/releases/release26/migration26/Migration26.md#router-tags-are-now-attributes

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayHttpServerDecorator.java

@@ -130,8 +130,8 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final Request<?> connection,
       final Request<?> request,
-      final Context context) {
-    super.onRequest(span, connection, request, context);
+      final Context parentContext) {
+    super.onRequest(span, connection, request, parentContext);
     if (request != null) {
       // more about routes here:
       // https://github.com/playframework/playframework/blob/master/documentation/manual/releases/release26/migration26/Migration26.md

dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Decorator.java

@@ -79,12 +79,12 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final HttpServletRequest connection,
       final HttpServletRequest request,
-      final Context context) {
+      final Context parentContext) {
     assert span != null;
     if (request != null) {
       span.setTag("servlet.context", request.getContextPath());
       span.setTag("servlet.path", request.getServletPath());
     }
-    return super.onRequest(span, connection, request, context);
+    return super.onRequest(span, connection, request, parentContext);
   }
 }

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Decorator.java

@@ -82,7 +82,7 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final HttpServletRequest connection,
       final HttpServletRequest request,
-      final Context context) {
+      final Context parentContext) {
     assert span != null;
     ClassloaderConfigurationOverrides.maybeEnrichSpan(span);
     if (request != null) {
@@ -97,7 +97,7 @@ public AgentSpan onRequest(
       request.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath);
       request.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath);
     }
-    return super.onRequest(span, connection, request, context);
+    return super.onRequest(span, connection, request, parentContext);
   }
 
   @Override

dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/SpringWebHttpServerDecorator.java

@@ -97,7 +97,7 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final HttpServletRequest connection,
       final HttpServletRequest request,
-      final Context context) {
+      final Context parentContext) {
     if (request != null) {
       final String method = request.getMethod();
       final Object bestMatchingPattern =

dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/SpringWebHttpServerDecorator.java

@@ -101,7 +101,7 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final HttpServletRequest connection,
       final HttpServletRequest request,
-      final Context context) {
+      final Context parentContext) {
     // FIXME: adding a filter to avoid resource name to be overridden on redirect and forwards.
     // Remove myself when jakarta.servlet will be available
     if (request != null && request.getAttribute(DD_FILTERED_SPRING_ROUTE_ALREADY_APPLIED) == null) {

dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java

@@ -72,7 +72,7 @@ public static ContextScope beginRequest(
       ContextScope scope = context.attach();
       AgentSpan span = spanFromContext(context);
       DECORATE.afterStart(span);
-      DECORATE.onRequest(span, connection, request, context);
+      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);

dd-java-agent/instrumentation/tomcat-5.5-common/src/main/java/datadog/trace/instrumentation/tomcat/TomcatDecorator.java

@@ -94,7 +94,7 @@ public AgentSpan onRequest(
       final AgentSpan span,
       final Request connection,
       final Request request,
-      final Context context) {
+      final Context parentContext) {
     if (request != null) {
       String contextPath = request.getContextPath();
       String servletPath = request.getServletPath();
@@ -112,7 +112,7 @@ public AgentSpan onRequest(
       request.setAttribute(DD_CONTEXT_PATH_ATTRIBUTE, contextPath);
       request.setAttribute(DD_SERVLET_PATH_ATTRIBUTE, servletPath);
     }
-    return super.onRequest(span, connection, request, context);
+    return super.onRequest(span, connection, request, parentContext);
   }
 
   @Override