Fix AWS API Gateway endpoints correlation HTTP span tags - Inferred Lambda Spans

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

@@ -174,6 +174,14 @@ public Context startSpan(REQUEST_CARRIER carrier, Context parentContext) {
     extracted = startInferredProxySpan(parentContext, extracted);
     AgentSpan span =
         tracer().startSpan(instrumentationName, spanName(), extracted).setMeasured(true);
+    // Register service-entry span with inferred proxy span (if present) so that premature
+    // finish calls from child spans (e.g., Spring MVC handler) are deferred until the
+    // service-entry span finishes (after the response status is known).
+    registerServiceEntrySpanInInferredProxy(parentContext, span);
+    // Reset service name inherited from inferred proxy parent: the inferred span uses the
+    // gateway domain name as service name, but the service-entry span should identify
+    // the application (configured DD_SERVICE), not the upstream gateway.
+    resetServiceNameIfUnderInferredProxy(parentContext, span);
     // Apply RequestBlockingAction if any
     Flow<Void> flow = callIGCallbackRequestHeaders(span, carrier);
     if (flow.getAction() instanceof RequestBlockingAction) {
@@ -193,6 +201,20 @@ protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanCont
     return span.start(extracted);
   }
 
+  private void registerServiceEntrySpanInInferredProxy(
+      Context parentContext, AgentSpan serviceEntrySpan) {
+    InferredProxySpan inferredProxy = InferredProxySpan.fromContext(parentContext);
+    if (inferredProxy != null) {
+      inferredProxy.registerServiceEntrySpan(serviceEntrySpan);
+    }
+  }
+
+  private void resetServiceNameIfUnderInferredProxy(Context parentContext, AgentSpan span) {
+    if (InferredProxySpan.fromContext(parentContext) != null) {
+      span.setServiceName(Config.get().getServiceName());
+    }
+  }
+
   private final DataStreamsTransactionTracker.TransactionSourceReader
       DSM_TRANSACTION_SOURCE_READER =
           (source, headerName) -> {

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java

@@ -18,6 +18,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.gateway.InferredLambdaSpan;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -55,7 +56,7 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".LambdaHandlerDecorator",
+      packageName + ".LambdaHandlerDecorator", "datadog.trace.api.gateway.InferredLambdaSpan",
     };
   }
 
@@ -77,6 +78,9 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class ExtensionCommunicationAdvice {
+    // ThreadLocal to store InferredLambdaSpan for cleanup in exit advice
+    public static final ThreadLocal<InferredLambdaSpan> INFERRED_LAMBDA_SPAN = new ThreadLocal<>();
+
     @OnMethodEnter
     static AgentScope enter(
         @This final Object that,
@@ -88,8 +92,23 @@ static AgentScope enter(
       if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
         return null;
       }
+
       String lambdaRequestId = awsContext.getAwsRequestId();
       AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in, lambdaRequestId);
+
+      // Try to create inferred lambda span if input is an API Gateway event
+      InferredLambdaSpan inferredSpan = InferredLambdaSpan.fromEvent(in);
+      if (inferredSpan.isValid()) {
+        // Start the inferred span and use its context as parent for lambda span
+        AgentSpanContext inferredContext = inferredSpan.start(lambdaContext);
+        if (inferredContext != null && inferredContext != lambdaContext) {
+          lambdaContext = inferredContext;
+          // Store for cleanup in exit
+          INFERRED_LAMBDA_SPAN.set(inferredSpan);
+        }
+      }
+
+      // Create lambda invocation span (may be child of inferred span)
       final AgentSpan span;
       if (null == lambdaContext) {
         span = startSpan(INVOCATION_SPAN_NAME);
@@ -127,6 +146,16 @@ static void exit(
         AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
       } finally {
         scope.close();
+
+        // Finish inferred lambda span if it was created
+        InferredLambdaSpan inferredSpan = INFERRED_LAMBDA_SPAN.get();
+        if (inferredSpan != null) {
+          try {
+            inferredSpan.finish();
+          } finally {
+            INFERRED_LAMBDA_SPAN.remove();
+          }
+        }
       }
     }
   }

dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy

@@ -523,6 +523,8 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
             "$Tags.HTTP_ROUTE" "/success"
             "stage" "test"
             "_dd.inferred_span" 1
+            "$Tags.HTTP_STATUS" SUCCESS.status
+            "$Tags.HTTP_USER_AGENT" String
             // Standard tags that are automatically added
             "_dd.agent_psr" Number
             "_dd.base_service" String
@@ -541,10 +543,8 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
           }
         }
         // Server span should be a child of the inferred proxy span
-        // When there's an inferred proxy span parent, the server span inherits the parent's service name
         span {
-          // Service name is inherited from the inferred proxy span parent
-          serviceName "api.example.com"
+          serviceName expectedServiceName()
           operationName operation()
           resourceName expectedResourceName(SUCCESS, "GET", address)
           spanType DDSpanTypes.HTTP_SERVER
@@ -568,9 +568,8 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
           }
         }
         if (hasHandlerSpan()) {
-          // Handler span inherits service name from inferred proxy span parent
           it.span {
-            serviceName "api.example.com"
+            serviceName expectedServiceName()
             operationName "spring.handler"
             resourceName "TestController.success"
             spanType DDSpanTypes.HTTP_SERVER
@@ -583,9 +582,8 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
             }
           }
         }
-        // Controller span also inherits service name
         it.span {
-          serviceName "api.example.com"
+          serviceName expectedServiceName()
           operationName "controller"
           resourceName "controller"
           errored false

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_4.java

@@ -71,7 +71,10 @@ public void accept(Metadata metadata) {
 
       TagMap tags = metadata.getTags();
 
-      final boolean writeSamplingPriority = firstSpanInTrace || lastSpanInTrace;
+      // Also write on top-level spans so that inferred proxy spans (which may be in the middle
+      // of the serialized list due to phased-finish ordering) always carry the sampling decision.
+      final boolean writeSamplingPriority =
+          firstSpanInTrace || lastSpanInTrace || metadata.topLevel();
       final UTF8BytesString processTags = firstSpanInPayload ? metadata.processTags() : null;
       int metaSize =
           metadata.getBaggage().size()

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_5.java

@@ -202,7 +202,10 @@ MetaWriter forSpan(boolean firstInTrace, boolean lastInTrace, boolean firstInPay
 
     @Override
     public void accept(Metadata metadata) {
-      final boolean writeSamplingPriority = firstSpanInTrace || lastSpanInTrace;
+      // Also write on top-level spans so that inferred proxy spans (which may be in the middle
+      // of the serialized list due to phased-finish ordering) always carry the sampling decision.
+      final boolean writeSamplingPriority =
+          firstSpanInTrace || lastSpanInTrace || metadata.topLevel();
       final UTF8BytesString processTags = firstSpanInPayload ? metadata.processTags() : null;
 
       TagMap tags = metadata.getTags();