Improve tag propagation between test event levels

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

@@ -30,6 +30,7 @@ public abstract class AbstractTestModule {
   protected final Codeowners codeowners;
   protected final LinesResolver linesResolver;
   private final Consumer<AgentSpan> onSpanFinish;
+  protected final SpanTagsPropagator tagsPropagator;
 
   public AbstractTestModule(
       AgentSpanContext sessionSpanContext,
@@ -63,6 +64,7 @@ public AbstractTestModule(
     }
 
     span = spanBuilder.start();
+    tagsPropagator = new SpanTagsPropagator(span);
 
     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

@@ -46,6 +46,7 @@ public abstract class AbstractTestSession {
   protected final SourcePathResolver sourcePathResolver;
   protected final Codeowners codeowners;
   protected final LinesResolver linesResolver;
+  protected final SpanTagsPropagator tagPropagator;
 
   public AbstractTestSession(
       String projectName,
@@ -93,6 +94,7 @@ public AbstractTestSession(
     }
 
     span = spanBuilder.start();
+    tagPropagator = new SpanTagsPropagator(span);
 
     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/utils/SpanUtils.java → dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/domain/SpanTagsPropagator.java

@@ -1,4 +1,4 @@
-package datadog.trace.civisibility.utils;
+package datadog.trace.civisibility.domain;
 
 import datadog.trace.api.civisibility.execution.TestStatus;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -14,28 +14,51 @@
 import java.util.function.BinaryOperator;
 import java.util.function.Consumer;
 
-public class SpanUtils {
-  public static final Consumer<AgentSpan> DO_NOT_PROPAGATE_CI_VISIBILITY_TAGS = span -> {};
+public class SpanTagsPropagator {
+  public static final Consumer<AgentSpan> NOOP_PROPAGATOR = span -> {};
 
-  public static Consumer<AgentSpan> propagateCiVisibilityTagsTo(AgentSpan parentSpan) {
-    return childSpan -> propagateCiVisibilityTags(parentSpan, childSpan);
+  private final AgentSpan parentSpan;
+  private final Object tagPropagationLock = new Object();
+
+  public SpanTagsPropagator(AgentSpan parentSpan) {
+    this.parentSpan = parentSpan;
+  }
+
+  public void propagateCiVisibilityTags(AgentSpan childSpan) {
+    mergeTestFrameworks(getFrameworks(childSpan));
+    propagateStatus(childSpan);
+  }
+
+  public void propagateStatus(AgentSpan childSpan) {
+    synchronized (tagPropagationLock) {
+      unsafePropagateStatus(childSpan);
+    }
+  }
+
+  public void mergeTestFrameworks(Collection<TestFramework> testFrameworks) {
+    synchronized (tagPropagationLock) {
+      unsafeMergeTestFrameworks(testFrameworks);
+    }
   }
 
-  public static void propagateCiVisibilityTags(AgentSpan parentSpan, AgentSpan childSpan) {
-    mergeTestFrameworks(parentSpan, getFrameworks(childSpan));
-    propagateStatus(parentSpan, childSpan);
+  public void propagateTags(AgentSpan childSpan, TagMergeSpec<?>... specs) {
+    synchronized (tagPropagationLock) {
+      for (TagMergeSpec<?> spec : specs) {
+        unsafePropagateTag(childSpan, spec);
+      }
+    }
   }
 
-  public static void mergeTestFrameworks(AgentSpan span, Collection<TestFramework> testFrameworks) {
-    Collection<TestFramework> spanFrameworks = getFrameworks(span);
-    Collection<TestFramework> merged = merge(spanFrameworks, testFrameworks);
-    setFrameworks(span, merged);
+  private void unsafeMergeTestFrameworks(Collection<TestFramework> childFrameworks) {
+    Collection<TestFramework> parentFrameworks = getFrameworks(parentSpan);
+    Collection<TestFramework> merged = merge(parentFrameworks, childFrameworks);
+    setFrameworks(merged);
   }
 
-  private static Collection<TestFramework> getFrameworks(AgentSpan span) {
+  static Collection<TestFramework> getFrameworks(AgentSpan span) {
     Object nameTag = span.getTag(Tags.TEST_FRAMEWORK);
     Object versionTag = span.getTag(Tags.TEST_FRAMEWORK_VERSION);
-    if (nameTag == null && versionTag == null) {
+    if (nameTag == null) {
       return Collections.emptyList();
     }
 
@@ -45,9 +68,11 @@ private static Collection<TestFramework> getFrameworks(AgentSpan span) {
 
     } else if (nameTag instanceof Collection) {
       Iterator<String> names = ((Collection<String>) nameTag).iterator();
-      Iterator<String> versions = ((Collection<String>) versionTag).iterator();
+      Iterator<String> versions =
+          versionTag != null ? ((Collection<String>) versionTag).iterator() : null;
       while (names.hasNext()) {
-        frameworks.add(new TestFramework(names.next(), versions.next()));
+        String version = (versions != null && versions.hasNext()) ? versions.next() : null;
+        frameworks.add(new TestFramework(names.next(), version));
       }
 
     } else {
@@ -79,7 +104,7 @@ private static Collection<TestFramework> merge(
     return merged;
   }
 
-  private static void setFrameworks(AgentSpan span, Collection<TestFramework> frameworks) {
+  private void setFrameworks(Collection<TestFramework> frameworks) {
     if (frameworks.isEmpty()) {
       return;
     }
@@ -90,7 +115,7 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
       if (framework.getVersion() != null) {
         tags.put(Tags.TEST_FRAMEWORK_VERSION, framework.getVersion());
       }
-      span.setAllTags(tags);
+      parentSpan.setAllTags(tags);
       return;
     }
     Collection<String> names = new ArrayList<>(frameworks.size());
@@ -102,10 +127,10 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
     Map<String, Collection<String>> tags = new HashMap<>();
     tags.put(Tags.TEST_FRAMEWORK, names);
     tags.put(Tags.TEST_FRAMEWORK_VERSION, versions);
-    span.setAllTags(tags);
+    parentSpan.setAllTags(tags);
   }
 
-  private static void propagateStatus(AgentSpan parentSpan, AgentSpan childSpan) {
+  private void unsafePropagateStatus(AgentSpan childSpan) {
     TestStatus childStatus = (TestStatus) childSpan.getTag(Tags.TEST_STATUS);
     if (childStatus == null) {
       return;
@@ -131,25 +156,32 @@ private static void propagateStatus(AgentSpan parentSpan, AgentSpan childSpan) {
     }
   }
 
-  public static void propagateTags(AgentSpan parentSpan, AgentSpan childSpan, String... tagNames) {
-    for (String tagName : tagNames) {
-      parentSpan.setTag(tagName, childSpan.getTag(tagName));
+  public static class TagMergeSpec<T> {
+    private final String tagKey;
+    private final BinaryOperator<T> mergeFunction;
+
+    TagMergeSpec(String tagKey, BinaryOperator<T> mergeFunction) {
+      this.tagKey = tagKey;
+      this.mergeFunction = mergeFunction;
+    }
+
+    public static <T> TagMergeSpec<T> of(String key, BinaryOperator<T> mergeFunction) {
+      return new TagMergeSpec<>(key, mergeFunction);
     }
-  }
 
-  public static <T> void propagateTag(AgentSpan parentSpan, AgentSpan childSpan, String tagName) {
-    propagateTag(parentSpan, childSpan, tagName, (p, c) -> c);
+    public static TagMergeSpec<Object> of(String tagKey) {
+      return new TagMergeSpec<>(tagKey, (parent, child) -> child);
+    }
   }
 
-  public static <T> void propagateTag(
-      AgentSpan parentSpan, AgentSpan childSpan, String tagName, BinaryOperator<T> mergeStrategy) {
-    T childTag = (T) childSpan.getTag(tagName);
+  private <T> void unsafePropagateTag(AgentSpan childSpan, TagMergeSpec<T> spec) {
+    T childTag = (T) childSpan.getTag(spec.tagKey);
     if (childTag != null) {
-      T parentTag = (T) parentSpan.getTag(tagName);
+      T parentTag = (T) parentSpan.getTag(spec.tagKey);
       if (parentTag == null) {
-        parentSpan.setTag(tagName, childTag);
+        parentSpan.setTag(spec.tagKey, childTag);
       } else {
-        parentSpan.setTag(tagName, mergeStrategy.apply(parentTag, childTag));
+        parentSpan.setTag(spec.tagKey, spec.mergeFunction.apply(parentTag, childTag));
       }
     }
   }

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

@@ -24,7 +24,6 @@
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.source.SourceResolutionException;
 import datadog.trace.civisibility.test.ExecutionResults;
-import datadog.trace.civisibility.utils.SpanUtils;
 import java.lang.reflect.Method;
 import java.util.Collection;
 import java.util.function.Consumer;
@@ -56,6 +55,7 @@ public class TestSuiteImpl implements DDTestSuite {
   private final boolean parallelized;
   private final Collection<LibraryCapability> capabilities;
   private final Consumer<AgentSpan> onSpanFinish;
+  private final SpanTagsPropagator tagsPropagator;
 
   public TestSuiteImpl(
       AgentSpanContext moduleSpanContext,
@@ -106,6 +106,7 @@ public TestSuiteImpl(
     }
 
     span = spanBuilder.start();
+    tagsPropagator = new SpanTagsPropagator(span);
 
     span.setSpanType(InternalSpanTypes.TEST_SUITE_END);
     span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_TEST_SUITE);
@@ -264,6 +265,6 @@ public TestImpl testStart(
         coverageStoreFactory,
         executionResults,
         capabilities,
-        SpanUtils.propagateCiVisibilityTagsTo(span));
+        tagsPropagator::propagateStatus);
   }
 }

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

@@ -30,7 +30,6 @@
 import datadog.trace.civisibility.ipc.SignalType;
 import datadog.trace.civisibility.source.LinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
-import datadog.trace.civisibility.utils.SpanUtils;
 import datadog.trace.util.Strings;
 import java.net.InetSocketAddress;
 import java.nio.file.Path;
@@ -289,7 +288,7 @@ private SignalResponse onModuleExecutionResultReceived(ModuleExecutionResult res
 
     testsSkipped.add(result.getTestsSkippedTotal());
 
-    SpanUtils.mergeTestFrameworks(span, result.getTestFrameworks());
+    tagsPropagator.mergeTestFrameworks(result.getTestFrameworks());
 
     return AckResponse.INSTANCE;
   }

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

@@ -1,6 +1,7 @@
 package datadog.trace.civisibility.domain.buildsystem;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.civisibility.domain.SpanTagsPropagator.TagMergeSpec;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
@@ -35,7 +36,6 @@
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.source.index.RepoIndex;
 import datadog.trace.civisibility.source.index.RepoIndexProvider;
-import datadog.trace.civisibility.utils.SpanUtils;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -54,7 +54,6 @@ public class BuildSystemSessionImpl<T extends CoverageCalculator> extends Abstra
   private final CoverageCalculator.Factory<T> coverageCalculatorFactory;
   private final T coverageCalculator;
   private final BuildSessionSettings settings;
-  private final Object tagPropagationLock = new Object();
 
   public BuildSystemSessionImpl(
       String projectName,
@@ -183,22 +182,17 @@ public AgentSpan testTaskStart(String taskName) {
 
   private void onModuleFinish(AgentSpan moduleSpan) {
     // multiple modules can finish in parallel
-    synchronized (tagPropagationLock) {
-      SpanUtils.propagateCiVisibilityTags(span, moduleSpan);
-
-      SpanUtils.propagateTag(span, moduleSpan, Tags.TEST_EARLY_FLAKE_ENABLED, Boolean::logicalOr);
-      SpanUtils.propagateTag(span, moduleSpan, Tags.TEST_EARLY_FLAKE_ABORT_REASON);
-
-      SpanUtils.propagateTag(span, moduleSpan, Tags.TEST_CODE_COVERAGE_ENABLED, Boolean::logicalOr);
-      SpanUtils.propagateTag(
-          span, moduleSpan, Tags.TEST_ITR_TESTS_SKIPPING_ENABLED, Boolean::logicalOr);
-      SpanUtils.propagateTag(span, moduleSpan, Tags.TEST_ITR_TESTS_SKIPPING_TYPE);
-      SpanUtils.propagateTag(span, moduleSpan, Tags.TEST_ITR_TESTS_SKIPPING_COUNT, Long::sum);
-      SpanUtils.propagateTag(span, moduleSpan, DDTags.CI_ITR_TESTS_SKIPPED, Boolean::logicalOr);
-
-      SpanUtils.propagateTag(
-          span, moduleSpan, Tags.TEST_TEST_MANAGEMENT_ENABLED, Boolean::logicalOr);
-    }
+    tagPropagator.propagateCiVisibilityTags(moduleSpan);
+    tagPropagator.propagateTags(
+        moduleSpan,
+        TagMergeSpec.of(Tags.TEST_EARLY_FLAKE_ENABLED, Boolean::logicalOr),
+        TagMergeSpec.of(Tags.TEST_EARLY_FLAKE_ABORT_REASON),
+        TagMergeSpec.of(Tags.TEST_CODE_COVERAGE_ENABLED, Boolean::logicalOr),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_ENABLED, Boolean::logicalOr),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_TYPE),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_COUNT, Long::sum),
+        TagMergeSpec.of(DDTags.CI_ITR_TESTS_SKIPPED, Boolean::logicalOr),
+        TagMergeSpec.of(Tags.TEST_TEST_MANAGEMENT_ENABLED, Boolean::logicalOr));
   }
 
   @Override

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

@@ -27,7 +27,6 @@
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.test.ExecutionResults;
 import datadog.trace.civisibility.test.ExecutionStrategy;
-import datadog.trace.civisibility.utils.SpanUtils;
 import java.util.Collection;
 import java.util.function.Consumer;
 import javax.annotation.Nonnull;
@@ -183,6 +182,6 @@ public TestSuiteImpl testSuiteStart(
         coverageStoreFactory,
         executionResults,
         capabilities,
-        SpanUtils.propagateCiVisibilityTagsTo(span));
+        tagsPropagator::propagateCiVisibilityTags);
   }
 }

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

@@ -1,5 +1,7 @@
 package datadog.trace.civisibility.domain.headless;
 
+import static datadog.trace.civisibility.domain.SpanTagsPropagator.TagMergeSpec;
+
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.civisibility.config.LibraryCapability;
@@ -19,7 +21,6 @@
 import datadog.trace.civisibility.source.LinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.test.ExecutionStrategy;
-import datadog.trace.civisibility.utils.SpanUtils;
 import java.util.Collection;
 import java.util.Collections;
 import javax.annotation.Nonnull;
@@ -82,22 +83,21 @@ public HeadlessTestModule testModuleStart(String moduleName, @Nullable Long star
         coverageStoreFactory,
         executionStrategy,
         capabilities,
-        this::propagateModuleTags);
+        this::onModuleFinish);
   }
 
-  private void propagateModuleTags(AgentSpan moduleSpan) {
-    SpanUtils.propagateCiVisibilityTags(span, moduleSpan);
-    SpanUtils.propagateTags(
-        span,
+  private void onModuleFinish(AgentSpan moduleSpan) {
+    tagPropagator.propagateCiVisibilityTags(moduleSpan);
+    tagPropagator.propagateTags(
         moduleSpan,
-        Tags.TEST_CODE_COVERAGE_ENABLED,
-        Tags.TEST_ITR_TESTS_SKIPPING_ENABLED,
-        Tags.TEST_ITR_TESTS_SKIPPING_TYPE,
-        Tags.TEST_ITR_TESTS_SKIPPING_COUNT,
-        Tags.TEST_EARLY_FLAKE_ENABLED,
-        Tags.TEST_EARLY_FLAKE_ABORT_REASON,
-        DDTags.CI_ITR_TESTS_SKIPPED,
-        Tags.TEST_TEST_MANAGEMENT_ENABLED);
+        TagMergeSpec.of(Tags.TEST_CODE_COVERAGE_ENABLED),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_ENABLED),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_TYPE),
+        TagMergeSpec.of(Tags.TEST_ITR_TESTS_SKIPPING_COUNT),
+        TagMergeSpec.of(Tags.TEST_EARLY_FLAKE_ENABLED),
+        TagMergeSpec.of(Tags.TEST_EARLY_FLAKE_ABORT_REASON),
+        TagMergeSpec.of(DDTags.CI_ITR_TESTS_SKIPPED),
+        TagMergeSpec.of(Tags.TEST_TEST_MANAGEMENT_ENABLED));
   }
 
   @Override

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

@@ -15,7 +15,6 @@
 import datadog.trace.civisibility.source.LinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.test.ExecutionResults;
-import datadog.trace.civisibility.utils.SpanUtils;
 import java.util.Collections;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;
@@ -81,6 +80,6 @@ public TestSuiteImpl testSuiteStart(
         coverageStoreFactory,
         executionResults,
         Collections.emptyList(),
-        SpanUtils.propagateCiVisibilityTagsTo(span));
+        tagsPropagator::propagateCiVisibilityTags);
   }
 }

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

@@ -11,7 +11,6 @@
 import datadog.trace.civisibility.domain.InstrumentationType;
 import datadog.trace.civisibility.source.LinesResolver;
 import datadog.trace.civisibility.source.SourcePathResolver;
-import datadog.trace.civisibility.utils.SpanUtils;
 import javax.annotation.Nullable;
 
 /**
@@ -60,6 +59,6 @@ public ManualApiTestModule testModuleStart(String moduleName, @Nullable Long sta
         codeowners,
         linesResolver,
         coverageStoreFactory,
-        SpanUtils.propagateCiVisibilityTagsTo(span));
+        tagPropagator::propagateCiVisibilityTags);
   }
 }