Merge a trace context without building its propagation tags

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

@@ -223,7 +223,25 @@ public <C> TagContext extract(
       ExtractionCache<C> extractionCache = new ExtractionCache<>(carrier, getter);
 
       for (final Extractor extractor : this.extractors) {
-        TagContext extracted = extractor.extract(extractionCache, extractionCache);
+        TagContext extracted;
+        if (context != null && extractor instanceof W3CHttpCodec.TraceContextExtractor) {
+          // Merging a trace context only needs a few of its parts, so skip building it whole
+          W3CHttpCodec.TraceContextExtractor traceContextExtractor =
+              (W3CHttpCodec.TraceContextExtractor) extractor;
+          W3CHttpCodec.TraceContextMerge merge =
+              traceContextExtractor.extractForMerge(extractionCache, extractionCache);
+          if (merge == null) {
+            continue;
+          }
+          if (traceIdMatch(context.getTraceId(), merge.traceId)) {
+            applyTraceContextToFirstContext(context, merge, extractionCache);
+            continue;
+          }
+          // A mismatched trace context becomes a span link, which needs all of it
+          extracted = traceContextExtractor.completeExtraction();
+        } else {
+          extracted = extractor.extract(extractionCache, extractionCache);
+        }
         // Check if context is valid
         if (extracted instanceof ExtractedContext) {
           ExtractedContext extractedContext = (ExtractedContext) extracted;
@@ -237,12 +255,7 @@ public <C> TagContext extract(
           }
           // If another valid context is extracted
           else {
-            if (traceIdMatch(context.getTraceId(), extractedContext.getTraceId())) {
-              boolean comingFromTraceContext = extracted.getPropagationStyle() == TRACECONTEXT;
-              if (comingFromTraceContext) {
-                applyTraceContextToFirstContext(context, extractedContext, extractionCache);
-              }
-            } else {
+            if (!traceIdMatch(context.getTraceId(), extractedContext.getTraceId())) {
               // Terminate extracted context and add it as span link
               context.addTerminatedSpanLink(
                   DDSpanLink.from(
@@ -283,16 +296,18 @@ else if (extracted != null && partialContext == null) {
      */
     private <C> void applyTraceContextToFirstContext(
         ExtractedContext firstContext,
-        ExtractedContext traceContext,
+        W3CHttpCodec.TraceContextMerge traceContext,
         ExtractionCache<C> extractionCache) {
       // Propagate newly extracted W3C tracestate to first valid context
-      firstContext.getPropagationTags().updateW3CTracestateFrom(traceContext.getPropagationTags());
+      firstContext
+          .getPropagationTags()
+          .updateW3CTracestateFrom(traceContext.tracestate, traceContext.overridingPriority);
       // Check if parent spans differ to reconcile them
-      if (firstContext.getSpanId() != traceContext.getSpanId()) {
+      if (firstContext.getSpanId() != traceContext.spanId) {
         // Override parent span id with W3C one
-        firstContext.overrideSpanId(traceContext.getSpanId());
+        firstContext.overrideSpanId(traceContext.spanId);
         // Add last parent ID as a span tag (check W3C first, else Datadog)
-        CharSequence lastParentId = traceContext.getPropagationTags().getLastParentId();
+        CharSequence lastParentId = traceContext.tracestate.getLastParentId();
         if (lastParentId == null) {
           lastParentId = extractionCache.getDatadogSpanIdHex();
         }

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

@@ -7,6 +7,7 @@
 import datadog.trace.api.ProductTraceSource;
 import datadog.trace.api.llmobs.LLMObsPropagationValues;
 import datadog.trace.core.propagation.ptags.PTagsFactory;
+import datadog.trace.core.propagation.ptags.W3CTracestate;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -188,10 +189,14 @@ public abstract boolean tryUpdateProbabilitySamplingDecision(
    */
   public abstract void updateW3CTracestate(String tracestate);
 
-  /** Updates the original W3C tracestate header from {@code source}. */
-  public void updateW3CTracestateFrom(PropagationTags source) {
-    updateW3CTracestate(source.getW3CTracestate());
-  }
+  /**
+   * Takes the W3C tracestate from a trace context extracted alongside this one.
+   *
+   * @param source the trace context's tracestate
+   * @param overridingPriority the trace context's priority when it overrides the tracestate's own
+   *     ({@code s}), otherwise {@link datadog.trace.api.sampling.PrioritySampling#UNSET}
+   */
+  public abstract void updateW3CTracestateFrom(W3CTracestate source, int overridingPriority);
 
   /**
    * Constructs a header value that includes valid propagated _dd.p.* tags and possibly a new

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

@@ -18,9 +18,20 @@ public TagContextExtractor(
 
   @Override
   public <C> TagContext extract(final C carrier, final AgentPropagation.ContextVisitor<C> getter) {
+    return interpret(carrier, getter).build();
+  }
+
+  /** This thread's interpreter, as the last {@link #interpret} left it. */
+  protected final ContextInterpreter interpreter() {
+    return this.ctxInterpreter.get();
+  }
+
+  /** Runs this thread's interpreter over the carrier, ready to build. */
+  protected final <C> ContextInterpreter interpret(
+      final C carrier, final AgentPropagation.ContextVisitor<C> getter) {
     ContextInterpreter interpreter = this.ctxInterpreter.get().reset(traceConfigSupplier.get());
     getter.forEachKey(carrier, interpreter);
-    return interpreter.build();
+    return interpreter;
   }
 
   @Override

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

@@ -22,9 +22,11 @@
 import datadog.trace.api.propagation.W3CTraceParent;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.sampling.SamplingMechanism;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.core.DDSpanContext;
 import datadog.trace.core.propagation.PropagationTags.SamplingState;
+import datadog.trace.core.propagation.ptags.W3CTracestate;
 import java.util.Map;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
@@ -118,7 +120,55 @@ private <C> void injectBaggage(DDSpanContext context, C carrier, CarrierSetter<C
 
   public static HttpCodec.Extractor newExtractor(
       Config config, Supplier<TraceConfig> traceConfigSupplier) {
-    return new TagContextExtractor(traceConfigSupplier, () -> new W3CContextInterpreter(config));
+    return new TraceContextExtractor(traceConfigSupplier, config);
+  }
+
+  /**
+   * Whether the trace context's priority, from the {@code traceparent} sampled flag, overrides the
+   * tracestate's own Datadog priority ({@code s}).
+   */
+  static boolean overridesTracestatePriority(int traceparentPriority, int tracestatePriority) {
+    return (traceparentPriority == SAMPLER_DROP && tracestatePriority > 0)
+        || (traceparentPriority == SAMPLER_KEEP && tracestatePriority <= 0)
+        || tracestatePriority == PrioritySampling.UNSET;
+  }
+
+  /** What merging a trace context into a context extracted before it reads from it. */
+  static final class TraceContextMerge {
+    final DDTraceId traceId;
+    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;
+      this.spanId = spanId;
+      this.tracestate = tracestate;
+      this.overridingPriority = overridingPriority;
+    }
+  }
+
+  static final class TraceContextExtractor extends TagContextExtractor {
+    private TraceContextExtractor(Supplier<TraceConfig> traceConfigSupplier, Config config) {
+      super(traceConfigSupplier, () -> new W3CContextInterpreter(config));
+    }
+
+    /**
+     * Extracts only what merging into another extracted context needs, without building propagation
+     * tags. Returns null where {@link #extract} wouldn't return an {@link ExtractedContext}.
+     */
+    <C> TraceContextMerge extractForMerge(
+        final C carrier, final AgentPropagation.ContextVisitor<C> getter) {
+      return ((W3CContextInterpreter) interpret(carrier, getter)).buildForMerge();
+    }
+
+    /** Builds the whole context that {@link #extractForMerge} last returned the merge parts of. */
+    TagContext completeExtraction() {
+      return interpreter().build();
+    }
   }
 
   private static class W3CContextInterpreter extends ContextInterpreter {
@@ -315,6 +365,36 @@ void parseTraceParentHeader(String tp) {
       }
     }
 
+    /** Like {@link #build()}, but for {@link TraceContextExtractor#extractForMerge}. */
+    TraceContextMerge buildForMerge() {
+      if (traceparentHeader == null && tracestateHeader == null) {
+        onlyTagContext();
+      }
+      if (!valid || !fullContext) {
+        return null;
+      }
+      W3CTracestate tracestate;
+      try {
+        if (traceparentHeader == null) {
+          throw new IllegalStateException(
+              "Found no traceparent header but tracestate header '" + tracestateHeader + "'");
+        }
+        parseTraceParentHeader(traceparentHeader);
+        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);
+    }
+
     void parseTraceStateHeader(String tracestate) {
       if (tracestate == null || tracestate.isEmpty()) {
         this.propagationTags = this.propagationTagsFactory.empty();
@@ -324,9 +404,7 @@ void parseTraceStateHeader(String tracestate) {
       this.lastParentId = this.propagationTags.getLastParentId();
       int ptagsPriority = this.propagationTags.getSamplingPriority();
       int contextPriority = this.samplingPriority;
-      if ((contextPriority == SAMPLER_DROP && ptagsPriority > 0)
-          || (contextPriority == SAMPLER_KEEP && ptagsPriority <= 0)
-          || ptagsPriority == PrioritySampling.UNSET) {
+      if (overridesTracestatePriority(contextPriority, ptagsPriority)) {
         // Override Datadog sampling priority with W3C one
         this.propagationTags.updateTraceSamplingPriority(
             contextPriority, SamplingMechanism.EXTERNAL_OVERRIDE);

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

@@ -34,7 +34,7 @@ public class PTagsFactory implements PropagationTags.Factory {
   public PTagsFactory(int xDatadogTagsLimit) {
     this.xDatadogTagsLimit = xDatadogTagsLimit;
     DEC_ENC_MAP.put(DATADOG, new DatadogPTagsCodec(xDatadogTagsLimit));
-    DEC_ENC_MAP.put(W3C, new W3CPTagsCodec());
+    DEC_ENC_MAP.put(W3C, new W3CPTagsCodec(this));
   }
 
   boolean isPropagationTagsDisabled() {
@@ -888,19 +888,19 @@ public void updateW3CTracestate(String tracestate) {
     }
 
     @Override
-    public void updateW3CTracestateFrom(PropagationTags source) {
-      if (!(source instanceof PTags)) {
-        super.updateW3CTracestateFrom(source);
-        return;
+    public void updateW3CTracestateFrom(W3CTracestate source, int overridingPriority) {
+      OtelTraceState otelTraceState = source.otelTraceState;
+      if (overridingPriority != PrioritySampling.UNSET) {
+        // as the trace context's own priority override reconciles its member
+        otelTraceState =
+            reconcileOtelTraceState(
+                otelTraceState, overridingPriority, SamplingMechanism.EXTERNAL_OVERRIDE);
       }
-      PTags sourcePTags = (PTags) source;
-      SamplingState sourceState = sourcePTags.samplingState();
-      CharSequence sourceOtelTraceState = sourceState.getOtelTraceState();
       setW3CTracestate(
-          sourceState.getTracestate(),
-          sourceOtelTraceState instanceof OtelTraceState
-              ? (OtelTraceState) sourceOtelTraceState
-              : W3CPTagsCodec.extractOtelTraceState(sourceState.getTracestate()));
+          source.tracestate,
+          otelTraceState != null
+              ? otelTraceState
+              : W3CPTagsCodec.extractOtelTraceState(source.tracestate));
     }
 
     private void setW3CTracestate(String tracestate, OtelTraceState otelTraceState) {