Review feedback: support collecting multiple traces into one payload

Author: mcculls

dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceCollector.java

@@ -1,12 +1,14 @@
 package datadog.trace.core.otlp.trace;
 
-import datadog.trace.core.DDSpan;
+import datadog.trace.core.CoreSpan;
 import datadog.trace.core.otlp.common.OtlpPayload;
 import java.util.List;
 
-/** Collects trace spans ready for export. */
+/** Collects traces ready for export. */
 public interface OtlpTraceCollector {
-  OtlpTraceCollector NOOP_COLLECTOR = spans -> OtlpPayload.EMPTY;
+  OtlpTraceCollector NOOP_COLLECTOR = () -> OtlpPayload.EMPTY;
 
-  OtlpPayload collectSpans(List<DDSpan> spans);
+  OtlpPayload collectTraces();
+
+  default void addTrace(List<? extends CoreSpan<?>> spans) {}
 }

dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceProto.java

@@ -259,19 +259,27 @@ private static int spanKind(CharSequence spanKind) {
 
   public static class MetaWriter implements MetadataConsumer {
     private final StreamingBuffer buf;
-    private boolean firstSpan = true;
+
+    private boolean includeProcessTags;
+    private boolean includeSamplingTags;
 
     public MetaWriter(StreamingBuffer buf) {
       this.buf = buf;
     }
 
-    public void reset() {
-      firstSpan = true;
+    /** Call this to ensure process tags are written out for the next span. */
+    public void includeProcessTags() {
+      includeProcessTags = true;
+    }
+
+    /** Call this to ensure sampling tags are written out for the next span. */
+    public void includeSamplingTags() {
+      includeSamplingTags = true;
     }
 
     @Override
     public void accept(Metadata metadata) {
-      if ((firstSpan || metadata.topLevel()) && metadata.hasSamplingPriority()) {
+      if ((includeSamplingTags || metadata.topLevel()) && metadata.hasSamplingPriority()) {
         writeSpanTag(buf, SAMPLING_PRIORITY_KEY, metadata.samplingPriority());
       }
       if (metadata.measured()) {
@@ -297,13 +305,15 @@ public void accept(Metadata metadata) {
       if (metadata.getOrigin() != null) {
         writeSpanTag(buf, ORIGIN_KEY, metadata.getOrigin());
       }
-      if (firstSpan && metadata.processTags() != null) {
+      if (includeProcessTags && metadata.processTags() != null) {
         writeSpanTag(buf, PROCESS_TAGS_KEY, metadata.processTags());
       }
 
       metadata.getTags().forEach(tagEntry -> writeSpanTag(buf, tagEntry));
 
-      firstSpan = false;
+      // reset for next span
+      includeProcessTags = false;
+      includeSamplingTags = false;
     }
   }
 }

dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceProtoCollector.java

@@ -9,6 +9,7 @@
 import datadog.communication.serialization.GrowableBuffer;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.bootstrap.otel.common.OtelInstrumentationScope;
+import datadog.trace.core.CoreSpan;
 import datadog.trace.core.DDSpan;
 import datadog.trace.core.otlp.common.OtlpCommonProto;
 import datadog.trace.core.otlp.common.OtlpPayload;
@@ -48,6 +49,8 @@ public final class OtlpTraceProtoCollector implements OtlpTraceCollector {
   private final List<byte[]> scopedChunks = new ArrayList<>();
   private final List<byte[]> spanChunks = new ArrayList<>();
 
+  private boolean payloadStarted;
+
   // total number of chunked bytes at different nesting levels
   private int payloadBytes;
   private int scopedBytes;
@@ -56,19 +59,31 @@ public final class OtlpTraceProtoCollector implements OtlpTraceCollector {
   private OtelInstrumentationScope currentScope;
   private DDSpan currentSpan;
 
+  /** Adds the given trace spans to the collector. */
+  @Override
+  public void addTrace(List<? extends CoreSpan<?>> spans) {
+    if (!payloadStarted) {
+      start();
+      metaWriter.includeProcessTags();
+      payloadStarted = true;
+    }
+
+    for (int i = 0, len = spans.size(); i < len; i++) {
+      if (i == 0 || i == len - 1) {
+        metaWriter.includeSamplingTags();
+      }
+      visitSpan(spans.get(i));
+    }
+  }
+
   /**
-   * Collects trace spans and marshalls them into a chunked payload.
+   * Marshalls the traces collected so far into a chunked payload.
    *
    * <p>This payload is only valid for the calling thread until the next collection.
    */
   @Override
-  public OtlpPayload collectSpans(List<DDSpan> spans) {
-    OtlpCommonProto.recalibrateCaches();
-    start();
+  public OtlpPayload collectTraces() {
     try {
-      // for now put all spans under the default scope
-      visitScopedSpans(DEFAULT_TRACE_SCOPE);
-      spans.forEach(this::visitSpan);
       return completePayload();
     } finally {
       stop();
@@ -77,14 +92,22 @@ public OtlpPayload collectSpans(List<DDSpan> spans) {
 
   /** Prepare temporary elements to collect trace data. */
   private void start() {
+
     // clear payloadChunks in case it wasn't fully consumed via OtlpPayload
     payloadChunks.clear();
+
+    // remove stale entries from caches
+    OtlpCommonProto.recalibrateCaches();
+
+    // for now put all spans under the default scope
+    visitScopedSpans(DEFAULT_TRACE_SCOPE);
   }
 
   /** Cleanup elements used to collect trace data. */
   private void stop() {
+    payloadStarted = false;
+
     buf.reset();
-    metaWriter.reset();
 
     // leave payloadChunks in place so it can be consumed via OtlpPayload
     scopedChunks.clear();
@@ -105,12 +128,12 @@ private void visitScopedSpans(OtelInstrumentationScope scope) {
     currentScope = scope;
   }
 
-  private void visitSpan(DDSpan span) {
+  private void visitSpan(CoreSpan<?> span) {
     if (currentSpan != null) {
       completeSpan();
     }
-    currentSpan = span;
-    span.getLinks().forEach(this::visitSpanLink);
+    currentSpan = (DDSpan) span;
+    currentSpan.getLinks().forEach(this::visitSpanLink);
   }
 
   private void visitSpanLink(AgentSpanLink spanLink) {

dd-trace-core/src/test/java/datadog/trace/core/otlp/trace/OtlpTraceProtoTest.java

@@ -45,7 +45,7 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 /**
- * Tests for {@link OtlpTraceProto} via {@link OtlpTraceProtoCollector#collectSpans}.
+ * Tests for {@link OtlpTraceProto} via {@link OtlpTraceProtoCollector#collectTraces}.
  *
  * <p>Each test case builds real {@link DDSpan} instances via a shared {@link CoreTracer}, collects
  * them using {@link OtlpTraceProtoCollector}, drains the resulting chunked payload into a
@@ -550,10 +550,11 @@ static Stream<Arguments> cases() {
 
   @ParameterizedTest(name = "{0}")
   @MethodSource("cases")
-  void testCollectSpans(String caseName, List<SpanSpec> specs) throws IOException {
+  void testCollectTraces(String caseName, List<SpanSpec> specs) throws IOException {
     List<DDSpan> spans = buildSpans(specs);
 
-    OtlpPayload payload = OtlpTraceProtoCollector.INSTANCE.collectSpans(spans);
+    OtlpTraceProtoCollector.INSTANCE.addTrace(spans);
+    OtlpPayload payload = OtlpTraceProtoCollector.INSTANCE.collectTraces();
 
     if (spans.isEmpty()) {
       assertEquals(0, payload.getContentLength(), "empty span list must produce empty payload");