Add baggage span tags functionality and tests

Author: rachelyangdog

dd-java-agent/instrumentation/okhttp-3/src/test/groovy/OkHttp3Test.groovy

@@ -78,65 +78,6 @@ abstract class OkHttp3Test extends HttpClientTest {
     method = "GET"
     url = server.address.resolve(path)
   }
-
-  def "baggage span tags are properly added"() {
-    setup:
-    // W3C Baggage header format: key1=value1,key2=value2,key3=value3
-    def baggageHeader = "user.id=bark,session.id=test-sess1,account.id=fetch,language=en"
-    def sentHeaders = [:]
-
-    when:
-    def status
-    // Capture the headers that OkHttp3 sends
-    def interceptor = new Interceptor() {
-        @Override
-        Response intercept(Chain chain) throws IOException {
-          def request = chain.request()
-          // Capture all headers sent
-          request.headers().names().each { headerName ->
-            sentHeaders[headerName] = request.header(headerName)
-          }
-          return chain.proceed(request)
-        }
-      }
-
-    def testClient = new OkHttpClient.Builder()
-      .connectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS)
-      .readTimeout(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS)
-      .writeTimeout(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS)
-      .addInterceptor(interceptor)
-      .build()
-
-    def request = new Request.Builder()
-      .url(server.address.resolve("/success").toURL())
-      .header("baggage", baggageHeader)  // Pass baggage directly in header
-      .get()
-      .build()
-    def response = testClient.newCall(request).execute()
-    status = response.code()
-
-    then:
-    status == 200
-    // Verify baggage header was sent
-    sentHeaders["baggage"] == baggageHeader
-    sentHeaders["baggage"].contains("user.id=bark")
-    sentHeaders["baggage"].contains("session.id=test-sess1")
-    sentHeaders["baggage"].contains("account.id=fetch")
-    sentHeaders["baggage"].contains("language=en")
-
-    // Verify the resulting span has the correct baggage tags (only default configured keys)
-    assertTraces(1) {
-      trace(1) {
-        clientSpan(it, null, "GET", false, false, server.address.resolve("/success"), 200, false, null, false, [
-          // Should have baggage tags for keys in default configuration
-          "baggage.user.id": "bark",
-          "baggage.session.id": "test-sess1",
-          "baggage.account.id": "fetch",
-          // "baggage.language" should NOT be present since it's not in default config
-        ])
-      }
-    }
-  }
 }
 
 @Timeout(5)

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

@@ -785,6 +785,53 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
     'GET'  | null | 'x-datadog-test-request-header' | 'bar' | ['request_header_tag': 'bar']
   }
 
+  def "test baggage span tags are properly added"() {
+    setup:
+    // 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)
+    .header("baggage", baggageHeader)
+    .build()
+    def response = client.newCall(request).execute()
+    if (isDataStreamsEnabled()) {
+      TEST_DATA_STREAMS_WRITER.waitForGroups(1)
+    }
+
+    expect:
+    response.code() == SUCCESS.status
+    response.body().string() == SUCCESS.body
+
+    and:
+    assertTraces(1) {
+      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
+        ])
+        if (hasHandlerSpan()) {
+          handlerSpan(it)
+        }
+        controllerSpan(it)
+        if (hasResponseSpan(SUCCESS)) {
+          responseSpan(it, SUCCESS)
+        }
+      }
+    }
+
+    and:
+    if (isDataStreamsEnabled()) {
+      StatsGroup first = TEST_DATA_STREAMS_WRITER.groups.find { it.parentHash == 0 }
+      verifyAll(first) {
+        edgeTags.containsAll(DSM_EDGE_TAGS)
+        edgeTags.size() == DSM_EDGE_TAGS.size()
+      }
+    }
+  }
+
   @Flaky(value = "https://github.com/DataDog/dd-trace-java/issues/4690", suites = ["MuleHttpServerForkedTest"])
   def "test QUERY_ENCODED_BOTH with response header x-ig-response-header tag mapping"() {
     setup:

dd-trace-core/src/test/groovy/datadog/trace/core/CoreSpanBuilderTest.groovy

@@ -13,7 +13,6 @@ import datadog.trace.api.TagMap
 import datadog.trace.api.gateway.RequestContextSlot
 import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.api.naming.SpanNaming
-import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.bootstrap.instrumentation.api.AgentScope
 import datadog.trace.api.datastreams.NoopPathwayContext
 import datadog.trace.bootstrap.instrumentation.api.TagContext