Support distributed traces in tests (#8078)

Author: nikita-tkachenko-datadog

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

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.domain;
 
+import static datadog.trace.api.TracePropagationStyle.NONE;
 import static datadog.trace.api.civisibility.CIConstants.CI_VISIBILITY_INSTRUMENTATION_NAME;
 
 import datadog.trace.api.Config;
@@ -17,9 +18,11 @@
 import datadog.trace.api.civisibility.telemetry.tag.IsHeadless;
 import datadog.trace.api.civisibility.telemetry.tag.IsUnsupportedCI;
 import datadog.trace.api.civisibility.telemetry.tag.Provider;
+import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.civisibility.codeowners.Codeowners;
 import datadog.trace.civisibility.decorator.TestDecorator;
@@ -65,7 +68,16 @@ public AbstractTestSession(
     // CI Test Cycle protocol requires session's trace ID and span ID to be the same
     IdGenerationStrategy idGenerationStrategy = config.getIdGenerationStrategy();
     DDTraceId traceId = idGenerationStrategy.generateTraceId();
-    AgentSpan.Context traceContext = new TraceContext(traceId);
+    AgentSpan.Context traceContext =
+        new TagContext(
+            CIConstants.CIAPP_TEST_ORIGIN,
+            Collections.emptyMap(),
+            null,
+            null,
+            PrioritySampling.UNSET,
+            null,
+            NONE,
+            traceId);
 
     AgentTracer.SpanBuilder spanBuilder =
         AgentTracer.get()

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

@@ -27,6 +27,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.civisibility.codeowners.Codeowners;
 import datadog.trace.civisibility.decorator.TestDecorator;
@@ -35,6 +36,7 @@
 import datadog.trace.civisibility.source.SourceResolutionException;
 import java.lang.reflect.Method;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -85,11 +87,13 @@ public TestImpl(
 
     this.context = new TestContextImpl(coverageStore);
 
+    AgentSpan.Context traceContext =
+        new TagContext(CIConstants.CIAPP_TEST_ORIGIN, Collections.emptyMap());
     AgentTracer.SpanBuilder spanBuilder =
         AgentTracer.get()
             .buildSpan(CI_VISIBILITY_INSTRUMENTATION_NAME, testDecorator.component() + ".test")
             .ignoreActiveSpan()
-            .asChildOf(null)
+            .asChildOf(traceContext)
             .withRequestContextData(RequestContextSlot.CI_VISIBILITY, context);
 
     if (startTime != null) {

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

@@ -1,10 +1,10 @@
 package datadog.trace.civisibility.interceptor;
 
-import datadog.trace.api.DDSpanTypes;
+import static datadog.trace.api.civisibility.CIConstants.CIAPP_TEST_ORIGIN;
+
 import datadog.trace.api.DDTags;
 import datadog.trace.api.interceptor.AbstractTraceInterceptor;
 import datadog.trace.api.interceptor.MutableSpan;
-import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.core.DDSpan;
 import datadog.trace.core.DDTraceCoreInfo;
 import java.util.Collection;
@@ -15,8 +15,6 @@ public class CiVisibilityTraceInterceptor extends AbstractTraceInterceptor {
   public static final CiVisibilityTraceInterceptor INSTANCE =
       new CiVisibilityTraceInterceptor(Priority.CI_VISIBILITY_TRACE);
 
-  static final UTF8BytesString CIAPP_TEST_ORIGIN = UTF8BytesString.create("ciapp-test");
-
   protected CiVisibilityTraceInterceptor(Priority priority) {
     super(priority);
   }
@@ -33,20 +31,13 @@ public Collection<? extends MutableSpan> onTraceComplete(
 
     final DDSpan spanToCheck = null == localRootSpan ? firstSpan : localRootSpan;
 
-    // If root span is not a CI visibility span, we drop the full trace.
-    CharSequence type = spanToCheck.getType(); // Don't null pointer if there is no type
-    if (type == null
-        || (!DDSpanTypes.TEST.contentEquals(type)
-            && !DDSpanTypes.TEST_SUITE_END.contentEquals(type)
-            && !DDSpanTypes.TEST_MODULE_END.contentEquals(type)
-            && !DDSpanTypes.TEST_SESSION_END.contentEquals(type))) {
+    // If root span does not originate from CI visibility, we drop the full trace.
+    CharSequence origin = spanToCheck.getOrigin();
+    if (origin == null || !CIAPP_TEST_ORIGIN.contentEquals(origin)) {
       return Collections.emptyList();
     }
 
-    // If the trace belongs to a "test", we need to set the origin to `ciapp-test` and the
-    // `library_version` tag for all spans.
-    firstSpan.context().setOrigin(CIAPP_TEST_ORIGIN);
-    firstSpan.setTag(DDTags.LIBRARY_VERSION_TAG_KEY, DDTraceCoreInfo.VERSION);
+    // If the trace belongs to a "test", we need to set the `library_version` tag for all spans.
     for (MutableSpan span : trace) {
       span.setTag(DDTags.LIBRARY_VERSION_TAG_KEY, DDTraceCoreInfo.VERSION);
     }

dd-trace-core/src/main/java/datadog/trace/civisibility/interceptor/CiVisibilityTraceInterceptor.java

@@ -269,7 +269,8 @@ protected TagContext build() {
             baggage,
             samplingPriorityOrDefault(traceId, samplingPriority),
             traceConfig,
-            style());
+            style(),
+            DDTraceId.ZERO);
       }
     }
     return null;

dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java

@@ -49,7 +49,15 @@ public ExtractedContext(
       final PropagationTags propagationTags,
       final TraceConfig traceConfig,
       final TracePropagationStyle propagationStyle) {
-    super(origin, tags, httpHeaders, baggage, samplingPriority, traceConfig, propagationStyle);
+    super(
+        origin,
+        tags,
+        httpHeaders,
+        baggage,
+        samplingPriority,
+        traceConfig,
+        propagationStyle,
+        DDTraceId.ZERO);
     this.traceId = traceId;
     this.spanId = spanId;
     this.endToEndStartTime = endToEndStartTime;

dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java

@@ -2,7 +2,9 @@ package datadog.trace.civisibility.interceptor
 
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
+import datadog.trace.api.civisibility.CIConstants
 import datadog.trace.common.writer.ListWriter
+import datadog.trace.core.DDSpanContext
 import datadog.trace.core.test.DDCoreSpecification
 import spock.lang.Timeout
 
@@ -16,31 +18,49 @@ class CiVisibilityTraceInterceptorTest extends DDCoreSpecification {
     tracer?.close()
   }
 
-  def "discard a trace that does not come from a test"() {
+  def "discard a trace that does not come from ci app"() {
     tracer.addTraceInterceptor(CiVisibilityTraceInterceptor.INSTANCE)
     tracer.buildSpan("sample-span").start().finish()
 
     expect:
     writer.size() == 0
   }
 
-  def "add ciapp origin and tracer version to spans of type #spanType"() {
+  def "do not discard a trace that comes from ci app"() {
+    tracer.addTraceInterceptor(CiVisibilityTraceInterceptor.INSTANCE)
+
+    def span = tracer.buildSpan("sample-span").start()
+    ((DDSpanContext) span.context()).origin = CIConstants.CIAPP_TEST_ORIGIN
+    span.finish()
+
+    expect:
+    writer.size() == 1
+  }
+
+  def "add tracer version to spans of type #spanType"() {
     setup:
     tracer.addTraceInterceptor(CiVisibilityTraceInterceptor.INSTANCE)
 
-    tracer.buildSpan("sample-span").withSpanType(spanType).start().finish()
+
+    def span = tracer.buildSpan("sample-span").withSpanType(spanType).start()
+    ((DDSpanContext) span.context()).origin = CIConstants.CIAPP_TEST_ORIGIN
+    span.finish()
     writer.waitForTraces(1)
 
     expect:
     def trace = writer.firstTrace()
     trace.size() == 1
 
-    def span = trace[0]
+    def receivedSpan = trace[0]
 
-    span.context().origin == CiVisibilityTraceInterceptor.CIAPP_TEST_ORIGIN
-    span.getTag(DDTags.LIBRARY_VERSION_TAG_KEY) != null
+    receivedSpan.getTag(DDTags.LIBRARY_VERSION_TAG_KEY) != null
 
     where:
-    spanType << [DDSpanTypes.TEST, DDSpanTypes.TEST_SUITE_END, DDSpanTypes.TEST_MODULE_END]
+    spanType << [
+      DDSpanTypes.TEST,
+      DDSpanTypes.TEST_SUITE_END,
+      DDSpanTypes.TEST_MODULE_END,
+      DDSpanTypes.TEST_SESSION_END
+    ]
   }
 }

dd-trace-core/src/test/groovy/datadog/trace/civisibility/interceptor/CiVisibilityTraceInterceptorTest.groovy

@@ -12,4 +12,6 @@ public interface CIConstants {
   String CI_VISIBILITY_INSTRUMENTATION_NAME = "civisibility";
 
   String FAIL_FAST_TEST_ORDER = "FAILFAST";
+
+  String CIAPP_TEST_ORIGIN = "ciapp-test";
 }

internal-api/src/main/java/datadog/trace/api/civisibility/CIConstants.java

@@ -34,13 +34,14 @@ public class TagContext implements AgentSpan.Context.Extracted {
   private final int samplingPriority;
   private final TraceConfig traceConfig;
   private final TracePropagationStyle propagationStyle;
+  private final DDTraceId traceId;
 
   public TagContext() {
     this(null, null);
   }
 
-  public TagContext(final String origin, final Map<String, String> tags) {
-    this(origin, tags, null, null, PrioritySampling.UNSET, null, NONE);
+  public TagContext(final CharSequence origin, final Map<String, String> tags) {
+    this(origin, tags, null, null, PrioritySampling.UNSET, null, NONE, DDTraceId.ZERO);
   }
 
   public TagContext(
@@ -50,7 +51,8 @@ public TagContext(
       final Map<String, String> baggage,
       final int samplingPriority,
       final TraceConfig traceConfig,
-      final TracePropagationStyle propagationStyle) {
+      final TracePropagationStyle propagationStyle,
+      final DDTraceId traceId) {
     this.origin = origin;
     this.tags = tags;
     this.terminatedContextLinks = null;
@@ -59,6 +61,7 @@ public TagContext(
     this.samplingPriority = samplingPriority;
     this.traceConfig = traceConfig;
     this.propagationStyle = propagationStyle;
+    this.traceId = traceId;
   }
 
   public TraceConfig getTraceConfig() {
@@ -187,7 +190,7 @@ public Iterable<Map.Entry<String, String>> baggageItems() {
 
   @Override
   public DDTraceId getTraceId() {
-    return DDTraceId.ZERO;
+    return traceId;
   }
 
   @Override