Use W3C Trace Context trace ID as parent ID regardless of propagation style order (#7355)

* Update HttpCodec.java
* feat(core): Simplify W3C phase 3 implementation
* Added unit tests
* feat(context): Improve W3C support
Store Datadog last parent ID
Allow tags update from TagContext
Add Datadog parent header validation
* feat(context): Add W3C header override test
* fix lastParentId assignment for case when no p tag available
* feat(context): Make sure last parent id can't be invalid
* Changed comment to be clearer
* added more unit test cases for better codecov
* Clarified wording on applyTraceContextToFirstContext function

Co-authored-by: Bruce Bujon <[email protected]>
Co-authored-by: Bruce Bujon <[email protected]>

Author: mtoffl01

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

@@ -12,9 +12,9 @@
  */
 public class ExtractedContext extends TagContext {
   private final DDTraceId traceId;
-  private final long spanId;
   private final long endToEndStartTime;
   private final PropagationTags propagationTags;
+  private long spanId;
 
   public ExtractedContext(
       final DDTraceId traceId,
@@ -66,6 +66,10 @@ public final long getSpanId() {
     return spanId;
   }
 
+  public final void overrideSpanId(final long spanId) {
+    this.spanId = spanId;
+  }
+
   public final long getEndToEndStartTime() {
     return endToEndStartTime;
   }

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

@@ -1,10 +1,13 @@
 package datadog.trace.core.propagation;
 
+import static datadog.trace.api.DDTags.PARENT_ID;
 import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
+import static datadog.trace.core.propagation.DatadogHttpCodec.SPAN_ID_KEY;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DD64bTraceId;
+import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.TracePropagationStyle;
@@ -233,10 +236,7 @@ public <C> TagContext extract(
             if (traceIdMatch(context.getTraceId(), extractedContext.getTraceId())) {
               boolean comingFromTraceContext = extracted.getPropagationStyle() == TRACECONTEXT;
               if (comingFromTraceContext) {
-                // Propagate newly extracted W3C tracestate to first valid context
-                String extractedTracestate =
-                    extractedContext.getPropagationTags().getW3CTracestate();
-                context.getPropagationTags().updateW3CTracestate(extractedTracestate);
+                applyTraceContextToFirstContext(context, extractedContext, extractionCache);
               }
             } else {
               // Terminate extracted context and add it as span link
@@ -262,13 +262,48 @@ else if (extracted != null && partialContext == null) {
         return null;
       }
     }
+
+    /**
+     * Applies span ID from W3C trace context over any other valid context previously found.
+     *
+     * @param firstContext The first valid context found.
+     * @param traceContext The trace context to apply.
+     * @param extractionCache The extraction cache to get quick access to any extra information.
+     * @param <C> The carrier type.
+     */
+    private <C> void applyTraceContextToFirstContext(
+        ExtractedContext firstContext,
+        ExtractedContext traceContext,
+        ExtractionCache<C> extractionCache) {
+      // Propagate newly extracted W3C tracestate to first valid context
+      String extractedTracestate = traceContext.getPropagationTags().getW3CTracestate();
+      firstContext.getPropagationTags().updateW3CTracestate(extractedTracestate);
+      // Check if parent spans differ to reconcile them
+      if (firstContext.getSpanId() != traceContext.getSpanId()) {
+        // Override parent span id with W3C one
+        firstContext.overrideSpanId(traceContext.getSpanId());
+        // Add last parent ID as a span tag (check W3C first, else Datadog)
+        CharSequence lastParentId = traceContext.getPropagationTags().getLastParentId();
+        if (lastParentId == null) {
+          lastParentId = extractionCache.getDatadogSpanIdHex();
+        }
+        if (lastParentId != null) {
+          firstContext.putTag(PARENT_ID, lastParentId.toString());
+        }
+      }
+    }
   }
 
   private static class ExtractionCache<C>
       implements AgentPropagation.KeyClassifier,
           AgentPropagation.ContextVisitor<ExtractionCache<?>> {
     /** Cached context key-values (even indexes are header names, odd indexes are header values). */
     private final List<String> keysAndValues;
+    /**
+     * The parent span identifier from {@link DatadogHttpCodec#SPAN_ID_KEY} header formatted as 16
+     * hexadecimal characters, {@code null} if absent or invalid.
+     */
+    private String datadogSpanIdHex;
 
     public ExtractionCache(C carrier, AgentPropagation.ContextVisitor<C> getter) {
       this.keysAndValues = new ArrayList<>(32);
@@ -279,9 +314,24 @@ public ExtractionCache(C carrier, AgentPropagation.ContextVisitor<C> getter) {
     public boolean accept(String key, String value) {
       this.keysAndValues.add(key);
       this.keysAndValues.add(value);
+      cacheDatadogSpanId(key, value);
       return true;
     }
 
+    private void cacheDatadogSpanId(String key, String value) {
+      if (SPAN_ID_KEY.equalsIgnoreCase(key)) {
+        try {
+          // Parse numeric header value to format it as 16 hexadecimal character format
+          this.datadogSpanIdHex = DDSpanId.toHexStringPadded(DDSpanId.from(value));
+        } catch (NumberFormatException ignored) {
+        }
+      }
+    }
+
+    private String getDatadogSpanIdHex() {
+      return this.datadogSpanIdHex;
+    }
+
     @Override
     public void forEachKey(ExtractionCache<?> carrier, AgentPropagation.KeyClassifier classifier) {
       List<String> keysAndValues = carrier.keysAndValues;

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

@@ -110,6 +110,10 @@ static class PTags extends PropagationTags {
      */
     protected volatile String error;
 
+    /**
+     * The last parent span id using the 16-characters zero padded hexadecimal representation,
+     * {@code null} if not set.
+     */
     private volatile CharSequence lastParentId;
 
     public PTags(
@@ -271,7 +275,6 @@ public CharSequence getLastParentId() {
 
     @Override
     public void updateLastParentId(CharSequence lastParentId) {
-      lastParentId = "0000000000000000".equals(lastParentId) ? null : lastParentId;
       if (!Objects.equals(this.lastParentId, lastParentId)) {
         clearCachedHeader(W3C);
         this.lastParentId = TagValue.from(lastParentId);

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

@@ -26,7 +26,6 @@ public class W3CPTagsCodec extends PTagsCodec {
   private static final int MIN_ALLOWED_CHAR = 32;
   private static final int MAX_ALLOWED_CHAR = 126;
   private static final int MAX_MEMBER_COUNT = 32;
-  private static final String LAST_PARENT_ZERO = "0000000000000000";
 
   @Override
   PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
@@ -96,7 +95,7 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
     TagValue traceIdTagValue = null;
     boolean appsecPropagationEnabled = false;
     int maxUnknownSize = 0;
-    CharSequence lastParentId = LAST_PARENT_ZERO;
+    CharSequence lastParentId = null;
     while (tagPos < ddMemberValueEnd) {
       int tagKeyEndsAt =
           validateCharsUntilSeparatorOrEnd(
@@ -242,7 +241,7 @@ protected int appendPrefix(StringBuilder sb, PTags ptags) {
     }
     // append last ParentId (p)
     CharSequence lastParent = ptags.getLastParentId();
-    if (lastParent != null && !lastParent.equals(LAST_PARENT_ZERO)) {
+    if (lastParent != null) {
       if (sb.length() > EMPTY_SIZE) {
         sb.append(';');
       }