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

Author: PerfectSlayer

dd-trace-core/build.gradle

@@ -48,6 +48,10 @@ excludedClassesCoverage += [
   'datadog.trace.core.TracingConfigPoller.Updater',
   // covered with dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/CheckpointerTest.groovy
   'datadog.trace.core.datastreams.DefaultDataStreamsMonitoring',
+  // TODO CorePropagation will be removed during context refactoring
+  'datadog.trace.core.propagation.CorePropagation',
+  // TODO DSM propagator will be tested once fully migrated
+  'datadog.trace.core.datastreams.DataStreamPropagator'
 ]
 
 addTestSuite('traceAgentTest')

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;
@@ -713,22 +714,25 @@ 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;
+    TracingPropagator tracingPropagator = new TracingPropagator(injector, tracingExtractor);
     // Check if standalone AppSec is enabled:
     // If enabled, use the standalone AppSec propagator by default that will limit tracing concern
     // injection and delegate to the tracing propagator if needed,
     // If disabled, the most common case, use the usual tracing propagator by default.
     boolean standaloneAppSec = config.isAppSecStandaloneEnabled();
+    boolean dsm = config.isDataStreamsEnabled();
     Propagators.register(STANDALONE_ASM_CONCERN, new StandaloneAsmPropagator(), standaloneAppSec);
-    Propagators.register(
-        TRACING_CONCERN, new TracingPropagator(injector, extractor), !standaloneAppSec);
+    Propagators.register(TRACING_CONCERN, tracingPropagator, !standaloneAppSec);
     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,83 @@
+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.TraceConfig;
+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.PathwayContext;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+// TODO Javadoc
+@ParametersAreNonnullByDefault
+public class DataStreamPropagator implements Propagator {
+  private final Supplier<TraceConfig> traceConfigSupplier;
+  private final TimeSource timeSource;
+  private final long hashOfKnownTags;
+  private final String serviceNameOverride;
+
+  public DataStreamPropagator(
+      Supplier<TraceConfig> traceConfigSupplier,
+      TimeSource timeSource,
+      long hashOfKnownTags,
+      String serviceNameOverride) {
+    this.traceConfigSupplier = traceConfigSupplier;
+    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) {
+    // TODO Pathway context needs to be stored into its own context element
+    // Get span context to store pathway context into
+    TagContext spanContext = getSpanContextOrNull(context);
+    PathwayContext pathwayContext;
+    // Ensure if DSM is enabled and look for pathway context
+    if (isDsmEnabled(spanContext)
+        && (pathwayContext = extractDsmPathwayContext(carrier, visitor)) != null) {
+      // Store pathway context into span context
+      if (spanContext == null) {
+        spanContext = new TagContext();
+        AgentSpan span = AgentSpan.fromSpanContext(spanContext);
+        context = Context.root().with(span);
+      }
+      spanContext.withPathwayContext(pathwayContext);
+    }
+    return context;
+  }
+
+  private TagContext getSpanContextOrNull(Context context) {
+    AgentSpan extractedSpan = AgentSpan.fromContext(context);
+    AgentSpanContext extractedSpanContext;
+    if (extractedSpan != null
+        && (extractedSpanContext = extractedSpan.context()) instanceof TagContext) {
+      return (TagContext) extractedSpanContext;
+    }
+    return null;
+  }
+
+  private boolean isDsmEnabled(@Nullable TagContext tagContext) {
+    TraceConfig traceConfig = tagContext == null ? null : tagContext.getTraceConfig();
+    if (traceConfig == null) {
+      traceConfig = this.traceConfigSupplier.get();
+    }
+    return traceConfig.isDataStreamsEnabled();
+  }
+
+  private <C> PathwayContext extractDsmPathwayContext(C carrier, CarrierVisitor<C> visitor) {
+    return DefaultPathwayContext.extract(
+        carrier, visitor, this.timeSource, this.hashOfKnownTags, this.serviceNameOverride);
+  }
+}

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,12 @@ public PathwayContext newPathwayContext() {
     }
   }
 
+  @Override
+  public Propagator propagator() {
+    return new DataStreamPropagator(
+        this.traceConfigSupplier, 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

@@ -7,18 +7,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;
   }