Fix propagation tag serialization for trace chunks

Author: sarahchen6

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

@@ -337,7 +337,8 @@ public void map(List<? extends CoreSpan<?>> trace, final Writable writable) {
       span.processTagsAndBaggage(
           metaWriter
               .withWritable(writable)
-              .forSpan(i == 0, i == trace.size() - 1, !firstSpanWritten));
+              .forSpan(i == 0, i == trace.size() - 1, !firstSpanWritten),
+          i == 0);
       if (!metaStruct.isEmpty()) {
         /* 13 */
         metaStructWriter.withWritable(writable).write(metaStruct);

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

@@ -82,7 +82,8 @@ public void map(final List<? extends CoreSpan<?>> trace, final Writable writable
       span.processTagsAndBaggage(
           metaWriter
               .withWritable(writable)
-              .forSpan(i == 0, i == trace.size() - 1, !firstSpanWritten));
+              .forSpan(i == 0, i == trace.size() - 1, !firstSpanWritten),
+          i == 0);
       /* 12 */
       writeDictionaryEncoded(writable, span.getType());
       firstSpanWritten = true;

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

@@ -89,7 +89,7 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
     }
 
     CoreSpan<?> firstSpan = trace.get(0);
-    firstSpan.processTagsAndBaggageWithStructuredLinks(spanMetadata);
+    firstSpan.processTagsAndBaggageWithStructuredLinks(spanMetadata, true);
     Metadata firstSpanMeta = spanMetadata.metadata;
 
     // encoded fields: 1..7, but skipping #5, as not required by tracers and set by the agent.
@@ -106,7 +106,7 @@ public void map(List<? extends CoreSpan<?>> trace, Writable writable) {
     // traceID = 6, the ID of the trace to which all spans in this chunk belong
     encodeTraceId(writable, 6, firstSpan.getTraceId());
     // samplingMechanism = 7, uint32
-    encodeInt(writable, 7, parseSamplingMechanism(firstSpanMeta.getTags()));
+    encodeInt(writable, 7, parseSamplingMechanism(firstSpanMeta.getBaggage()));
   }
 
   private Map<String, Object> buildChunkAttributes(List<? extends CoreSpan<?>> trace) {
@@ -126,9 +126,10 @@ private void encodeSpans(Writable writable, int fieldId, List<? extends CoreSpan
 
     // spanMetadata will already have data from first span.
     Metadata meta = spanMetadata.metadata;
-    for (CoreSpan<?> span : spans) {
+    for (int i = 0; i < spans.size(); i++) {
+      CoreSpan<?> span = spans.get(i);
       if (meta == null) {
-        span.processTagsAndBaggageWithStructuredLinks(spanMetadata);
+        span.processTagsAndBaggageWithStructuredLinks(spanMetadata, i == 0);
         meta = spanMetadata.metadata;
       }
       TagMap tags = meta.getTags();
@@ -586,8 +587,8 @@ static int getSpanKindValue(CharSequence spanKind) {
    * <p>V1 payload expects only the numeric mechanism, so we normalize both forms to a positive
    * integer and fall back to {@link SamplingMechanism#DEFAULT} when absent or malformed.
    */
-  private int parseSamplingMechanism(TagMap tags) {
-    String decisionMaker = tags.getString(KEY_DECISION_MAKER);
+  private int parseSamplingMechanism(Map<String, String> propagationMetadata) {
+    String decisionMaker = propagationMetadata.get(KEY_DECISION_MAKER);
     if (decisionMaker == null || decisionMaker.isEmpty()) {
       return SamplingMechanism.DEFAULT;
     }

dd-trace-core/src/main/java/datadog/trace/core/CoreSpan.java

@@ -98,6 +98,10 @@ default String getSpanKindString() {
 
   void processTagsAndBaggage(MetadataConsumer consumer);
 
+  default void processTagsAndBaggage(MetadataConsumer consumer, boolean firstInChunk) {
+    processTagsAndBaggage(consumer);
+  }
+
   /**
    * Variant of {@link #processTagsAndBaggage(MetadataConsumer)} for protocols that serialize span
    * links as first-class structured data rather than tags. Baggage tag injection still follows the
@@ -110,6 +114,11 @@ default void processTagsAndBaggageWithStructuredLinks(MetadataConsumer consumer)
     processTagsAndBaggage(consumer);
   }
 
+  default void processTagsAndBaggageWithStructuredLinks(
+      MetadataConsumer consumer, boolean firstInChunk) {
+    processTagsAndBaggageWithStructuredLinks(consumer);
+  }
+
   T setSamplingPriority(int samplingPriority, int samplingMechanism);
 
   T setSamplingPriority(

dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java

@@ -780,11 +780,23 @@ public void processTagsAndBaggage(final MetadataConsumer consumer) {
     context.processTagsAndBaggage(consumer, longRunningVersion, this);
   }
 
+  @Override
+  public void processTagsAndBaggage(final MetadataConsumer consumer, final boolean firstInChunk) {
+    context.processTagsAndBaggage(consumer, longRunningVersion, this, firstInChunk);
+  }
+
   @Override
   public void processTagsAndBaggageWithStructuredLinks(final MetadataConsumer consumer) {
     context.processTagsAndBaggageWithStructuredLinks(consumer, longRunningVersion, this);
   }
 
+  @Override
+  public void processTagsAndBaggageWithStructuredLinks(
+      final MetadataConsumer consumer, final boolean firstInChunk) {
+    context.processTagsAndBaggageWithStructuredLinks(
+        consumer, longRunningVersion, this, firstInChunk);
+  }
+
   @Override
   public boolean isError() {
     return context.getErrorFlag();

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -1196,7 +1196,26 @@ void earlyProcessTags(AppendableSpanLinks links) {
   void processTagsAndBaggage(
       final MetadataConsumer consumer, int longRunningVersion, DDSpan restrictedSpan) {
     processTagsAndBaggage(
-        consumer, longRunningVersion, restrictedSpan, injectLinksAsTags, injectBaggageAsTags);
+        consumer,
+        longRunningVersion,
+        restrictedSpan,
+        injectLinksAsTags,
+        injectBaggageAsTags,
+        propagationTags);
+  }
+
+  void processTagsAndBaggage(
+      final MetadataConsumer consumer,
+      int longRunningVersion,
+      DDSpan restrictedSpan,
+      boolean firstInChunk) {
+    processTagsAndBaggage(
+        consumer,
+        longRunningVersion,
+        restrictedSpan,
+        injectLinksAsTags,
+        injectBaggageAsTags,
+        firstInChunk ? getPropagationTags() : null);
   }
 
   /**
@@ -1210,15 +1229,31 @@ void processTagsAndBaggageWithStructuredLinks(
         longRunningVersion,
         restrictedSpan,
         false, // injectLinksAsTags
-        injectBaggageAsTags);
+        injectBaggageAsTags,
+        propagationTags);
+  }
+
+  void processTagsAndBaggageWithStructuredLinks(
+      final MetadataConsumer consumer,
+      int longRunningVersion,
+      DDSpan restrictedSpan,
+      boolean firstInChunk) {
+    processTagsAndBaggage(
+        consumer,
+        longRunningVersion,
+        restrictedSpan,
+        false, // injectLinksAsTags
+        injectBaggageAsTags,
+        firstInChunk ? getPropagationTags() : null);
   }
 
   void processTagsAndBaggage(
       final MetadataConsumer consumer,
       int longRunningVersion,
       DDSpan restrictedSpan,
       boolean injectLinksAsTags,
-      boolean injectBaggageAsTags) {
+      boolean injectBaggageAsTags,
+      PropagationTags serializedPropagationTags) {
     // NOTE: The span is passed for the sole purpose of allowing updating & reading of the span
     // links
     // This is a compromise to avoid...
@@ -1243,9 +1278,14 @@ void processTagsAndBaggage(
         if (w3cBaggage != null) {
           injectW3CBaggageTags(baggageItemsWithPropagationTags);
         }
-        propagationTags.fillTagMap(baggageItemsWithPropagationTags);
+        if (serializedPropagationTags != null) {
+          serializedPropagationTags.fillTagMap(baggageItemsWithPropagationTags);
+        }
       } else {
-        baggageItemsWithPropagationTags = propagationTags.createTagMap();
+        baggageItemsWithPropagationTags =
+            serializedPropagationTags == null
+                ? EMPTY_BAGGAGE
+                : serializedPropagationTags.createTagMap();
       }
 
       consumer.accept(

dd-trace-core/src/test/groovy/datadog/trace/common/writer/ddagent/TraceMapperV1PayloadTest.groovy

@@ -99,8 +99,7 @@ class TraceMapperV1PayloadTest extends DDSpecification {
       (Tags.SPAN_KIND): Tags.SPAN_KIND_CLIENT,
       "attr.string": "value",
       "attr.bool"  : true,
-      "attr.number": 12.5d,
-      "_dd.p.dm"   : "-3"
+      "attr.number": 12.5d
     ]
     def span = new TraceGenerator.PojoSpan(
       "service-a",
@@ -112,7 +111,7 @@ class TraceMapperV1PayloadTest extends DDSpecification {
       1000L,
       2000L,
       1,
-      [:],
+      ["_dd.p.dm": "-3"],
       tags,
       "web",
       false,
@@ -184,8 +183,8 @@ class TraceMapperV1PayloadTest extends DDSpecification {
       1000L,
       2000L,
       0,
-      [:],
       decisionMakerTag == null ? [:] : ["_dd.p.dm": decisionMakerTag],
+      [:],
       "custom",
       false,
       PrioritySampling.SAMPLER_KEEP,
@@ -884,7 +883,7 @@ class TraceMapperV1PayloadTest extends DDSpecification {
     assertEquals(1, chunkAttributes.size())
     assertEqualsWithNullAsEmpty(firstSpan.getLocalRootSpan().getServiceName(), chunkAttributes.get("service"))
     assertArrayEquals(traceIdBytes(firstSpan.getTraceId()), traceId)
-    assertEquals(expectedSamplingMechanism(firstSpan.getTags()), samplingMechanism)
+    assertEquals(expectedSamplingMechanism(firstSpan.getBaggage()), samplingMechanism)
   }
 
   private static byte[] traceIdBytes(DDTraceId traceId) {
@@ -1064,8 +1063,8 @@ class TraceMapperV1PayloadTest extends DDSpecification {
     }
   }
 
-  private static int expectedSamplingMechanism(Map<String, Object> tags) {
-    Object decisionMakerRaw = tags.get("_dd.p.dm")
+  private static int expectedSamplingMechanism(Map<String, String> propagationMetadata) {
+    Object decisionMakerRaw = propagationMetadata.get("_dd.p.dm")
     if (decisionMakerRaw == null) {
       return SamplingMechanism.DEFAULT
     }

dd-trace-core/src/test/java/datadog/trace/common/writer/ddagent/V1PayloadReader.java

@@ -110,9 +110,44 @@ private V1PayloadReader() {}
 
   /** Decodes the first span of the first chunk of an encoded V1 payload. */
   public static V1Span readFirstSpan(byte[] encoded) throws IOException {
+    V1Chunk chunk = readFirstChunk(encoded);
+    assertEquals(1, chunk.getSpans().size());
+    return chunk.getSpans().get(0);
+  }
+
+  /** Decodes the first chunk of an encoded V1 payload. */
+  public static V1Chunk readFirstChunk(byte[] encoded) throws IOException {
     MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(new ArrayBufferInput(encoded));
     List<String> stringTable = newStringTable();
-    return readFirstSpan(unpacker, stringTable);
+    int payloadFieldCount = unpacker.unpackMapHeader();
+    for (int i = 0; i < payloadFieldCount; i++) {
+      int payloadFieldId = unpacker.unpackInt();
+      if (payloadFieldId != PayloadField.CHUNKS) {
+        skipPayloadField(unpacker, payloadFieldId, stringTable);
+        continue;
+      }
+      int chunkCount = unpacker.unpackArrayHeader();
+      assertEquals(1, chunkCount);
+      int chunkFieldCount = unpacker.unpackMapHeader();
+      List<V1Span> spans = Collections.emptyList();
+      int samplingMechanism = 0;
+      for (int j = 0; j < chunkFieldCount; j++) {
+        int chunkFieldId = unpacker.unpackInt();
+        if (chunkFieldId == ChunkField.SPANS) {
+          int spanCount = unpacker.unpackArrayHeader();
+          spans = new ArrayList<>(spanCount);
+          for (int k = 0; k < spanCount; k++) {
+            spans.add(decodeSpan(unpacker, stringTable));
+          }
+        } else if (chunkFieldId == ChunkField.SAMPLING_MECHANISM) {
+          samplingMechanism = unpacker.unpackInt();
+        } else {
+          skipChunkField(unpacker, chunkFieldId, stringTable);
+        }
+      }
+      return new V1Chunk(spans, samplingMechanism);
+    }
+    throw new AssertionError("Could not find first chunk in v1 payload");
   }
 
   /** Creates a string table seeded with the empty string at index 0, as the writer expects. */
@@ -489,6 +524,24 @@ public List<V1SpanEvent> getEvents() {
     }
   }
 
+  public static final class V1Chunk {
+    private final List<V1Span> spans;
+    private final int samplingMechanism;
+
+    private V1Chunk(List<V1Span> spans, int samplingMechanism) {
+      this.spans = spans;
+      this.samplingMechanism = samplingMechanism;
+    }
+
+    public List<V1Span> getSpans() {
+      return spans;
+    }
+
+    public int getSamplingMechanism() {
+      return samplingMechanism;
+    }
+  }
+
   /** A decoded V1 structured span link. */
   public static final class V1SpanLink {
     private final byte[] traceId;