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, value);
+        }
+      }
+    }
+  }
+
+  public void propagateCustomTags(Map<String, Object> tags) {
+    if (propagatedTagKeys.isEmpty() || tags == null || tags.isEmpty()) {
+      return;
+    }
+    synchronized (tagPropagationLock) {
+      for (String key : propagatedTagKeys) {
+        Object 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, Object> 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, 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, Object> 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, Object> 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, Object> 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.writeObjectMap(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, Object> propagatedTags = Serializer.readObjectMap(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/main/java/datadog/trace/civisibility/ipc/serialization/Serializer.java

@@ -16,6 +16,13 @@
 import java.util.function.Supplier;
 
 public class Serializer {
+  private static final byte NULL = 0;
+  private static final byte STRING = 1;
+  private static final byte BOOLEAN = 2;
+  private static final byte INTEGER = 3;
+  private static final byte LONG = 4;
+  private static final byte FLOAT = 5;
+  private static final byte DOUBLE = 6;
 
   private final ByteArrayOutputStream baos = new ByteArrayOutputStream();
 
@@ -45,6 +52,14 @@ public void write(long l) {
     baos.write((int) l);
   }
 
+  public void write(float f) {
+    write(Float.floatToIntBits(f));
+  }
+
+  public void write(double d) {
+    write(Double.doubleToLongBits(d));
+  }
+
   public void write(String s) {
     if (s == null) {
       write(-1);
@@ -81,6 +96,10 @@ public void write(Map<String, String> m) {
     write(m, Serializer::write, Serializer::write);
   }
 
+  public void writeObjectMap(Map<String, Object> m) {
+    write(m, Serializer::write, Serializer::writeObject);
+  }
+
   public <K, V> void write(
       Map<K, V> m,
       BiConsumer<Serializer, K> keySerializer,
@@ -104,6 +123,32 @@ public void write(BitSet bitSet) {
     }
   }
 
+  private void writeObject(Object value) {
+    if (value == null) {
+      write(NULL);
+    } else if (value instanceof String) {
+      write(STRING);
+      write((String) value);
+    } else if (value instanceof Boolean) {
+      write(BOOLEAN);
+      write((boolean) value);
+    } else if (value instanceof Integer) {
+      write(INTEGER);
+      write((int) value);
+    } else if (value instanceof Long) {
+      write(LONG);
+      write((long) value);
+    } else if (value instanceof Float) {
+      write(FLOAT);
+      write((float) value);
+    } else if (value instanceof Double) {
+      write(DOUBLE);
+      write((double) value);
+    } else {
+      throw new IllegalArgumentException("Unsupported value type: " + value.getClass());
+    }
+  }
+
   public int length() {
     return baos.size();
   }
@@ -133,6 +178,14 @@ public static long readLong(ByteBuffer byteBuffer) {
     return byteBuffer.getLong();
   }
 
+  public static float readFloat(ByteBuffer byteBuffer) {
+    return Float.intBitsToFloat(readInt(byteBuffer));
+  }
+
+  public static double readDouble(ByteBuffer byteBuffer) {
+    return Double.longBitsToDouble(readLong(byteBuffer));
+  }
+
   public static String readString(ByteBuffer byteBuffer) {
     byte[] b = readByteArray(byteBuffer);
     return b != null ? new String(b, StandardCharsets.UTF_8) : null;
@@ -175,6 +228,10 @@ public static Map<String, String> readStringMap(ByteBuffer byteBuffer) {
     return readMap(byteBuffer, Serializer::readString, Serializer::readString);
   }
 
+  public static Map<String, Object> readObjectMap(ByteBuffer byteBuffer) {
+    return readMap(byteBuffer, Serializer::readString, Serializer::readObject);
+  }
+
   public static <K, V> Map<K, V> readMap(
       ByteBuffer byteBuffer,
       Function<ByteBuffer, K> keyDeserializer,
@@ -212,6 +269,28 @@ private static <K, V> Map<K, V> fillMap(
     return m;
   }
 
+  private static Object readObject(ByteBuffer byteBuffer) {
+    byte type = readByte(byteBuffer);
+    switch (type) {
+      case NULL:
+        return null;
+      case STRING:
+        return readString(byteBuffer);
+      case BOOLEAN:
+        return readBoolean(byteBuffer);
+      case INTEGER:
+        return readInt(byteBuffer);
+      case LONG:
+        return readLong(byteBuffer);
+      case FLOAT:
+        return readFloat(byteBuffer);
+      case DOUBLE:
+        return readDouble(byteBuffer);
+      default:
+        throw new IllegalArgumentException("Unsupported value type: " + type);
+    }
+  }
+
   public static BitSet readBitSet(ByteBuffer byteBuffer) {
     byte[] bytes = readByteArray(byteBuffer);
     return bytes != null ? BitSet.valueOf(bytes) : null;