Enable Inferred Proxy Span Support (#9958)

Author: zarirhamza

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

@@ -544,7 +544,7 @@ public Context beforeFinish(Context context) {
     }
 
     // Close Serverless Gateway Inferred Span if any
-    // finishInferredProxySpan(context);
+    finishInferredProxySpan(context);
 
     return super.beforeFinish(context);
   }

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

@@ -5,6 +5,7 @@ import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
+import datadog.trace.api.config.TracerConfig
 import datadog.trace.api.iast.IastContext
 import datadog.trace.api.iast.InstrumentationBridge
 import datadog.trace.api.iast.SourceTypes
@@ -61,6 +62,7 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
   protected void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig('dd.iast.enabled', 'true')
+    injectSysConfig(TracerConfig.TRACE_INFERRED_PROXY_SERVICES_ENABLED, 'true')
   }
 
   @Override
@@ -479,4 +481,121 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
       }
     }
   }
+
+  def "test inferred proxy span is finished"() {
+    setup:
+    def request = request(SUCCESS, "GET", null)
+      .header("x-dd-proxy", "aws-apigateway")
+      .header("x-dd-proxy-request-time-ms", "12345")
+      .header("x-dd-proxy-path", "/success")
+      .header("x-dd-proxy-httpmethod", "GET")
+      .header("x-dd-proxy-domain-name", "api.example.com")
+      .header("x-dd-proxy-stage", "test")
+      .build()
+
+    when:
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == SUCCESS.status
+
+    and:
+    // Verify that inferred proxy span was created and finished
+    // It should appear in the trace as an additional span
+    assertTraces(1) {
+      trace(spanCount(SUCCESS) + 1) {
+        sortSpansByStart()
+        // The inferred proxy span should be the first span (earliest start time)
+        // Verify it exists and was finished (appears in trace)
+        // Operation name is the proxy system name (aws.apigateway), not inferred_proxy
+        span {
+          operationName "aws.apigateway"
+          serviceName "api.example.com"
+          // Resource Name: httpmethod + " " + path
+          resourceName "GET /success"
+          spanType "web"
+          parent()
+          tags {
+            "$Tags.COMPONENT" "aws-apigateway"
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_URL" "api.example.com/success"
+            "$Tags.HTTP_ROUTE" "/success"
+            "stage" "test"
+            "_dd.inferred_span" 1
+            // Standard tags that are automatically added
+            "_dd.agent_psr" Number
+            "_dd.base_service" String
+            "_dd.dsm.enabled" Number
+            "_dd.profiling.ctx" String
+            "_dd.profiling.enabled" Number
+            "_dd.trace_span_attribute_schema" Number
+            "_dd.tracer_host" String
+            "_sample_rate" Number
+            "language" "jvm"
+            "process_id" Number
+            "runtime-id" String
+            "thread.id" Number
+            "thread.name" String
+          }
+        }
+        // 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"
+          operationName operation()
+          resourceName expectedResourceName(SUCCESS, "GET", address)
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          childOfPrevious()
+          tags {
+            "$Tags.COMPONENT" component
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "$Tags.PEER_HOST_IPV4" "127.0.0.1"
+            "$Tags.PEER_PORT" Integer
+            "$Tags.HTTP_CLIENT_IP" "127.0.0.1"
+            "$Tags.HTTP_HOSTNAME" address.host
+            "$Tags.HTTP_URL" String
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_STATUS" SUCCESS.status
+            "$Tags.HTTP_USER_AGENT" String
+            "$Tags.HTTP_ROUTE" "/success"
+            "servlet.context" "/boot-context"
+            "servlet.path" "/success"
+            defaultTags()
+          }
+        }
+        if (hasHandlerSpan()) {
+          // Handler span inherits service name from inferred proxy span parent
+          it.span {
+            serviceName "api.example.com"
+            operationName "spring.handler"
+            resourceName "TestController.success"
+            spanType DDSpanTypes.HTTP_SERVER
+            errored false
+            childOfPrevious()
+            tags {
+              "$Tags.COMPONENT" SpringWebHttpServerDecorator.DECORATE.component()
+              "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+              defaultTags()
+            }
+          }
+        }
+        // Controller span also inherits service name
+        it.span {
+          serviceName "api.example.com"
+          operationName "controller"
+          resourceName "controller"
+          errored false
+          childOfPrevious()
+          tags {
+            defaultTags()
+          }
+        }
+        if (hasResponseSpan(SUCCESS)) {
+          responseSpan(it, SUCCESS)
+        }
+      }
+    }
+  }
 }

internal-api/src/main/java/datadog/trace/api/gateway/InferredProxySpan.java

@@ -2,13 +2,16 @@
 
 import static datadog.context.ContextKey.named;
 import static datadog.trace.api.DDTags.SPAN_TYPE;
+import static datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities.MANUAL_INSTRUMENTATION;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_METHOD;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_ROUTE;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.HTTP_URL;
 
 import datadog.context.Context;
 import datadog.context.ContextKey;
 import datadog.context.ImplicitContextKeyed;
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -69,15 +72,45 @@ public AgentSpanContext start(AgentSpanContext extracted) {
 
     String proxySystem = header(PROXY_SYSTEM);
     String proxy = SUPPORTED_PROXIES.get(proxySystem);
+    String httpMethod = header(PROXY_HTTP_METHOD);
+    String path = header(PROXY_PATH);
+    String domainName = header(PROXY_DOMAIN_NAME);
+
     AgentSpan span = AgentTracer.get().startSpan(INSTRUMENTATION_NAME, proxy, extracted, startTime);
-    span.setServiceName(header(PROXY_DOMAIN_NAME));
+
+    // Service: value of x-dd-proxy-domain-name or global config if not found
+    String serviceName =
+        domainName != null && !domainName.isEmpty() ? domainName : Config.get().getServiceName();
+    span.setServiceName(serviceName);
+
+    // Component: aws-apigateway
     span.setTag(COMPONENT, proxySystem);
+
+    // SpanType: web
     span.setTag(SPAN_TYPE, "web");
-    span.setTag(HTTP_METHOD, header(PROXY_HTTP_METHOD));
-    span.setTag(HTTP_URL, header(PROXY_DOMAIN_NAME) + header(PROXY_PATH));
+
+    // Http.method - value of x-dd-proxy-httpmethod
+    span.setTag(HTTP_METHOD, httpMethod);
+
+    // Http.url - value of x-dd-proxy-domain-name + x-dd-proxy-path
+    span.setTag(HTTP_URL, domainName != null ? domainName + path : path);
+
+    // Http.route - value of x-dd-proxy-path
+    span.setTag(HTTP_ROUTE, path);
+
+    // "stage" - value of x-dd-proxy-stage
     span.setTag("stage", header(STAGE));
+
+    // _dd.inferred_span = 1 (indicates that this is an inferred span)
     span.setTag("_dd.inferred_span", 1);
 
+    // Resource Name: value of x-dd-proxy-httpmethod + " " + value of x-dd-proxy-path
+    // Use MANUAL_INSTRUMENTATION priority to prevent TagInterceptor from overriding
+    String resourceName = httpMethod != null && path != null ? httpMethod + " " + path : null;
+    if (resourceName != null) {
+      span.setResourceName(resourceName, MANUAL_INSTRUMENTATION);
+    }
+
     // Free collected headers
     this.headers.clear();
     // Store inferred span

internal-api/src/test/java/datadog/trace/api/gateway/InferredProxySpanTests.java

@@ -92,4 +92,108 @@ void testStoreAndFromContext() {
 
     assertNull(fromContext(root()), "fromContext on empty context should be null");
   }
+
+  @Test
+  @DisplayName("Invalid start time should return extracted context")
+  void testInvalidStartTime() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "invalid-number");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertTrue(inferredProxySpan.isValid());
+    assertNull(inferredProxySpan.start(null), "Invalid start time should return null");
+  }
+
+  @Test
+  @DisplayName("Service name should fallback to config when domain name is null")
+  void testServiceNameFallbackNull() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    headers.put(InferredProxySpan.PROXY_HTTP_METHOD, "GET");
+    headers.put(InferredProxySpan.PROXY_PATH, "/test");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    // Service name should use Config.get().getServiceName() when domain name is null
+  }
+
+  @Test
+  @DisplayName("Service name should fallback to config when domain name is empty")
+  void testServiceNameFallbackEmpty() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    headers.put(InferredProxySpan.PROXY_DOMAIN_NAME, "");
+    headers.put(InferredProxySpan.PROXY_HTTP_METHOD, "GET");
+    headers.put(InferredProxySpan.PROXY_PATH, "/test");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    // Service name should use Config.get().getServiceName() when domain name is empty
+  }
+
+  @Test
+  @DisplayName("HTTP URL should use path only when domain name is null")
+  void testHttpUrlWithoutDomain() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    headers.put(InferredProxySpan.PROXY_HTTP_METHOD, "GET");
+    headers.put(InferredProxySpan.PROXY_PATH, "/test");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    // HTTP URL should be just the path when domain name is null
+  }
+
+  @Test
+  @DisplayName("Resource name should be null when httpMethod is null")
+  void testResourceNameNullMethod() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    headers.put(InferredProxySpan.PROXY_PATH, "/test");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    // Resource name should be null when httpMethod is null
+  }
+
+  @Test
+  @DisplayName("Resource name should be null when path is null")
+  void testResourceNameNullPath() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+    headers.put(InferredProxySpan.PROXY_HTTP_METHOD, "GET");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    // Resource name should be null when path is null
+  }
+
+  @Test
+  @DisplayName("Finish should handle null span gracefully")
+  void testFinishWithNullSpan() {
+    InferredProxySpan inferredProxySpan = fromHeaders(null);
+    // Should not throw exception when span is null
+    inferredProxySpan.finish();
+    assertFalse(inferredProxySpan.isValid());
+  }
+
+  @Test
+  @DisplayName("Finish should clear span after finishing")
+  void testFinishClearsSpan() {
+    Map<String, String> headers = new HashMap<>();
+    headers.put(PROXY_START_TIME_MS, "12345");
+    headers.put(PROXY_SYSTEM, "aws-apigateway");
+
+    InferredProxySpan inferredProxySpan = fromHeaders(headers);
+    assertNotNull(inferredProxySpan.start(null));
+    inferredProxySpan.finish();
+    // Span should be cleared after finish, so calling finish again should be safe
+    inferredProxySpan.finish();
+  }
 }