feat(otlp): cache OpenTelemetry trace state in OTLP export

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

@@ -3,12 +3,15 @@
 import static datadog.trace.core.DDSpanContext.SPAN_SAMPLING_MECHANISM_TAG;
 
 import datadog.trace.core.CoreSpan;
+import datadog.trace.core.DDSpan;
 import datadog.trace.core.otlp.common.OtlpPayload;
 import java.util.List;
 
 /** Collects traces ready for export. */
 public abstract class OtlpTraceCollector {
 
+  private static final String OTEL_TRACE_STATE_PREFIX = "ot=";
+
   /** Adds spans from the given trace to the collector. */
   public abstract void addTrace(List<? extends CoreSpan<?>> spans);
 
@@ -22,4 +25,15 @@ 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
   }
+
+  protected final String getOtlpOtelTraceState(List<? extends CoreSpan<?>> spans) {
+    for (CoreSpan<?> span : spans) {
+      if (span instanceof DDSpan) {
+        String otelTraceState =
+            ((DDSpan) span).spanContext().getPropagationTags().getOtelTraceState();
+        return otelTraceState == null ? null : OTEL_TRACE_STATE_PREFIX + otelTraceState;
+      }
+    }
+    return null;
+  }
 }

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

@@ -32,7 +32,6 @@
 import datadog.trace.core.Metadata;
 import datadog.trace.core.MetadataConsumer;
 import datadog.trace.core.PendingTrace;
-import datadog.trace.core.propagation.PropagationTags;
 import java.util.List;
 import java.util.Map;
 
@@ -48,17 +47,23 @@ private OtlpTraceJson() {}
 
   /** Writes one complete {@code Span} JSON object. */
   public static void writeSpan(
-      JsonWriter writer, DDSpan span, MetaWriter metaWriter, List<? extends AgentSpanLink> links) {
-    PropagationTags propagationTags = span.spanContext().getPropagationTags();
+      JsonWriter writer,
+      DDSpan span,
+      MetaWriter metaWriter,
+      List<? extends AgentSpanLink> links,
+      String otelTraceState) {
 
     writer.beginObject();
 
     writer.name("traceId").value(hexTraceId(span.getTraceId()));
     writer.name("spanId").value(hexSpanId(span.getSpanId()));
 
-    String tracestate = propagationTags.getW3CTracestate();
-    if (tracestate != null) {
-      writer.name("traceState").value(tracestate);
+    String traceState = span.spanContext().getPropagationTags().getW3CTracestate();
+    if (traceState == null) {
+      traceState = otelTraceState;
+    }
+    if (traceState != null) {
+      writer.name("traceState").value(traceState);
     }
 
     if (span.getParentId() != 0) {

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

@@ -46,6 +46,7 @@ public final class OtlpTraceJsonCollector extends OtlpTraceCollector {
   private OtelInstrumentationScope currentScope;
   private DDSpan currentSpan;
   private List<? extends AgentSpanLink> currentSpanLinks = Collections.emptyList();
+  private String currentOtelTraceState;
 
   /** Adds the given trace spans to the collector. */
   @Override
@@ -56,8 +57,9 @@ public void addTrace(List<? extends CoreSpan<?>> spans) {
     }
 
     try {
+      String otelTraceState = getOtlpOtelTraceState(spans);
       for (CoreSpan<?> span : spans) {
-        visitSpan(span);
+        visitSpan(span, otelTraceState);
       }
     } catch (Throwable e) {
       // reset the buffer for subsequent traces
@@ -114,6 +116,7 @@ private void stop() {
     currentScope = null;
     currentSpan = null;
     currentSpanLinks = Collections.emptyList();
+    currentOtelTraceState = null;
   }
 
   private void visitScopedSpans(OtelInstrumentationScope scope) {
@@ -128,7 +131,7 @@ private void visitScopedSpans(OtelInstrumentationScope scope) {
     writer.name("spans").beginArray();
   }
 
-  private void visitSpan(CoreSpan<?> span) {
+  private void visitSpan(CoreSpan<?> span, String otelTraceState) {
     if (!shouldExport(span)) {
       return;
     }
@@ -141,6 +144,7 @@ private void visitSpan(CoreSpan<?> span) {
     }
     currentSpan = (DDSpan) span;
     currentSpanLinks = currentSpan.getLinks();
+    currentOtelTraceState = otelTraceState;
   }
 
   // called once we've processed all scopes and span messages
@@ -185,12 +189,13 @@ private void completeSpan() {
       metaWriter.includeProcessTags();
       firstSpanInScope = false;
     }
-    writeSpan(writer, currentSpan, metaWriter, currentSpanLinks);
+    writeSpan(writer, currentSpan, metaWriter, currentSpanLinks, currentOtelTraceState);
     anySpanWritten = true;
 
     // reset temporary elements for next span
     currentSpan = null;
     currentSpanLinks = Collections.emptyList();
+    currentOtelTraceState = null;
 
     if (writer.size() > MAX_CAPACITY_BYTES) {
       throw new IllegalStateException(

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

@@ -46,7 +46,6 @@
 import datadog.trace.core.MetadataConsumer;
 import datadog.trace.core.PendingTrace;
 import datadog.trace.core.otlp.common.OtlpProtoBuffer;
-import datadog.trace.core.propagation.PropagationTags;
 
 /** Provides optimized writers for OpenTelemetry's "trace.proto" wire protocol. */
 public final class OtlpTraceProto {
@@ -81,19 +80,22 @@ public static int recordSpanMessage(
       DDSpan span,
       MetaWriter metaWriter,
       int nestedSpanLinkBytes,
-      OtlpProtoBuffer protobuf) {
-    PropagationTags propagationTags = span.spanContext().getPropagationTags();
+      OtlpProtoBuffer protobuf,
+      String otelTraceState) {
 
     writeTag(buf, 1, LEN_WIRE_TYPE);
     writeTraceId(buf, span.getTraceId());
 
     writeTag(buf, 2, LEN_WIRE_TYPE);
     writeSpanId(buf, span.getSpanId());
 
-    String tracestate = propagationTags.getW3CTracestate();
-    if (tracestate != null) {
+    String traceState = span.spanContext().getPropagationTags().getW3CTracestate();
+    if (traceState == null) {
+      traceState = otelTraceState;
+    }
+    if (traceState != null) {
       writeTag(buf, 3, LEN_WIRE_TYPE);
-      writeString(buf, tracestate);
+      writeString(buf, traceState);
     }
 
     if (span.getParentId() != 0) {

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

@@ -47,6 +47,7 @@ public final class OtlpTraceProtoCollector extends OtlpTraceCollector {
 
   private OtelInstrumentationScope currentScope;
   private DDSpan currentSpan;
+  private String currentOtelTraceState;
 
   /** Adds the given trace spans to the collector. */
   @Override
@@ -57,9 +58,10 @@ public void addTrace(List<? extends CoreSpan<?>> spans) {
     }
 
     try {
+      String otelTraceState = getOtlpOtelTraceState(spans);
       // OtlpProtoBuffer collects spans in reverse
       for (int i = spans.size() - 1; i >= 0; i--) {
-        visitSpan(spans.get(i));
+        visitSpan(spans.get(i), otelTraceState);
       }
     } catch (Throwable e) {
       // reset the buffer for subsequent traces
@@ -109,6 +111,7 @@ private void stop() {
 
     currentScope = null;
     currentSpan = null;
+    currentOtelTraceState = null;
   }
 
   private void visitScopedSpans(OtelInstrumentationScope scope) {
@@ -118,7 +121,7 @@ private void visitScopedSpans(OtelInstrumentationScope scope) {
     currentScope = scope;
   }
 
-  private void visitSpan(CoreSpan<?> span) {
+  private void visitSpan(CoreSpan<?> span, String otelTraceState) {
     if (!shouldExport(span)) {
       return;
     }
@@ -131,6 +134,7 @@ private void visitSpan(CoreSpan<?> span) {
       completeSpan();
     }
     currentSpan = (DDSpan) span;
+    currentOtelTraceState = otelTraceState;
     currentSpan.getLinks().forEach(this::visitSpanLink);
   }
 
@@ -178,10 +182,12 @@ private void completeScope() {
   // called once we've processed all span-links in a specific span
   private void completeSpan() {
 
-    scopedBytes += recordSpanMessage(buf, currentSpan, metaWriter, spanBytes, protobuf);
+    scopedBytes +=
+        recordSpanMessage(buf, currentSpan, metaWriter, spanBytes, protobuf, currentOtelTraceState);
 
     // reset temporary elements for next span
     currentSpan = null;
+    currentOtelTraceState = null;
     spanBytes = 0;
   }
 }

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

@@ -90,6 +90,13 @@ public abstract void updateOtelTraceState(
    */
   public abstract String getW3CTracestate();
 
+  /**
+   * Gets the OpenTelemetry {@code ot} tracestate member value, without the member key.
+   *
+   * @return The OpenTelemetry tracestate value, or {@code null} if none is present.
+   */
+  public abstract String getOtelTraceState();
+
   /**
    * Stores the original <a href="https://www.w3.org/TR/trace-context/#tracestate-header">W3C
    * tracestate header</a> value.

dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java

@@ -575,10 +575,16 @@ private void setW3CTracestate(String tracestate, OtelTraceState otelTraceState)
       this.otelTraceState = otelTraceState;
     }
 
-    OtelTraceState getOtelTraceState() {
+    OtelTraceState getOtelTraceStateForW3C() {
       return otelTraceState;
     }
 
+    @Override
+    public String getOtelTraceState() {
+      OtelTraceState state = otelTraceState;
+      return state == null ? null : state.getValue();
+    }
+
     void setOtelTraceState(OtelTraceState otelTraceState) {
       if (this.otelTraceState != otelTraceState) {
         this.otelTraceState = otelTraceState;

dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/W3CPTagsCodec.java

@@ -270,7 +270,7 @@ protected int estimateHeaderSize(PTags pTags) {
       size += pTags.tracestate.length();
       includesOriginalTracestate = true;
     }
-    OtelTraceState otelTraceState = pTags.getOtelTraceState();
+    OtelTraceState otelTraceState = pTags.getOtelTraceStateForW3C();
     if (otelTraceState != null) {
       size -= includesOriginalTracestate ? otelTraceState.getOriginalMemberContributionSize() : 0;
       size += OTEL_MEMBER_KEY.length() + otelTraceState.length() + 1;
@@ -749,7 +749,7 @@ private static int cleanUpAndAppendUnknown(StringBuilder sb, W3CPTags w3CPTags,
   private static boolean appendOtelAndVendorMembers(
       StringBuilder sb, PTags ptags, boolean hasDatadogMember) {
     String original = ptags.tracestate;
-    OtelTraceState otelTraceState = ptags.getOtelTraceState();
+    OtelTraceState otelTraceState = ptags.getOtelTraceStateForW3C();
     int remainingMembers = MAX_MEMBER_COUNT - (hasDatadogMember ? 1 : 0);
     int otherMemberPosition = 0;
     boolean otelTraceStateAppended = false;

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

@@ -130,27 +130,70 @@ void spanTraceStateOmittedWhenNotPropagated() throws IOException {
   }
 
   @Test
-  void spanTraceStateIncludedWhenPropagated() throws IOException {
+  void spansExportFullTraceStateWithOtelMember() throws IOException {
     PropagationTags propagationTags = PropagationTags.factory().empty();
-    propagationTags.updateW3CTracestate("vendor=state");
+    propagationTags.updateW3CTracestate("vendor=state,ot=rv:0abcdef1234567;th:80000000000000");
+    propagationTags.updateTraceSamplingPriority(
+        PrioritySampling.USER_KEEP, SamplingMechanism.EXTERNAL_OVERRIDE);
     ExtractedContext parent =
         new ExtractedContext(
             DDTraceId.ONE,
             0L,
-            PrioritySampling.UNSET,
+            PrioritySampling.USER_KEEP,
             null,
             propagationTags,
             TracePropagationStyle.DATADOG);
 
-    AgentSpan agentSpan = TRACER.startSpan("test", "op.tracestate", parent);
-    agentSpan.setResourceName("op.tracestate");
-    agentSpan.finish();
+    AgentSpan root = TRACER.startSpan("test", "op.tracestate.root", parent);
+    root.setResourceName("op.tracestate.root");
+    AgentSpan child = TRACER.startSpan("test", "op.tracestate.child", root.spanContext());
+    child.setResourceName("op.tracestate.child");
+    child.finish();
+    root.finish();
 
     OtlpTraceJsonCollector collector = new OtlpTraceJsonCollector();
-    collector.addTrace(asList((CoreSpan<?>) agentSpan));
-    Map<String, Object> parsedSpan = onlySpan(collector.collectTraces());
+    collector.addTrace(asList((CoreSpan<?>) root, (CoreSpan<?>) child));
+    List<Map<String, Object>> parsedSpans = allSpans(collector.collectTraces());
+
+    assertEquals(2, parsedSpans.size());
+    for (Map<String, Object> parsedSpan : parsedSpans) {
+      assertEquals(
+          "vendor=state,ot=rv:0abcdef1234567;th:80000000000000", parsedSpan.get("traceState"));
+    }
+  }
+
+  @Test
+  void spansExportOtelTraceStateWithoutW3CTraceState() throws IOException {
+    PropagationTags propagationTags = PropagationTags.factory().empty();
+    propagationTags.updateTraceSamplingPriority(
+        PrioritySampling.USER_KEEP, SamplingMechanism.EXTERNAL_OVERRIDE);
+    propagationTags.updateOtelTraceState(
+        DDTraceId.ONE.toLong(), 0.5, true, PrioritySampling.USER_KEEP);
+    String traceState = "ot=" + propagationTags.getOtelTraceState();
+    ExtractedContext parent =
+        new ExtractedContext(
+            DDTraceId.ONE,
+            0L,
+            PrioritySampling.USER_KEEP,
+            null,
+            propagationTags,
+            TracePropagationStyle.DATADOG);
 
-    assertEquals("vendor=state", parsedSpan.get("traceState"));
+    AgentSpan root = TRACER.startSpan("test", "op.otel.root", parent);
+    root.setResourceName("op.otel.root");
+    AgentSpan child = TRACER.startSpan("test", "op.otel.child", root.spanContext());
+    child.setResourceName("op.otel.child");
+    child.finish();
+    root.finish();
+
+    OtlpTraceJsonCollector collector = new OtlpTraceJsonCollector();
+    collector.addTrace(asList((CoreSpan<?>) root, (CoreSpan<?>) child));
+    List<Map<String, Object>> parsedSpans = allSpans(collector.collectTraces());
+
+    assertEquals(2, parsedSpans.size());
+    for (Map<String, Object> parsedSpan : parsedSpans) {
+      assertEquals(traceState, parsedSpan.get("traceState"));
+    }
   }
 
   @Test