Support recording of baggage, that would be sent to agent as trace metadata, during tests

Author: mcculls

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy

@@ -39,6 +39,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
 import datadog.trace.bootstrap.instrumentation.api.URIUtils
 import datadog.trace.core.DDSpan
+import datadog.trace.core.Metadata
 import datadog.trace.core.datastreams.StatsGroup
 import datadog.trace.test.util.Flaky
 import groovy.json.JsonOutput
@@ -780,6 +781,15 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
 
   def "test baggage span tags are properly added"() {
     setup:
+    def recordedBaggage = [:]
+    TEST_WRITER.metadataConsumer = { Metadata md ->
+      md.baggage.forEach { k, v ->
+        // record non-internal baggage sent to agent as trace metadata
+        if (!k.startsWith("_dd.")) {
+          recordedBaggage.put(k, v)
+        }
+      }
+    }
     // Use default configuration for TRACE_BAGGAGE_TAG_KEYS (user.id, session.id, account.id)
     def baggageHeader = "user.id=test-user,session.id=test-session,account.id=test-account,language=en"
     def request = request(SUCCESS, 'GET', null)
@@ -799,12 +809,7 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
       trace(spanCount(SUCCESS)) {
         sortSpansByStart()
         // Verify baggage tags are added for default configured keys only
-        serverSpan(it, null, null, 'GET', SUCCESS, [
-          "baggage.user.id": "test-user",
-          "baggage.session.id": "test-session",
-          "baggage.account.id": "test-account"
-          // "baggage.language" should NOT be present since it's not in default config
-        ])
+        serverSpan(it, null, null, 'GET', SUCCESS)
         if (hasHandlerSpan()) {
           handlerSpan(it)
         }
@@ -814,6 +819,12 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
         }
       }
     }
+    recordedBaggage == [
+      "baggage.user.id": "test-user",
+      "baggage.session.id": "test-session",
+      "baggage.account.id": "test-account"
+      // "baggage.language" should NOT be present since it's not in default config
+    ]
 
     and:
     if (isDataStreamsEnabled()) {

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

@@ -5,6 +5,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import datadog.trace.core.DDSpan;
+import datadog.trace.core.MetadataConsumer;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
@@ -25,6 +26,8 @@ public class ListWriter extends CopyOnWriteArrayList<List<DDSpan>> implements Wr
 
   private Filter filter = ACCEPT_ALL;
 
+  private MetadataConsumer metadataConsumer = MetadataConsumer.NO_OP;
+
   public List<DDSpan> firstTrace() {
     return get(0);
   }
@@ -37,18 +40,7 @@ public void write(List<DDSpan> trace) {
     for (DDSpan span : trace) {
       // This is needed to properly do all delayed processing to make this writer even
       // remotely realistic so the test actually test something
-      span.processTagsAndBaggage(
-          metadata -> {
-            // surface injected baggage metadata as span tags so they can be asserted
-            metadata
-                .getBaggage()
-                .forEach(
-                    (k, v) -> {
-                      if (!k.startsWith("_dd.") && span.getTag(k) == null) {
-                        span.setTag(k, v);
-                      }
-                    });
-          });
+      span.processTagsAndBaggage(metadataConsumer);
     }
 
     add(trace);
@@ -127,6 +119,11 @@ public void setFilter(Filter filter) {
     this.filter = filter;
   }
 
+  /** Set a {@link MetadataConsumer} to capture what trace metadata would be sent to the agent. */
+  public void setMetadataConsumer(MetadataConsumer metadataConsumer) {
+    this.metadataConsumer = metadataConsumer;
+  }
+
   private boolean isReported(DDSpan span) {
     for (List<DDSpan> trace : this) {
       for (DDSpan aSpan : trace) {