Do tags and baggage processing in ListWriter

bantonsson · bantonsson · commit d43a099a06e8 · 2023-07-05T15:13:39.000+02:00
diff --git a/dd-java-agent/instrumentation/jsp-2.3/src/test/groovy/JSPInstrumentationBasicTests.groovy b/dd-java-agent/instrumentation/jsp-2.3/src/test/groovy/JSPInstrumentationBasicTests.groovy
@@ -110,7 +110,7 @@ class JSPInstrumentationBasicTests extends JSPTestBase {
             "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
             "$Tags.PEER_HOST_IPV4" "127.0.0.1"
             "$Tags.PEER_PORT" Integer
-            "$Tags.HTTP_URL" "http://localhost:$port/$jspWebappContext/getQuery.jsp"
+            "$Tags.HTTP_URL" "http://localhost:$port/$jspWebappContext/getQuery.jsp?HELLO"
             "$DDTags.HTTP_QUERY" "HELLO"
             "$Tags.HTTP_HOSTNAME" "localhost"
             "$Tags.HTTP_METHOD" "GET"
diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/SpringWebfluxHttpClientBase.groovy b/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/SpringWebfluxHttpClientBase.groovy
@@ -6,6 +6,7 @@ import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator
 import datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator
 import org.springframework.http.HttpMethod
@@ -64,6 +65,11 @@ abstract class SpringWebfluxHttpClientBase extends HttpClientTest implements Tes
     def leafParentId = trace.spanAssertCount.get()
     super.clientSpan(trace, parentSpan, method, renameService, tagQueryString, uri, status, error, exception)
     if (!exception) {
+      def expectedQuery = tagQueryString ? uri.query : null
+      def expectedUrl = URIUtils.buildURL(uri.scheme, uri.host, uri.port, uri.path)
+      if (expectedQuery != null && !expectedQuery.empty) {
+        expectedUrl = "$expectedUrl?$expectedQuery"
+      }
       trace.span {
         childOf(trace.span(leafParentId))
         if (renameService) {
@@ -80,13 +86,13 @@ abstract class SpringWebfluxHttpClientBase extends HttpClientTest implements Tes
           "$Tags.PEER_HOSTNAME" "localhost"
           "$Tags.PEER_PORT" uri.port
           "$Tags.PEER_HOST_IPV4" { it == null || it == "127.0.0.1" } // Optional
-          "$Tags.HTTP_URL" "${uri.resolve(uri.path)}"
+          "$Tags.HTTP_URL" expectedUrl
           "$Tags.HTTP_METHOD" method
           if (status) {
             "$Tags.HTTP_STATUS" status
           }
           if (tagQueryString) {
-            "$DDTags.HTTP_QUERY" uri.query
+            "$DDTags.HTTP_QUERY" expectedQuery
             "$DDTags.HTTP_FRAGMENT" { it == null || it == uri.fragment } // Optional
           }
           if (exception) {
diff --git a/dd-java-agent/instrumentation/spring-webflux-6/src/test/groovy/dd/trace/instrumentation/springwebflux6/client/SpringWebfluxHttpClientBase.groovy b/dd-java-agent/instrumentation/spring-webflux-6/src/test/groovy/dd/trace/instrumentation/springwebflux6/client/SpringWebfluxHttpClientBase.groovy
@@ -7,6 +7,7 @@ import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import datadog.trace.api.naming.SpanNaming
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator
 import datadog.trace.instrumentation.springwebflux6.client.SpringWebfluxHttpClientDecorator
 import org.springframework.http.HttpMethod
@@ -63,6 +64,11 @@ abstract class SpringWebfluxHttpClientBase extends HttpClientTest implements Tes
     def leafParentId = trace.spanAssertCount.get()
     super.clientSpan(trace, parentSpan, method, renameService, tagQueryString, uri, status, error, exception)
     if (!exception) {
+      def expectedQuery = tagQueryString ? uri.query : null
+      def expectedUrl = URIUtils.buildURL(uri.scheme, uri.host, uri.port, uri.path)
+      if (expectedQuery != null && !expectedQuery.empty) {
+        expectedUrl = "$expectedUrl?$expectedQuery"
+      }
       trace.span {
         childOf(trace.span(leafParentId))
         if (renameService) {
@@ -79,13 +85,13 @@ abstract class SpringWebfluxHttpClientBase extends HttpClientTest implements Tes
           "$Tags.PEER_HOSTNAME" "localhost"
           "$Tags.PEER_PORT" uri.port
           "$Tags.PEER_HOST_IPV4" { it == null || it == "127.0.0.1" } // Optional
-          "$Tags.HTTP_URL" "${uri.resolve(uri.path)}"
+          "$Tags.HTTP_URL" expectedUrl
           "$Tags.HTTP_METHOD" method
           if (status) {
             "$Tags.HTTP_STATUS" status
           }
           if (tagQueryString) {
-            "$DDTags.HTTP_QUERY" uri.query
+            "$DDTags.HTTP_QUERY" expectedQuery
             "$DDTags.HTTP_FRAGMENT" { it == null || it == uri.fragment } // Optional
           }
           if (exception) {
diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy
@@ -8,6 +8,7 @@ import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import datadog.trace.api.config.TracerConfig
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.core.DDSpan
 import datadog.trace.core.datastreams.StatsGroup
 import datadog.trace.test.util.Flaky
@@ -748,6 +749,11 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
 
   // parent span must be cast otherwise it breaks debugging classloading (junit loads it early)
   void clientSpan(TraceAssert trace, Object parentSpan, String method = "GET", boolean renameService = false, boolean tagQueryString = false, URI uri = server.address.resolve("/success"), Integer status = 200, boolean error = false, Throwable exception = null, boolean ignorePeer = false) {
+    def expectedQuery = tagQueryString ? uri.query : null
+    def expectedUrl = URIUtils.buildURL(uri.scheme, uri.host, uri.port, uri.path)
+    if (expectedQuery != null && !expectedQuery.empty) {
+      expectedUrl = "$expectedUrl?$expectedQuery"
+    }
     trace.span {
       if (parentSpan == null) {
         parent()
@@ -768,13 +774,13 @@ abstract class HttpClientTest extends VersionedNamingTestBase {
         "$Tags.PEER_HOSTNAME" { it == uri.host || ignorePeer }
         "$Tags.PEER_HOST_IPV4" { it == null || it == "127.0.0.1" || ignorePeer } // Optional
         "$Tags.PEER_PORT" { it == null || it == uri.port || it == proxy.port || it == 443 || ignorePeer }
-        "$Tags.HTTP_URL" "${uri.resolve(uri.path)}" // remove fragment
+        "$Tags.HTTP_URL" expectedUrl
         "$Tags.HTTP_METHOD" method
         if (status) {
           "$Tags.HTTP_STATUS" status
         }
         if (tagQueryString) {
-          "$DDTags.HTTP_QUERY" uri.query
+          "$DDTags.HTTP_QUERY" expectedQuery
           "$DDTags.HTTP_FRAGMENT" { it == null || it == uri.fragment } // Optional
         }
         if ({ isDataStreamsEnabled() }) {
diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy
@@ -182,7 +182,12 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
   String expectedUrl(ServerEndpoint endpoint, URI address) {
     URI url = endpoint.resolve(address)
     def path = Config.get().isHttpServerRawResource() && supportsRaw() ? url.rawPath : url.path
-    return URIUtils.buildURL(url.scheme, url.host, url.port, path)
+    def uri = URIUtils.buildURL(url.scheme, url.host, url.port, path)
+    def qt = expectedQueryTag(endpoint)
+    if (qt && !qt.isEmpty()) {
+      uri = "$uri?$qt"
+    }
+    return uri
   }
 
   Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
diff --git a/dd-trace-core/src/main/java/datadog/trace/civisibility/writer/ddintake/CiTestCycleMapperV1.java b/dd-trace-core/src/main/java/datadog/trace/civisibility/writer/ddintake/CiTestCycleMapperV1.java
@@ -232,7 +232,7 @@ public String endpoint() {
     return "citestcycle/v1";
   }
 
-  private static final class MetaWriter extends MetadataConsumer {
+  private static final class MetaWriter implements MetadataConsumer {
 
     private Writable writable;
 
diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/ListWriter.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/ListWriter.java
@@ -1,13 +1,10 @@
 package datadog.trace.common.writer;
 
-import datadog.trace.api.DDTags;
-import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.core.DDSpan;
+import datadog.trace.core.MetadataConsumer;
 import datadog.trace.core.tagprocessor.PeerServiceCalculator;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -46,18 +43,9 @@ public void write(List<DDSpan> trace) {
       return;
     }
     for (DDSpan span : trace) {
-      // todo: this can be better with span.processTagsAndBaggage(NoopMetadataConsumer.INSTANCE);
-      // we cannot use the full postProcessor chain since it trigger also query obsfuscator and will
-      // make fail a lot of tests
-      // only kicks in peer.service computation (avoid scrambling queries and url)
-      final Map<String, ?> enriched =
-          peerServiceCalculator.processTags(new HashMap<>(span.getTags()));
-      final Object peerServiceSource = enriched.get(DDTags.PEER_SERVICE_SOURCE);
-      // we needed to copy the original tag map since it's unmodifiable
-      if (peerServiceSource != null) {
-        span.setTag(Tags.PEER_SERVICE, enriched.get(Tags.PEER_SERVICE));
-        span.setTag(DDTags.PEER_SERVICE_SOURCE, peerServiceSource);
-      }
+      // This is needed to properly do all delayed processing to make this writer even
+      // remotely realistic so the test actually test something
+      span.processTagsAndBaggage(MetadataConsumer.NO_OP);
     }
     traceCount.incrementAndGet();
     synchronized (latches) {
diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_4.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_4.java
@@ -29,7 +29,7 @@ public TraceMapperV0_4() {
     this(5 << 20);
   }
 
-  private static final class MetaWriter extends MetadataConsumer {
+  private static final class MetaWriter implements MetadataConsumer {
 
     private Writable writable;
     private boolean writeSamplingPriority;
diff --git a/dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_5.java b/dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapperV0_5.java
@@ -176,7 +176,7 @@ private List<ByteBuffer> toList() {
     }
   }
 
-  private final class MetaWriter extends MetadataConsumer {
+  private final class MetaWriter implements MetadataConsumer {
 
     private Writable writable;
     private boolean writeSamplingPriority;
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/MetadataConsumer.java b/dd-trace-core/src/main/java/datadog/trace/core/MetadataConsumer.java
@@ -1,5 +1,11 @@
 package datadog.trace.core;
 
-public abstract class MetadataConsumer {
-  public abstract void accept(Metadata metadata);
+import java.util.function.Consumer;
+
+@FunctionalInterface
+public interface MetadataConsumer extends Consumer<Metadata> {
+
+  MetadataConsumer NO_OP = (md) -> {};
+
+  void accept(Metadata metadata);
 }
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/QueryObfuscator.java b/dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/QueryObfuscator.java
@@ -57,13 +57,13 @@ private String obfuscate(String query) {
   @Override
   public Map<String, Object> processTags(Map<String, Object> unsafeTags) {
     Object query = unsafeTags.get(DDTags.HTTP_QUERY);
-    if (query instanceof String) {
-      query = obfuscate((String) query);
+    if (query instanceof CharSequence) {
+      query = obfuscate(query.toString());
 
       unsafeTags.put(DDTags.HTTP_QUERY, query);
 
       Object url = unsafeTags.get(Tags.HTTP_URL);
-      if (url instanceof String) {
+      if (url instanceof CharSequence) {
         unsafeTags.put(Tags.HTTP_URL, url + "?" + query);
       }
     }
diff --git a/dd-trace-core/src/test/java/datadog/trace/SamplingPriorityMetadataChecker.java b/dd-trace-core/src/test/java/datadog/trace/SamplingPriorityMetadataChecker.java
@@ -3,7 +3,7 @@
 import datadog.trace.core.Metadata;
 import datadog.trace.core.MetadataConsumer;
 
-public class SamplingPriorityMetadataChecker extends MetadataConsumer {
+public class SamplingPriorityMetadataChecker implements MetadataConsumer {
   public volatile boolean hasSamplingPriority;
 
   @Override

Original file line number	Diff line number	Diff line change
`@@ -232,7 +232,7 @@ public String endpoint() {`
`232`	`232`	`return "citestcycle/v1";`
`233`	`233`	`}`
`234`	`234`
`235`		`- private static final class MetaWriter extends MetadataConsumer {`
	`235`	`+ private static final class MetaWriter implements MetadataConsumer {`
`236`	`236`
`237`	`237`	`private Writable writable;`
`238`	`238`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ public TraceMapperV0_4() {`
`29`	`29`	`this(5 << 20);`
`30`	`30`	`}`
`31`	`31`
`32`		`- private static final class MetaWriter extends MetadataConsumer {`
	`32`	`+ private static final class MetaWriter implements MetadataConsumer {`
`33`	`33`
`34`	`34`	`private Writable writable;`
`35`	`35`	`private boolean writeSamplingPriority;`
Original file line number	Diff line number	Diff line change
`@@ -176,7 +176,7 @@ private List<ByteBuffer> toList() {`
`176`	`176`	`}`
`177`	`177`	`}`
`178`	`178`
`179`		`- private final class MetaWriter extends MetadataConsumer {`
	`179`	`+ private final class MetaWriter implements MetadataConsumer {`
`180`	`180`
`181`	`181`	`private Writable writable;`
`182`	`182`	`private boolean writeSamplingPriority;`