Keep traces with different global tags in one OTLP payload rather than eagerly flushing

dd-trace-core/src/main/java/datadog/trace/common/writer/OtlpPayloadDispatcher.java

@@ -26,10 +26,6 @@ final class OtlpPayloadDispatcher implements PayloadDispatcher {
 
   @Override
   public void addTrace(List<? extends CoreSpan<?>> trace) {
-    // a payload carries one resource: start a new one when the trace's global tags differ
-    if (!collector.sharesResource(trace)) {
-      flush();
-    }
     collector.addTrace(trace);
     // flush proactively to keep payload size bounded
     if (collector.sizeInBytes() >= FLUSH_THRESHOLD_BYTES) {

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

@@ -13,12 +13,19 @@
 import datadog.trace.api.KnownTagCodec;
 import datadog.trace.api.KnownTags;
 import datadog.trace.api.TagMap;
+import datadog.trace.api.function.TriConsumer;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.core.Metadata;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
-/** Span attribute keys shared by the OTLP protobuf and JSON trace writers. */
+/**
+ * Span attributes shared by the OTLP protobuf and JSON trace writers: the keys and baggage rules
+ * they use, and a {@link Writer} for the tags a span carries as attributes.
+ */
 final class OtlpSpanAttributes {
 
   static final UTF8BytesString RESOURCE_NAME = UTF8BytesString.create("resource.name");
@@ -90,5 +97,65 @@ private static UTF8BytesString otelKey(long tagId) {
     return UTF8BytesString.create(KnownTagCodec.openTelemetryTagOf(tagId));
   }
 
+  /**
+   * Writes the tags an OTLP span carries as attributes: the span's direct tags, plus the span
+   * defaults ({@code DD_TRACE_SPAN_TAGS}, {@code env}, ...) it doesn't set itself. The global tags
+   * are left to the resource, which carries them once per payload.
+   *
+   * <p>A span default is only written while the span still carries it: one the span removed, or
+   * that interception consumed into a span field (such as {@code resource.name}), is skipped. When
+   * the tracer tags need interception they're copied into each span, so global tags then also
+   * appear on the span, as before.
+   *
+   * <p>Not thread-safe: each writer owns one.
+   */
+  static final class Writer<T> {
+    private final T target;
+    private final TriConsumer<T, Integer, TagMap.EntryReader> tagWriter;
+
+    // per-span state
+    private Integer direction;
+    private TagMap spanTags;
+    private TagMap spanDefaultTags;
+    private final List<String> overriddenDefaults = new ArrayList<>(4);
+
+    Writer(T target, TriConsumer<T, Integer, TagMap.EntryReader> tagWriter) {
+      this.target = target;
+      this.tagWriter = tagWriter;
+    }
+
+    /** Writes the span's attribute tags, naming them for the span's {@code direction}. */
+    void write(Metadata metadata, Integer direction) {
+      this.direction = direction;
+      this.spanTags = metadata.getTags();
+      this.spanDefaultTags = metadata.getSpanDefaultTags();
+      try {
+        spanTags.forEachLocal(this, Writer::writeDirectTag);
+        if (!spanDefaultTags.isEmpty()) {
+          spanDefaultTags.forEach(this, Writer::writeDefaultTag);
+        }
+      } finally {
+        this.spanTags = null;
+        this.spanDefaultTags = null;
+        overriddenDefaults.clear();
+      }
+    }
+
+    private void writeDirectTag(TagMap.EntryReader entry) {
+      if (!spanDefaultTags.isEmpty() && spanDefaultTags.containsKey(entry.tag())) {
+        overriddenDefaults.add(entry.tag());
+      }
+      tagWriter.accept(target, direction, entry);
+    }
+
+    private void writeDefaultTag(TagMap.EntryReader entry) {
+      // skip defaults the span set itself, or no longer carries (removed, or consumed by
+      // interception)
+      if (!overriddenDefaults.contains(entry.tag()) && spanTags.getEntry(entry.tag()) != null) {
+        tagWriter.accept(target, direction, entry);
+      }
+    }
+  }
+
   private OtlpSpanAttributes() {}
 }

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

@@ -16,7 +16,7 @@ public abstract class OtlpTraceCollector {
   private static final TagMap CONFIG_GLOBAL_TAGS =
       TagMap.fromMapImmutable(Config.get().getGlobalTags());
 
-  /** The global tags of the payload being collected; null until its first trace. */
+  /** The global tags of the resource being collected; null until the payload's first span. */
   protected TagMap payloadGlobalTags;
 
   /** Adds spans from the given trace to the collector. */
@@ -28,15 +28,6 @@ public abstract class OtlpTraceCollector {
   /** Returns the number of bytes buffered since the last collection. */
   public abstract int sizeInBytes();
 
-  /**
-   * Whether the given trace can join the payload being collected. A payload carries the global tags
-   * once, on its resource, so a trace with different global tags (after a remote config tag change)
-   * needs a payload of its own.
-   */
-  public final boolean sharesResource(List<? extends CoreSpan<?>> spans) {
-    return payloadGlobalTags == null || globalTagsOf(spans) == payloadGlobalTags;
-  }
-
   protected final boolean shouldExport(CoreSpan<?> span) {
     return span.samplingPriority() > 0 // trace-level sampling priority
         || span.getTag(SPAN_SAMPLING_MECHANISM_TAG) != null; // span-level sampling priority

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

@@ -186,14 +186,14 @@ private static void writeSpanTag(JsonWriter writer, UTF8BytesString key, long va
 
   public static class MetaWriter implements MetadataConsumer {
     private final JsonWriter writer;
-    private final OtlpSpanTags<JsonWriter> spanTags;
+    private final OtlpSpanAttributes.Writer<JsonWriter> spanTags;
 
     private boolean includeProcessTags;
     private boolean includeSamplingTags;
 
     public MetaWriter(JsonWriter writer) {
       this.writer = writer;
-      this.spanTags = new OtlpSpanTags<>(writer, OtlpTraceJson::writeSpanTag);
+      this.spanTags = new OtlpSpanAttributes.Writer<>(writer, OtlpTraceJson::writeSpanTag);
     }
 
     /** Call this to ensure process tags are written out for the next span. */

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

@@ -31,6 +31,9 @@
  * one-span lookahead: that decision (and the corresponding {@code MetadataConsumer.accept} call)
  * can only be made once we see the next span (or reach a scope/payload boundary), since it depends
  * on whether the held-back span is the last one in its trace.
+ *
+ * <p>Each resource carries the global tags of its traces. A trace whose global tags differ (after a
+ * remote config tag change) completes the current resource and starts a new one.
  */
 public final class OtlpTraceJsonCollector extends OtlpTraceCollector {
 
@@ -48,19 +51,22 @@ public final class OtlpTraceJsonCollector extends OtlpTraceCollector {
   private DDSpan currentSpan;
   private List<? extends AgentSpanLink> currentSpanLinks = Collections.emptyList();
 
-  // resource for the last global tags seen; they rarely change, so it's rarely rebuilt
-  private TagMap resourceGlobalTags;
-  private String resourceFragment;
+  // global tags of the trace being added
+  private TagMap traceGlobalTags;
+
+  // resources for the last two global tags seen, see CachedResource
+  private CachedResource currentResource;
+  private CachedResource previousResource;
 
   /** Adds the given trace spans to the collector. */
   @Override
   public void addTrace(List<? extends CoreSpan<?>> spans) {
     if (!payloadStarted) {
-      payloadGlobalTags = globalTagsOf(spans);
       start();
       payloadStarted = true;
     }
 
+    traceGlobalTags = globalTagsOf(spans);
     try {
       for (CoreSpan<?> span : spans) {
         visitSpan(span);
@@ -94,25 +100,23 @@ public OtlpPayload collectTraces() {
     }
   }
 
-  /** Prepare temporary elements to collect trace data. */
+  /**
+   * Prepare temporary elements to collect trace data. Lazily allocate Resource when spans are
+   * added.
+   */
   private void start() {
     writer = new JsonWriter();
     metaWriter = new OtlpTraceJson.MetaWriter(writer);
 
     writer.beginObject();
     writer.name("resourceSpans").beginArray();
-    writer.beginObject();
-    writer.name("resource").jsonValue(resourceFragment());
-    writer.name("scopeSpans").beginArray();
-
-    // 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;
     payloadGlobalTags = null;
+    traceGlobalTags = null;
     anySpanWritten = false;
 
     writer = null;
@@ -124,12 +128,31 @@ private void stop() {
   }
 
   private String resourceFragment() {
-    // Update resourceGlobalTags when RC updates the DD_TAGS value (reflected by payloadGlobalTags)
-    if (resourceFragment == null || resourceGlobalTags != payloadGlobalTags) {
-      resourceFragment = buildTraceResourceFragment(payloadGlobalTags);
-      resourceGlobalTags = payloadGlobalTags;
+    // Rebuild when the open resource's tags (payloadGlobalTags) differ from the cached one
+    if (currentResource == null || currentResource.globalTags != payloadGlobalTags) {
+      CachedResource replaced = currentResource;
+      currentResource =
+          previousResource != null && previousResource.globalTags == payloadGlobalTags
+              ? previousResource
+              : new CachedResource(
+                  payloadGlobalTags, buildTraceResourceFragment(payloadGlobalTags));
+      previousResource = replaced;
     }
-    return resourceFragment;
+    return currentResource.fragment;
+  }
+
+  private void visitResource(TagMap globalTags) {
+    if (currentScope != null) {
+      completeResource();
+    }
+    payloadGlobalTags = globalTags;
+
+    writer.beginObject();
+    writer.name("resource").jsonValue(resourceFragment());
+    writer.name("scopeSpans").beginArray();
+
+    // for now put all spans under the default scope
+    visitScopedSpans(DEFAULT_TRACE_SCOPE);
   }
 
   private void visitScopedSpans(OtelInstrumentationScope scope) {
@@ -148,6 +171,11 @@ private void visitSpan(CoreSpan<?> span) {
     if (!shouldExport(span)) {
       return;
     }
+    if (currentScope == null || traceGlobalTags != payloadGlobalTags) {
+      // if payload's first exported span, start new resource/scope. Otherwise (global tags
+      // changed), close old resource/scope and start new resource/scope
+      visitResource(traceGlobalTags);
+    }
     if (currentSpan != null) {
       // ensure last span written at trace boundary includes sampling tags
       if (!span.getTraceId().equals(currentSpan.getTraceId())) {
@@ -159,14 +187,12 @@ private void visitSpan(CoreSpan<?> span) {
     currentSpanLinks = currentSpan.getLinks();
   }
 
-  // called once we've processed all scopes and span messages
+  // called once we've processed all resources and span messages
   private OtlpPayload completePayload() {
     if (currentScope != null) {
-      completeScope();
+      completeResource();
     }
 
-    writer.endArray(); // scopeSpans
-    writer.endObject(); // resourceSpans[0]
     writer.endArray(); // resourceSpans
     writer.endObject(); // root
 
@@ -178,6 +204,14 @@ private OtlpPayload completePayload() {
     return new OtlpPayload(ByteBuffer.wrap(bytes), JSON_CONTENT_TYPE);
   }
 
+  // called once we've processed all spans with the same global tags
+  private void completeResource() {
+    completeScope();
+
+    writer.endArray(); // scopeSpans
+    writer.endObject(); // resourceSpans[i]
+  }
+
   // called once we've processed all spans in a specific scope
   private void completeScope() {
     if (currentSpan != null) {
@@ -213,4 +247,23 @@ private void completeSpan() {
           "OTLP payload exceeds maximum buffer size of " + MAX_CAPACITY_BYTES + " bytes");
     }
   }
+
+  /**
+   * A resource built for one global tags map, cached so it's only rebuilt when the tags change.
+   *
+   * <p>A trace takes its global tags from the config snapshot when it starts, so after a remote
+   * config tag change, traces started before the change (such as long-lived traces still sending
+   * partial flushes) keep arriving with the old tags, interleaved with new traces. The collector
+   * keeps the current and the previous resource, so switching back and forth between the two swaps
+   * them instead of rebuilding each time.
+   */
+  private static final class CachedResource {
+    final TagMap globalTags;
+    final String fragment;
+
+    CachedResource(TagMap globalTags, String fragment) {
+      this.globalTags = globalTags;
+      this.fragment = fragment;
+    }
+  }
 }

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

@@ -264,14 +264,14 @@ private static void writeSpanTag(StreamingBuffer buf, UTF8BytesString key, long
 
   public static class MetaWriter implements MetadataConsumer {
     private final StreamingBuffer buf;
-    private final OtlpSpanTags<StreamingBuffer> spanTags;
+    private final OtlpSpanAttributes.Writer<StreamingBuffer> spanTags;
 
     private boolean includeProcessTags;
     private boolean includeSamplingTags;
 
     public MetaWriter(StreamingBuffer buf) {
       this.buf = buf;
-      this.spanTags = new OtlpSpanTags<>(buf, OtlpTraceProto::writeSpanTag);
+      this.spanTags = new OtlpSpanAttributes.Writer<>(buf, OtlpTraceProto::writeSpanTag);
     }
 
     /** Call this to ensure process tags are written out for the next span. */