feat(core): Move DSM extractor to its own propagator

Author: PerfectSlayer

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

@@ -5,6 +5,7 @@
 import static datadog.trace.api.DDTags.DJM_ENABLED;
 import static datadog.trace.api.DDTags.DSM_ENABLED;
 import static datadog.trace.api.DDTags.PROFILING_CONTEXT_ENGINE;
+import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.STANDALONE_ASM_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.TRACING_CONCERN;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.XRAY_TRACING_CONCERN;
@@ -712,17 +713,21 @@ private CoreTracer(
 
     sharedCommunicationObjects.whenReady(this.dataStreamsMonitoring::start);
 
-    // Create default extractor from config if not provided and decorate it with DSM extractor
-    HttpCodec.Extractor builtExtractor =
-        extractor == null ? HttpCodec.createExtractor(config, this::captureTraceConfig) : extractor;
-    builtExtractor = this.dataStreamsMonitoring.extractor(builtExtractor);
-    // Store all propagators to propagation
-    this.propagation = new CorePropagation(builtExtractor, this.dataStreamsMonitoring.injector());
+    // Store all propagators to propagation -- only DSM injection left
+    this.propagation = new CorePropagation(this.dataStreamsMonitoring.injector());
 
+    // Register context propagators
+    HttpCodec.Extractor tracingExtractor =
+        extractor == null ? HttpCodec.createExtractor(config, this::captureTraceConfig) : extractor;
     boolean appSec = config.isAppSecStandaloneEnabled();
+    boolean dsm = config.isDataStreamsEnabled();
     Propagators.register(STANDALONE_ASM_CONCERN, new StandaloneAsmPropagator(), appSec);
-    Propagators.register(TRACING_CONCERN, new TracingPropagator(injector, extractor), !appSec);
+    Propagators.register(
+        TRACING_CONCERN, new TracingPropagator(injector, tracingExtractor), !appSec);
     Propagators.register(XRAY_TRACING_CONCERN, new XRayPropagator(config), false);
+    if (dsm) {
+      Propagators.register(DSM_CONCERN, this.dataStreamsMonitoring.propagator());
+    }
 
     this.tagInterceptor =
         null == tagInterceptor ? new TagInterceptor(new RuleFlags(config)) : tagInterceptor;

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamPropagator.java

@@ -0,0 +1,58 @@
+package datadog.trace.core.datastreams;
+
+import datadog.context.Context;
+import datadog.context.propagation.CarrierSetter;
+import datadog.context.propagation.CarrierVisitor;
+import datadog.context.propagation.Propagator;
+import datadog.trace.api.time.TimeSource;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+// TODO Javadoc
+@ParametersAreNonnullByDefault
+public class DataStreamPropagator implements Propagator {
+  private final TimeSource timeSource;
+  private final long hashOfKnownTags;
+  private final String serviceNameOverride;
+
+  public DataStreamPropagator(
+      TimeSource timeSource, long hashOfKnownTags, String serviceNameOverride) {
+    this.timeSource = timeSource;
+    this.hashOfKnownTags = hashOfKnownTags;
+    this.serviceNameOverride = serviceNameOverride;
+  }
+
+  @Override
+  public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
+    // TODO Still in CorePropagation, not migrated yet
+  }
+
+  @Override
+  public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
+    // Extract pathway context
+    DefaultPathwayContext pathwayContext =
+        DefaultPathwayContext.extract(
+            carrier, visitor, this.timeSource, this.hashOfKnownTags, this.serviceNameOverride);
+    if (pathwayContext == null) {
+      return context;
+    }
+    // TODO Pathway context needs to be store into its own context element
+    // Get span context to store pathway context into
+    TagContext spanContext;
+    AgentSpan extractedSpan = AgentSpan.fromContext(context);
+    AgentSpanContext extractedSpanContext;
+    if (extractedSpan != null
+        && (extractedSpanContext = extractedSpan.context()) instanceof TagContext) {
+      spanContext = (TagContext) extractedSpanContext;
+    } else {
+      spanContext = new TagContext();
+      AgentSpan span = AgentSpan.fromSpanContext(spanContext);
+      context = Context.root().with(span);
+    }
+    spanContext.withPathwayContext(pathwayContext);
+    // Return the original context, with mutated span context
+    return context;
+  }
+}

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DataStreamsMonitoring.java

@@ -1,5 +1,6 @@
 package datadog.trace.core.datastreams;
 
+import datadog.context.propagation.Propagator;
 import datadog.trace.api.experimental.DataStreamsContextCarrier;
 import datadog.trace.bootstrap.instrumentation.api.AgentDataStreamsMonitoring;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -10,6 +11,13 @@
 public interface DataStreamsMonitoring extends AgentDataStreamsMonitoring, AutoCloseable {
   void start();
 
+  /**
+   * Gets the propagator for DSM concern.
+   *
+   * @return The propagator for DSM concern.
+   */
+  Propagator propagator();
+
   /**
    * Get a context extractor that support {@link PathwayContext} extraction.
    *

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java

@@ -14,6 +14,7 @@
 
 import datadog.communication.ddagent.DDAgentFeaturesDiscovery;
 import datadog.communication.ddagent.SharedCommunicationObjects;
+import datadog.context.propagation.Propagator;
 import datadog.trace.api.Config;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.WellKnownTags;
@@ -200,6 +201,11 @@ public PathwayContext newPathwayContext() {
     }
   }
 
+  @Override
+  public Propagator propagator() {
+    return new DataStreamPropagator(this.timeSource, this.hashOfKnownTags, getThreadServiceName());
+  }
+
   @Override
   public HttpCodec.Extractor extractor(HttpCodec.Extractor delegate) {
     return new DataStreamContextExtractor(

dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultPathwayContext.java

@@ -7,10 +7,10 @@
 import com.datadoghq.sketch.ddsketch.encoding.ByteArrayInput;
 import com.datadoghq.sketch.ddsketch.encoding.GrowingByteArrayOutput;
 import com.datadoghq.sketch.ddsketch.encoding.VarEncodingHelper;
+import datadog.context.propagation.CarrierVisitor;
 import datadog.trace.api.Config;
 import datadog.trace.api.WellKnownTags;
 import datadog.trace.api.time.TimeSource;
-import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
 import datadog.trace.bootstrap.instrumentation.api.StatsPoint;
 import datadog.trace.util.FNV64Hash;
@@ -26,6 +26,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -267,7 +268,7 @@ public String toString() {
     }
   }
 
-  private static class PathwayContextExtractor implements AgentPropagation.KeyClassifier {
+  private static class PathwayContextExtractor implements BiConsumer<String, String> {
     private final TimeSource timeSource;
     private final long hashOfKnownTags;
     private final String serviceNameOverride;
@@ -281,27 +282,25 @@ private static class PathwayContextExtractor implements AgentPropagation.KeyClas
     }
 
     @Override
-    public boolean accept(String key, String value) {
+    public void accept(String key, String value) {
       if (PROPAGATION_KEY_BASE64.equalsIgnoreCase(key)) {
         try {
           extractedContext = decode(timeSource, hashOfKnownTags, serviceNameOverride, value);
-        } catch (IOException e) {
-          return false;
+        } catch (IOException ignored) {
         }
       }
-      return true;
     }
   }
 
   static <C> DefaultPathwayContext extract(
       C carrier,
-      AgentPropagation.ContextVisitor<C> getter,
+      CarrierVisitor<C> getter,
       TimeSource timeSource,
       long hashOfKnownTags,
       String serviceNameOverride) {
     PathwayContextExtractor pathwayContextExtractor =
         new PathwayContextExtractor(timeSource, hashOfKnownTags, serviceNameOverride);
-    getter.forEachKey(carrier, pathwayContextExtractor);
+    getter.forEachKeyValue(carrier, pathwayContextExtractor);
     if (pathwayContextExtractor.extractedContext == null) {
       log.debug("No context extracted");
     } else {

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

@@ -1,5 +1,7 @@
 package datadog.trace.core.propagation;
 
+import datadog.context.Context;
+import datadog.context.propagation.Propagators;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -8,18 +10,14 @@
 
 public class CorePropagation implements AgentPropagation {
   private final DataStreamContextInjector dataStreamContextInjector;
-  private final HttpCodec.Extractor extractor;
 
   /**
    * Constructor
    *
-   * @param extractor The context extractor.
    * @param dataStreamContextInjector The DSM context injector, as a specific object until generic
    *     context injection is available.
    */
-  public CorePropagation(
-      HttpCodec.Extractor extractor, DataStreamContextInjector dataStreamContextInjector) {
-    this.extractor = extractor;
+  public CorePropagation(DataStreamContextInjector dataStreamContextInjector) {
     this.dataStreamContextInjector = dataStreamContextInjector;
   }
 
@@ -50,6 +48,8 @@ public <C> void injectPathwayContextWithoutSendingStats(
 
   @Override
   public <C> AgentSpanContext.Extracted extract(final C carrier, final ContextVisitor<C> getter) {
-    return extractor.extract(carrier, getter);
+    Context extracted = Propagators.defaultPropagator().extract(Context.root(), carrier, getter);
+    AgentSpan extractedSpan = AgentSpan.fromContext(extracted);
+    return extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
   }
 }