Implement propagation of custom tags to all test event levels

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/AbstractTestModule.java

@@ -64,7 +64,7 @@ public AbstractTestModule(
     }
 
     span = spanBuilder.start();
-    tagsPropagator = new SpanTagsPropagator(span);
+    tagsPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
 
     span.setSpanType(InternalSpanTypes.TEST_MODULE_END);
     span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_MODULE);

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/AbstractTestSession.java

@@ -99,7 +99,7 @@ public AbstractTestSession(
     }
 
     span = spanBuilder.start();
-    tagPropagator = new SpanTagsPropagator(span);
+    tagPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
 
     span.setSpanType(InternalSpanTypes.TEST_SESSION_END);
     span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_SESSION);

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/SpanTagsPropagator.java

@@ -18,15 +18,23 @@ public class SpanTagsPropagator {
   public static final Consumer<AgentSpan> NOOP_PROPAGATOR = span -> {};
 
   private final AgentSpan parentSpan;
+  private final Collection<String> propagatedTagKeys;
   private final Object tagPropagationLock = new Object();
 
   public SpanTagsPropagator(AgentSpan parentSpan) {
+    this(parentSpan, Collections.emptyList());
+  }
+
+  public SpanTagsPropagator(AgentSpan parentSpan, Collection<String> propagatedTagKeys) {
     this.parentSpan = parentSpan;
+    this.propagatedTagKeys =
+        propagatedTagKeys != null ? propagatedTagKeys : Collections.emptyList();
   }
 
   public void propagateCiVisibilityTags(AgentSpan childSpan) {
     mergeTestFrameworks(getFrameworks(childSpan));
     propagateStatus(childSpan);
+    propagateCustomTags(childSpan);
   }
 
   public void propagateStatus(AgentSpan childSpan) {
@@ -49,6 +57,34 @@ public void propagateTags(AgentSpan childSpan, TagMergeSpec<?>... specs) {
     }
   }
 
+  public void propagateCustomTags(AgentSpan childSpan) {
+    if (propagatedTagKeys.isEmpty()) {
+      return;
+    }
+    synchronized (tagPropagationLock) {
+      for (String key : propagatedTagKeys) {
+        Object value = childSpan.getTag(key);
+        if (value != null) {
+          parentSpan.setTag(key, String.valueOf(value));
+        }
+      }
+    }
+  }
+
+  public void propagateCustomTags(Map<String, String> tags) {
+    if (propagatedTagKeys.isEmpty() || tags == null || tags.isEmpty()) {
+      return;
+    }
+    synchronized (tagPropagationLock) {
+      for (String key : propagatedTagKeys) {
+        String value = tags.get(key);
+        if (value != null) {
+          parentSpan.setTag(key, value);
+        }
+      }
+    }
+  }
+
   private void unsafeMergeTestFrameworks(Collection<TestFramework> childFrameworks) {
     Collection<TestFramework> parentFrameworks = getFrameworks(parentSpan);
     Collection<TestFramework> merged = merge(parentFrameworks, childFrameworks);

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/TestSuiteImpl.java

@@ -112,7 +112,7 @@ public TestSuiteImpl(
     }
 
     span = spanBuilder.start();
-    tagsPropagator = new SpanTagsPropagator(span);
+    tagsPropagator = new SpanTagsPropagator(span, config.getCiVisibilityPropagatedTagKeys());
 
     span.setSpanType(InternalSpanTypes.TEST_SUITE_END);
     span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_SUITE);
@@ -275,6 +275,11 @@ public TestImpl testStart(
         executionResults,
         configurationErrors,
         capabilities,
-        tagsPropagator::propagateStatus);
+        this::propagateTags);
+  }
+
+  private void propagateTags(AgentSpan childSpan) {
+    tagsPropagator.propagateStatus(childSpan);
+    tagsPropagator.propagateCustomTags(childSpan);
   }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/buildsystem/BuildSystemModuleImpl.java

@@ -290,6 +290,7 @@ private SignalResponse onModuleExecutionResultReceived(ModuleExecutionResult res
     testsSkipped.add(result.getTestsSkippedTotal());
 
     tagsPropagator.mergeTestFrameworks(result.getTestFrameworks());
+    tagsPropagator.propagateCustomTags(result.getPropagatedTags());
 
     return AckResponse.INSTANCE;
   }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/buildsystem/ProxyTestModule.java

@@ -31,6 +31,8 @@
 import datadog.trace.civisibility.test.ExecutionResults;
 import datadog.trace.civisibility.test.ExecutionStrategy;
 import java.util.Collection;
+import java.util.Map;
+import java.util.Set;
 import java.util.TreeSet;
 import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nonnull;
@@ -61,6 +63,8 @@ public class ProxyTestModule implements TestFrameworkModule {
   private final LinesResolver linesResolver;
   private final CoverageStore.Factory coverageStoreFactory;
   private final Collection<TestFramework> testFrameworks = ConcurrentHashMap.newKeySet();
+  private final Map<String, String> propagatedTags = new ConcurrentHashMap<>();
+  private final Set<String> propagatedTagKeys;
   private final Collection<LibraryCapability> capabilities;
 
   public ProxyTestModule(
@@ -91,6 +95,7 @@ public ProxyTestModule(
     this.linesResolver = linesResolver;
     this.coverageStoreFactory = coverageStoreFactory;
     this.capabilities = capabilities;
+    this.propagatedTagKeys = config.getCiVisibilityPropagatedTagKeys();
   }
 
   @Override
@@ -180,7 +185,8 @@ private void sendModuleExecutionResult() {
               testManagementEnabled,
               hasFailedTestReplayTests,
               testsSkippedTotal,
-              new TreeSet<>(testFrameworks)));
+              new TreeSet<>(testFrameworks),
+              propagatedTags));
 
     } catch (Exception e) {
       log.error("Error while reporting module execution result", e);
@@ -215,13 +221,24 @@ public TestSuiteImpl testSuiteStart(
         executionResults,
         executionStrategy.getExecutionSettings().getConfigurationErrors(),
         capabilities,
-        this::propagateTestFrameworkData);
+        this::propagateData);
   }
 
-  private void propagateTestFrameworkData(AgentSpan childSpan) {
+  private void propagateData(AgentSpan childSpan) {
     testFrameworks.add(
         new TestFramework(
             (String) childSpan.getTag(Tags.TEST_FRAMEWORK),
             (String) childSpan.getTag(Tags.TEST_FRAMEWORK_VERSION)));
+
+    if (propagatedTagKeys.isEmpty()) {
+      return;
+    }
+
+    for (String key : propagatedTagKeys) {
+      Object value = childSpan.getTag(key);
+      if (value != null) {
+        propagatedTags.put(key, String.valueOf(value));
+      }
+    }
   }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/ipc/ModuleExecutionResult.java

@@ -4,6 +4,8 @@
 import datadog.trace.civisibility.ipc.serialization.Serializer;
 import java.nio.ByteBuffer;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
 
 public class ModuleExecutionResult extends ModuleSignal {
@@ -23,6 +25,7 @@ public class ModuleExecutionResult extends ModuleSignal {
   private final boolean hasFailedTestReplayTests;
   private final long testsSkippedTotal;
   private final Collection<TestFramework> testFrameworks;
+  private final Map<String, String> propagatedTags;
 
   public ModuleExecutionResult(
       DDTraceId sessionId,
@@ -34,7 +37,8 @@ public ModuleExecutionResult(
       boolean testManagementEnabled,
       boolean hasFailedTestReplayTests,
       long testsSkippedTotal,
-      Collection<TestFramework> testFrameworks) {
+      Collection<TestFramework> testFrameworks,
+      Map<String, String> propagatedTags) {
     super(sessionId, moduleId);
     this.coverageEnabled = coverageEnabled;
     this.testSkippingEnabled = testSkippingEnabled;
@@ -44,6 +48,7 @@ public ModuleExecutionResult(
     this.hasFailedTestReplayTests = hasFailedTestReplayTests;
     this.testsSkippedTotal = testsSkippedTotal;
     this.testFrameworks = testFrameworks;
+    this.propagatedTags = propagatedTags != null ? propagatedTags : Collections.emptyMap();
   }
 
   public boolean isCoverageEnabled() {
@@ -78,6 +83,10 @@ public Collection<TestFramework> getTestFrameworks() {
     return testFrameworks;
   }
 
+  public Map<String, String> getPropagatedTags() {
+    return propagatedTags;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -94,7 +103,8 @@ public boolean equals(Object o) {
         && testSkippingEnabled == that.testSkippingEnabled
         && hasFailedTestReplayTests == that.hasFailedTestReplayTests
         && testsSkippedTotal == that.testsSkippedTotal
-        && Objects.equals(testFrameworks, that.testFrameworks);
+        && Objects.equals(testFrameworks, that.testFrameworks)
+        && Objects.equals(propagatedTags, that.propagatedTags);
   }
 
   @Override
@@ -106,7 +116,8 @@ public int hashCode() {
         testSkippingEnabled,
         hasFailedTestReplayTests,
         testsSkippedTotal,
-        testFrameworks);
+        testFrameworks,
+        propagatedTags);
   }
 
   @Override
@@ -161,6 +172,7 @@ public ByteBuffer serialize() {
 
     s.write(testsSkippedTotal);
     s.write(testFrameworks, TestFramework::serialize);
+    s.write(propagatedTags);
 
     return s.flush();
   }
@@ -180,6 +192,7 @@ public static ModuleExecutionResult deserialize(ByteBuffer buffer) {
     long testsSkippedTotal = Serializer.readLong(buffer);
     Collection<TestFramework> testFrameworks =
         Serializer.readList(buffer, TestFramework::deserialize);
+    Map<String, String> propagatedTags = Serializer.readStringMap(buffer);
 
     return new ModuleExecutionResult(
         sessionId,
@@ -191,6 +204,7 @@ public static ModuleExecutionResult deserialize(ByteBuffer buffer) {
         testManagementEnabled,
         hasFailedTestReplayTests,
         testsSkippedTotal,
-        testFrameworks);
+        testFrameworks,
+        propagatedTags);
   }
 }

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/SpanTagsPropagatorTest.groovy

@@ -3,6 +3,7 @@ package datadog.trace.civisibility.domain
 import static datadog.trace.civisibility.domain.SpanTagsPropagator.TagMergeSpec
 
 import datadog.trace.api.civisibility.execution.TestStatus
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.civisibility.ipc.TestFramework
 import datadog.trace.core.DDSpan
@@ -126,6 +127,61 @@ class SpanTagsPropagatorTest extends Specification {
     TagMergeSpec.of("tag", Boolean::logicalOr) | false      | false       | true           | false
   }
 
+  // Mocks AgentSpan (interface) rather than DDSpan because propagateCustomTags writes through
+  // the final DDSpan#setTag(String, String) overload, which Spock cannot intercept on a class mock.
+  def "test custom tag propagation from span: child=#childValue, parent=#parentValue, key=#key, allowlist=#allowlist"() {
+    given:
+    def parentSpan = Mock(AgentSpan)
+    parentSpan.getTag(key) >> parentValue
+
+    def childSpan = Mock(AgentSpan)
+    childSpan.getTag(key) >> childValue
+
+    def propagator = new SpanTagsPropagator(parentSpan, allowlist)
+
+    when:
+    propagator.propagateCustomTags(childSpan)
+
+    then:
+    if (expectedValue != null) {
+      1 * parentSpan.setTag(key, expectedValue)
+    } else {
+      0 * parentSpan.setTag(key, _)
+    }
+
+    where:
+    allowlist                | key                  | childValue | parentValue | expectedValue
+    ["bazel.shard_index"]    | "bazel.shard_index"  | "0"        | null        | "0"
+    ["bazel.shard_index"]    | "bazel.shard_index"  | "1"        | "0"         | "1"  // child overrides parent
+    ["bazel.shard_index"]    | "bazel.shard_index"  | null       | "0"         | null // missing on child, no-op
+    ["bazel.shard_index"]    | "bazel.total_shards" | "2"        | null        | null // not in allowlist
+    []                       | "bazel.shard_index"  | "0"        | null        | null // empty allowlist
+    null                     | "bazel.shard_index"  | "0"        | null        | null // null allowlist
+    ["bazel.shard_index"]    | "bazel.shard_index"  | 0L         | null        | "0"  // non-string child stringified
+    ["bazel.shard_index"]    | "bazel.shard_index"  | true       | null        | "true" // boolean stringified
+  }
+
+  def "test custom tag propagation from map: allowlist=#allowlist, tags=#tags"() {
+    given:
+    def parentSpan = Mock(AgentSpan)
+    def propagator = new SpanTagsPropagator(parentSpan, allowlist)
+
+    when:
+    propagator.propagateCustomTags(tags)
+
+    then:
+    expectedSets * parentSpan.setTag(_, _)
+
+    where:
+    allowlist                                       | tags                                                  | expectedSets
+    ["bazel.shard_index", "bazel.total_shards"]     | ["bazel.shard_index": "0", "bazel.total_shards": "2"] | 2
+    ["bazel.shard_index"]                           | ["bazel.shard_index": "0"]                            | 1
+    ["bazel.shard_index"]                           | ["bazel.shard_index": "0", "bazel.total_shards": "2"] | 1   // only allowlisted keys are copied
+    ["bazel.shard_index"]                           | [:]                                                   | 0   // empty tags
+    []                                              | ["bazel.shard_index": "0"]                            | 0   // empty allowlist
+    null                                            | ["bazel.shard_index": "0"]                            | 0   // null allowlist
+  }
+
   def "test synchronized propagation"() {
     given:
     def parentSpan = Mock(DDSpan)

dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/domain/headless/HeadlessTestSessionTest.groovy

@@ -24,6 +24,7 @@ class HeadlessTestSessionTest extends SpanWriterTest {
     def module = session.testModuleStart("module-name", null)
 
     when:
+    module.setTag("custom.propagated_tag", "value")
     module.end(null)
     session.end(null)
 
@@ -34,10 +35,14 @@ class HeadlessTestSessionTest extends SpanWriterTest {
           spanType DDSpanTypes.TEST_SESSION_END
           tags(false) {
             "$Tags.TEST_TEST_MANAGEMENT_ENABLED" true
+            "custom.propagated_tag" "value"
           }
         }
         span(1) {
           spanType DDSpanTypes.TEST_MODULE_END
+          tags(false) {
+            "custom.propagated_tag" "value"
+          }
         }
       }
     })
@@ -47,13 +52,16 @@ class HeadlessTestSessionTest extends SpanWriterTest {
     def executionSettings = Stub(ExecutionSettings)
     executionSettings.getTestManagementSettings() >> new TestManagementSettings(true, 10)
 
-    def executionStrategy = new ExecutionStrategy(Stub(Config), executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
+    def config = Stub(Config)
+    config.getCiVisibilityPropagatedTagKeys() >> ["custom.propagated_tag"]
+
+    def executionStrategy = new ExecutionStrategy(config, executionSettings, Stub(SourcePathResolver), Stub(LinesResolver))
 
     new HeadlessTestSession(
       "project-name",
       null,
       Provider.UNSUPPORTED,
-      Stub(Config),
+      config,
       Stub(CiVisibilityMetricCollector),
       Stub(TestDecorator),
       Stub(SourcePathResolver),