Build one trace id per extracted compound context

dd-trace-api/src/main/java/datadog/trace/api/DDTraceId.java

@@ -39,6 +39,19 @@ public static DDTraceId from(String s) throws NumberFormatException {
     return DD64bTraceId.create(LongStringUtils.parseUnsignedLong(s), s);
   }
 
+  /**
+   * Creates a new {@link DD64bTraceId 64-bit TraceId} from an id already parsed from its {@link
+   * #toString() String representation}, as {@link #from(String)} would without parsing it again.
+   *
+   * @param id The parsed id.
+   * @param s The {@link #toString() String representation} {@code id} was parsed from, which isn't
+   *     checked against it.
+   * @return A new {@link DDTraceId} instance.
+   */
+  public static DDTraceId from(long id, String s) {
+    return DD64bTraceId.create(id, s);
+  }
+
   /**
    * Creates a new {@link DDTraceId} from the given {@link #toHexString() hexadecimal
    * representation}.

dd-trace-api/src/test/java/datadog/trace/api/DDTraceIdTest.java

@@ -200,6 +200,25 @@ void failParsingIllegal128BitIdHexadecimalStringRepresentationFromPartialString(
         () -> DD128bTraceId.fromHex(hexId, start, length, lowerCaseOnly));
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"1", "123456789", "18446744073709551615", "0", "007"})
+  void createFromParsedIdAsFromString(String stringId) {
+    DDTraceId parsed = DDTraceId.from(stringId);
+    DDTraceId created = DDTraceId.from(parsed.toLong(), stringId);
+
+    assertEquals(parsed, created);
+    assertEquals(parsed.toString(), created.toString());
+    assertEquals(parsed.getClass(), created.getClass());
+  }
+
+  @Test
+  void createFromParsedIdKeepsTheString() {
+    String stringId = "123456789";
+
+    assertSame(stringId, DDTraceId.from(123456789L, stringId).toString());
+    assertSame(DDTraceId.ZERO, DDTraceId.from(0L, "0"));
+  }
+
   @Test
   void checkZeroConstantInitialization() {
     DDTraceId zero = DDTraceId.ZERO;

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

@@ -16,6 +16,7 @@
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.TracePropagationStyle;
+import datadog.trace.api.internal.util.LongStringUtils;
 import datadog.trace.api.llmobs.LLMObsInternal;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.core.DDSpanContext;
@@ -114,6 +115,14 @@ private static class DatadogContextInterpreter extends ContextInterpreter {
     /** The raw x-datadog-tags value, parsed in {@link #build()} once the priority is known. */
     private String datadogTags;
 
+    /**
+     * The x-datadog-trace-id value and its parsed low order bits. The trace id is built in {@link
+     * #build()}, once x-datadog-tags may have given its high order bits.
+     */
+    private String traceIdString;
+
+    private long traceIdLow;
+
     private DatadogContextInterpreter(Config config) {
       super(config);
       isAwsPropagationEnabled = config.isAwsPropagationEnabled();
@@ -127,6 +136,8 @@ public TracePropagationStyle style() {
     @Override
     public ContextInterpreter reset(TraceConfig traceConfig) {
       datadogTags = null;
+      traceIdString = null;
+      traceIdLow = 0;
       return super.reset(traceConfig);
     }
 
@@ -181,7 +192,10 @@ public boolean accept(String key, String value) {
           if (null != value) {
             switch (classification) {
               case TRACE_ID:
-                traceId = DDTraceId.from(firstHeaderValue(value));
+                // parsed as DDTraceId.from does, so an invalid id fails here as before
+                String traceIdValue = firstHeaderValue(value);
+                traceIdLow = LongStringUtils.parseUnsignedLong(traceIdValue);
+                traceIdString = traceIdValue;
                 break;
               case SPAN_ID:
                 spanId = DDSpanId.from(firstHeaderValue(value));
@@ -235,7 +249,7 @@ protected TagContext build() {
           log.debug("Exception when extracting context", e);
         }
       }
-      restore128bTraceId();
+      buildTraceId();
       return super.build();
     }
 
@@ -248,14 +262,18 @@ private long extractEndToEndStartTime(String value) {
       }
     }
 
-    private void restore128bTraceId() {
+    private void buildTraceId() {
+      if (traceIdString == null) {
+        return;
+      }
       long highOrderBits;
-      // Check if the low-order 64 bits of the TraceId, and propagation tags were parsed
-      if (traceId != DDTraceId.ZERO
+      if (traceIdLow != 0
           && propagationTags != null
           && (highOrderBits = propagationTags.getTraceIdHighOrderBits()) != 0) {
         // Restore the 128-bit TraceId
-        traceId = DD128bTraceId.from(highOrderBits, traceId.toLong());
+        traceId = DD128bTraceId.from(highOrderBits, traceIdLow);
+      } else {
+        traceId = DDTraceId.from(traceIdLow, traceIdString);
       }
     }
   }

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

@@ -233,7 +233,7 @@ public <C> TagContext extract(
           if (merge == null) {
             continue;
           }
-          if (traceIdMatch(context.getTraceId(), merge.traceId)) {
+          if (traceIdMatch(context.getTraceId(), merge.traceIdHigh, merge.traceIdLow)) {
             applyTraceContextToFirstContext(context, merge, extractionCache);
             continue;
           }
@@ -366,6 +366,14 @@ public void forEachKey(ExtractionCache<?> carrier, AgentPropagation.KeyClassifie
     }
   }
 
+  /** {@link #traceIdMatch(DDTraceId, DDTraceId)} against a 128-bit id given by its parts. */
+  private static boolean traceIdMatch(DDTraceId a, long bHigh, long bLow) {
+    if (a instanceof DD128bTraceId) {
+      return a.toHighOrderLong() == bHigh && a.toLong() == bLow;
+    }
+    return a.toLong() == bLow;
+  }
+
   /**
    * Checks if trace identifier matches, even if they are not encoded using the same size (64-bit vs
    * 128-bit).

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

@@ -13,7 +13,6 @@
 import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTags;
-import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.api.internal.VisibleForTesting;
@@ -135,16 +134,24 @@ static boolean overridesTracestatePriority(int traceparentPriority, int tracesta
 
   /** What merging a trace context into a context extracted before it reads from it. */
   static final class TraceContextMerge {
-    final DDTraceId traceId;
+    /** The 128-bit trace id from the traceparent, as its high and low order bits. */
+    final long traceIdHigh;
+
+    final long traceIdLow;
     final long spanId;
     final W3CTracestate tracestate;
 
     /** The traceparent priority when it overrides the tracestate's, otherwise UNSET. */
     final int overridingPriority;
 
     private TraceContextMerge(
-        DDTraceId traceId, long spanId, W3CTracestate tracestate, int overridingPriority) {
-      this.traceId = traceId;
+        long traceIdHigh,
+        long traceIdLow,
+        long spanId,
+        W3CTracestate tracestate,
+        int overridingPriority) {
+      this.traceIdHigh = traceIdHigh;
+      this.traceIdLow = traceIdLow;
       this.spanId = spanId;
       this.tracestate = tracestate;
       this.overridingPriority = overridingPriority;
@@ -330,7 +337,20 @@ protected TagContext build() {
       return super.build();
     }
 
+    /** The high and low order bits of the trace id {@link #parseTraceParent} last parsed. */
+    private long traceIdHigh;
+
+    private long traceIdLow;
+
     void parseTraceParentHeader(String tp) {
+      parseTraceParent(tp, true);
+    }
+
+    /**
+     * Validates the traceparent and reads its parts. The trace id is built only if {@code
+     * buildTraceId}; its high and low order bits are read either way.
+     */
+    private void parseTraceParent(String tp, boolean buildTraceId) {
       int length = tp == null ? 0 : tp.length();
       if (length < TRACE_PARENT_LENGTH) {
         throw new IllegalStateException("The length of traceparent '" + tp + "' is too short");
@@ -341,13 +361,26 @@ void parseTraceParentHeader(String tp) {
       } else if (version == 0 && length > TRACE_PARENT_LENGTH) {
         throw new IllegalStateException("The length of traceparent '" + tp + "' is too long");
       }
-      DDTraceId traceId = DD128bTraceId.fromHex(tp, TRACE_PARENT_TID_START, 32, true);
-      if (traceId.toLong() == 0) {
+      DD128bTraceId traceId = null;
+      if (buildTraceId) {
+        traceId = DD128bTraceId.fromHex(tp, TRACE_PARENT_TID_START, 32, true);
+        this.traceIdHigh = traceId.toHighOrderLong();
+        this.traceIdLow = traceId.toLong();
+      } else {
+        // as DD128bTraceId.fromHex(tp, TRACE_PARENT_TID_START, 32, true) parses them
+        this.traceIdHigh =
+            LongStringUtils.parseUnsignedLongHex(tp, TRACE_PARENT_TID_START, 16, true);
+        this.traceIdLow =
+            LongStringUtils.parseUnsignedLongHex(tp, TRACE_PARENT_TID_START + 16, 16, true);
+      }
+      if (traceIdLow == 0) {
         throw new IllegalStateException(
             "Illegal all zero 64 bit trace id "
                 + tp.substring(TRACE_PARENT_TID_START, TRACE_PARENT_TID_END));
       }
-      this.traceId = traceId;
+      if (traceId != null) {
+        this.traceId = traceId;
+      }
       this.spanId = DDSpanId.fromHex(tp, TRACE_PARENT_SID_START, 16, true);
       if (this.spanId == 0) {
         throw new IllegalStateException(
@@ -379,20 +412,18 @@ TraceContextMerge buildForMerge() {
           throw new IllegalStateException(
               "Found no traceparent header but tracestate header '" + tracestateHeader + "'");
         }
-        parseTraceParentHeader(traceparentHeader);
+        // the trace id only needs to be compared, so it's kept as its parts
+        parseTraceParent(traceparentHeader, false);
         tracestate = W3CTracestate.parse(tracestateHeader);
       } catch (RuntimeException e) {
         log.debug("Exception when building context", e);
         return null;
       }
-      if (DDTraceId.ZERO.equals(traceId)) {
-        return null;
-      }
       int overridingPriority =
           overridesTracestatePriority(samplingPriority, tracestate.getSamplingPriority())
               ? samplingPriority
               : PrioritySampling.UNSET;
-      return new TraceContextMerge(traceId, spanId, tracestate, overridingPriority);
+      return new TraceContextMerge(traceIdHigh, traceIdLow, spanId, tracestate, overridingPriority);
     }
 
     void parseTraceStateHeader(String tracestate) {

dd-trace-core/src/test/java/datadog/trace/core/propagation/CompoundExtractionGoldenTest.java

@@ -53,7 +53,15 @@ class CompoundExtractionGoldenTest extends DDCoreJavaSpecification {
   private static final String DD_SAME_PARENT = Long.toString(Long.parseLong(PARENT, 16));
   private static final String RV = "c8b1e2d3f4a596";
 
-  private static final String[] DD_PRESENCE = {"none", "match", "matchOtherParent", "mismatch"};
+  private static final String[] DD_PRESENCE = {
+    "none",
+    "match",
+    "matchOtherParent",
+    "mismatch",
+    "invalidTraceId",
+    "negativeTraceId",
+    "zeroTraceId"
+  };
 
   private static final String[] DD_PRIORITY_FULL = {null, "0", "1", "2"};
   private static final String[] DD_PRIORITY = {"1", "0"};
@@ -169,7 +177,7 @@ private static List<Map<String, String>> corpus(boolean full) {
               }
               Map<String, String> headers = new LinkedHashMap<>();
               if (!"none".equals(presence)) {
-                headers.put("x-datadog-trace-id", "mismatch".equals(presence) ? "42" : DD_TRACE_ID);
+                headers.put("x-datadog-trace-id", datadogTraceId(presence));
                 headers.put(
                     "x-datadog-parent-id",
                     "matchOtherParent".equals(presence) ? "7" : DD_SAME_PARENT);
@@ -187,6 +195,21 @@ private static List<Map<String, String>> corpus(boolean full) {
     return corpus;
   }
 
+  private static String datadogTraceId(String presence) {
+    switch (presence) {
+      case "mismatch":
+        return "42";
+      case "invalidTraceId":
+        return "abc";
+      case "negativeTraceId":
+        return "-1";
+      case "zeroTraceId":
+        return "0";
+      default:
+        return DD_TRACE_ID;
+    }
+  }
+
   private static String describe(CoreTracer tracer, Map<String, String> headers) {
     Context extracted =
         Propagators.defaultPropagator()