Port span link related tests

DataDog · Mar 3, 2025 · 52d9691 · 52d9691
1 parent 955e655
commit 52d9691
Show file tree

Hide file tree

Showing 10 changed files with 87 additions and 54 deletions.
diff --git a/...ent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy b/...ent/instrumentation/armeria-jetty/src/test/jetty11/groovy/ArmeriaJetty11ServerTest.groovy
@@ -119,4 +119,9 @@ class ArmeriaJetty11ServerTest extends HttpServerTest<ArmeriaServer> {
   boolean testSessionId() {
     true
   }
+
+  @Override
+  boolean testWebsockets() {
+    false
+  }
 }
diff --git a/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy b/dd-java-agent/instrumentation/aws-java-s3-2.0/src/test/groovy/S3ClientTest.groovy
@@ -1,7 +1,9 @@
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTraceId
+import datadog.trace.bootstrap.instrumentation.api.SpanAttributes
+import datadog.trace.bootstrap.instrumentation.api.SpanLink
 import datadog.trace.core.tagprocessor.SpanPointersProcessor
-import groovy.json.JsonSlurper
 import org.testcontainers.containers.GenericContainer
 import org.testcontainers.utility.DockerImageName
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
@@ -70,6 +72,14 @@ class S3ClientTest extends AgentTestRunner {
           operationName "aws.http"
           resourceName "S3.PutObject"
           spanType DDSpanTypes.HTTP_CLIENT
+          links {
+            link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS,
+              SpanAttributes.builder()
+              .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND)
+              .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
+              .put("ptr.hash","6d1a2fe194c6579187408f827f942be3")
+              .put("link.kind",SpanPointersProcessor.LINK_KIND).build())
+          }
           tags {
             defaultTags()
             tag "component", "java-aws-sdk"
@@ -87,19 +97,6 @@ class S3ClientTest extends AgentTestRunner {
             tag "peer.port", { it instanceof Integer }
             tag "span.kind", "client"
             tag "aws.requestId", { it != null }
-            tag "_dd.span_links", { it != null }
-
-            // Assert the span links
-            def spanLinks = tags["_dd.span_links"]
-            assert spanLinks != null
-            def links = new JsonSlurper().parseText(spanLinks)
-            assert links.size() == 1
-            def link = links[0]
-            assert link["attributes"] != null
-            assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND
-            assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION
-            assert link["attributes"]["ptr.hash"] == "6d1a2fe194c6579187408f827f942be3"
-            assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND
           }
         }
       }
@@ -134,6 +131,14 @@ class S3ClientTest extends AgentTestRunner {
           operationName "aws.http"
           resourceName "S3.PutObject"
           spanType DDSpanTypes.HTTP_CLIENT
+          links {
+            link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS,
+              SpanAttributes.builder()
+              .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND)
+              .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
+              .put("ptr.hash","6d1a2fe194c6579187408f827f942be3")
+              .put("link.kind",SpanPointersProcessor.LINK_KIND).build())
+          }
           tags {
             defaultTags()
             tag "component", "java-aws-sdk"
@@ -151,7 +156,6 @@ class S3ClientTest extends AgentTestRunner {
             tag "peer.port", { it instanceof Integer }
             tag "span.kind", "client"
             tag "aws.requestId", { it != null }
-            tag "_dd.span_links", { it != null }
           }
         }
       }
@@ -161,6 +165,14 @@ class S3ClientTest extends AgentTestRunner {
           operationName "aws.http"
           resourceName "S3.CopyObject"
           spanType DDSpanTypes.HTTP_CLIENT
+          links {
+            link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS,
+              SpanAttributes.builder()
+              .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND)
+              .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
+              .put("ptr.hash","1542053ce6d393c424b1374bac1fc0c5")
+              .put("link.kind",SpanPointersProcessor.LINK_KIND).build())
+          }
           tags {
             defaultTags()
             tag "component", "java-aws-sdk"
@@ -178,19 +190,6 @@ class S3ClientTest extends AgentTestRunner {
             tag "peer.port", { it instanceof Integer }
             tag "span.kind", "client"
             tag "aws.requestId", { it != null }
-            tag "_dd.span_links", { it != null }
-
-            // Assert the span links
-            def spanLinks = tags["_dd.span_links"]
-            assert spanLinks != null
-            def links = new JsonSlurper().parseText(spanLinks)
-            assert links.size() == 1
-            def link = links[0]
-            assert link["attributes"] != null
-            assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND
-            assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION
-            assert link["attributes"]["ptr.hash"] == "1542053ce6d393c424b1374bac1fc0c5"
-            assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND
           }
         }
       }
@@ -346,6 +345,14 @@ class S3ClientTest extends AgentTestRunner {
           operationName "aws.http"
           resourceName "S3.CompleteMultipartUpload"
           spanType DDSpanTypes.HTTP_CLIENT
+          links {
+            link(DDTraceId.ZERO, (long)0, SpanLink.DEFAULT_FLAGS,
+              SpanAttributes.builder()
+              .put("ptr.kind", SpanPointersProcessor.S3_PTR_KIND)
+              .put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
+              .put("ptr.hash","422412aa6b472a7194f3e24f4b12b4a6")
+              .put("link.kind",SpanPointersProcessor.LINK_KIND).build())
+          }
           tags {
             defaultTags()
             tag "component", "java-aws-sdk"
@@ -364,19 +371,6 @@ class S3ClientTest extends AgentTestRunner {
             tag "span.kind", "client"
             tag "aws.requestId", { it != null }
             tag "http.query.string", { it != null }
-            tag "_dd.span_links", { it != null }
-
-            // Assert the span links
-            def spanLinks = tags["_dd.span_links"]
-            assert spanLinks != null
-            def links = new JsonSlurper().parseText(spanLinks)
-            assert links.size() == 1
-            def link = links[0]
-            assert link["attributes"] != null
-            assert link["attributes"]["ptr.kind"] == SpanPointersProcessor.S3_PTR_KIND
-            assert link["attributes"]["ptr.dir"] == SpanPointersProcessor.DOWN_DIRECTION
-            assert link["attributes"]["ptr.hash"] == "422412aa6b472a7194f3e24f4b12b4a6"
-            assert link["attributes"]["link.kind"] == SpanPointersProcessor.LINK_KIND
           }
         }
       }

diff --git a/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy b/dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy
@@ -31,9 +31,12 @@ class JettyServer implements WebsocketServer {
   JettyServer(ServletContextHandler handler) {
     server.handler = handler
     server.addBean(errorHandler)
-    JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> {
-      container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build())
-    })
+    try {
+      JakartaWebSocketServletContainerInitializer.configure(handler, (servletContext, container) -> {
+        container.addEndpoint(ServerEndpointConfig.Builder.create(WsEndpoint.class, "/websocket").build())
+      })
+    } catch (Throwable ignored) {
+    }
   }
 
   @Override

diff --git a/...nstrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy b/...nstrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy
@@ -311,6 +311,7 @@ class OpenTelemetry14Test extends AgentTestRunner {
     assertTraces(1) {
       trace(1) {
         span {
+          ignoreSpanLinks() // check is done on the content of the tag below
           tags {
             defaultTags()
             "$SPAN_KIND" "$SPAN_KIND_INTERNAL"
@@ -345,6 +346,7 @@ class OpenTelemetry14Test extends AgentTestRunner {
     assertTraces(1) {
       trace(1) {
         span {
+          ignoreSpanLinks() // check is done on the content of the tag below
           tags {
             defaultTags()
             "$SPAN_KIND" "$SPAN_KIND_INTERNAL"
@@ -380,6 +382,7 @@ class OpenTelemetry14Test extends AgentTestRunner {
     assertTraces(1) {
       trace(1) {
         span {
+          ignoreSpanLinks() // check is done on the content of the tag below
           tags {
             defaultTags()
             "$SPAN_KIND" "$SPAN_KIND_INTERNAL"
@@ -421,6 +424,7 @@ class OpenTelemetry14Test extends AgentTestRunner {
     assertTraces(1) {
       trace(1) {
         span {
+          ignoreSpanLinks() // check is done on the content of the tag below
           tags {
             defaultTags()
             "$SPAN_KIND" "$SPAN_KIND_INTERNAL"

diff --git a/...es/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy b/...es/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy
@@ -2,9 +2,11 @@ package datadog.trace.instrumentation.spark
 
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDTags
+import datadog.trace.api.DDTraceId
 import datadog.trace.api.Platform
 import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.api.sampling.SamplingMechanism
+import datadog.trace.core.DDSpan
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.Encoders
 import org.apache.spark.sql.execution.streaming.MemoryStream
@@ -300,8 +302,9 @@ class AbstractSparkStructuredStreamingTest extends AgentTestRunner {
           resourceName "test-query"
           spanType "spark"
           parent()
-          assert span.tags.containsKey(DDTags.SPAN_LINKS)
-          assert span.tags[DDTags.SPAN_LINKS] != null
+          links({
+            link(DDTraceId.from((long)12052652441736835200), (long)-6394091631972716416)
+          })
         }
         span {
           operationName "spark.sql"

diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle b/dd-java-agent/instrumentation/undertow/undertow-2.0/build.gradle
@@ -42,6 +42,8 @@ dependencies {
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-2')
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-3')
   testRuntimeOnly project(':dd-java-agent:instrumentation:servlet:request-5')
+  testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:javax-websocket-1.0')
+  testRuntimeOnly project(':dd-java-agent:instrumentation:websocket:jakarta-websocket-2.0')
 
   latestDepTestImplementation group: 'io.undertow', name: 'undertow-servlet', version: '2.2.+'
   latestDepTestImplementation group: 'io.undertow', name: 'undertow-websockets-jsr', version: '2.2.+'

diff --git a/...va-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy b/...va-agent/instrumentation/undertow/undertow-2.0/src/test/groovy/UndertowServletTest.groovy
@@ -323,6 +323,10 @@ abstract class UndertowServletTest extends HttpServerTest<Undertow> {
 }
 
 class UndertowServletV0Test extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 {
+  @Override
+  boolean testWebsockets() {
+    false
+  }
 }
 
 class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implements TestingGenericHttpNamingConventions.ServerV0 {
@@ -331,6 +335,12 @@ class UndertowServletNoHttpRouteForkedTest extends UndertowServletTest implement
     false
   }
 
+  @Override
+  boolean testWebsockets() {
+    // only test websocket in this test class not to have the same done too many times
+    true
+  }
+
   @Override
   protected void configurePreAgent() {
     super.configurePreAgent()

diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/LinksAssert.groovy
@@ -1,6 +1,6 @@
 package datadog.trace.agent.test.asserts
 
-
+import datadog.trace.api.DDTraceId
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink
 import datadog.trace.bootstrap.instrumentation.api.SpanAttributes
 import datadog.trace.bootstrap.instrumentation.api.SpanLink
@@ -31,9 +31,12 @@ class LinksAssert {
   }
 
   def link(DDSpan linked, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') {
+    link(linked.traceId, linked.spanId, flags, attributes, traceState)
+  }
+
+  def link(DDTraceId traceId, def spanId, byte flags = SpanLink.DEFAULT_FLAGS, SpanAttributes attributes = SpanAttributes.EMPTY, String traceState = '') {
     def found = links.find {
-      it.spanId() == linked.spanId &&
-        it.traceId() == linked.traceId
+      it.spanId() == spanId && it.traceId() == traceId
     }
     assert found != null
     assert found.traceFlags() == flags

diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/SpanAssert.groovy
@@ -11,9 +11,12 @@ import groovy.transform.stc.SimpleType
 
 import java.util.regex.Pattern
 
+import static datadog.trace.agent.test.asserts.TagsAssert.assertTags
+
 class SpanAssert {
   private final DDSpan span
   private final DDSpan previous
+  private boolean checkLinks = true
   private final checked = [:]
 
   private SpanAssert(span, DDSpan previous) {
@@ -168,8 +171,10 @@ class SpanAssert {
     if (!checked.errored) {
       errored(false)
     }
-    if (!checked.links) {
-      assert span.tags['_dd.span_links'] == null
+    if (checkLinks) {
+      if (!checked.links) {
+        assert span.tags['_dd.span_links'] == null
+      }
     }
     hasServiceName()
   }
@@ -185,6 +190,10 @@ class SpanAssert {
     assertTags(span, spec, checkAllTags)
   }
 
+  void ignoreSpanLinks() {
+    this.checkLinks = false
+  }
+
   void links(@ClosureParams(value = SimpleType, options = ['datadog.trace.agent.test.asserts.LinksAssert'])
     @DelegatesTo(value = LinksAssert, strategy = Closure.DELEGATE_FIRST) Closure spec) {
     checked.links = true

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/SpanPointersProcessor.java
@@ -25,9 +25,9 @@ public class SpanPointersProcessor implements TagsPostProcessor {
 
   // The pointer direction will always be down. The serverless agent handles cases where the
   // direction is up.
-  static final String DOWN_DIRECTION = "d";
-  static final String S3_PTR_KIND = "aws.s3.object";
-  static final String LINK_KIND = "span-pointer";
+  public static final String DOWN_DIRECTION = "d";
+  public static final String S3_PTR_KIND = "aws.s3.object";
+  public static final String LINK_KIND = "span-pointer";
 
   @Override
   public Map<String, Object> processTags(