feat(dsm): Implement DSM context injection

Author: PerfectSlayer

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

@@ -0,0 +1,46 @@
+package datadog.trace.core.datastreams;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import datadog.context.ImplicitContextKeyed;
+import java.util.LinkedHashMap;
+
+public class DataStreamContext implements ImplicitContextKeyed {
+  private static final ContextKey<DataStreamContext> CONTEXT_KEY =
+      ContextKey.named("dsm-context-key");
+
+  final LinkedHashMap<String, String> sortedTags;
+  final long defaultTimestamp;
+  final long payloadSizeBytes;
+  final boolean sendCheckpoint;
+
+  public static DataStreamContext fromContext(Context context) {
+    return context.get(CONTEXT_KEY);
+  }
+
+  public static DataStreamContext fromTags(LinkedHashMap<String, String> sortedTags) {
+    return new DataStreamContext(sortedTags, 0, 0, true);
+  }
+
+  public static DataStreamContext fromTagsWithoutCheckpoint(
+      LinkedHashMap<String, String> sortedTags) {
+    return new DataStreamContext(sortedTags, 0, 0, false);
+  }
+
+  // That's basically a record for now
+  private DataStreamContext(
+      LinkedHashMap<String, String> sortedTags,
+      long defaultTimestamp,
+      long payloadSizeBytes,
+      boolean sendCheckpoint) {
+    this.sortedTags = sortedTags;
+    this.defaultTimestamp = defaultTimestamp;
+    this.payloadSizeBytes = payloadSizeBytes;
+    this.sendCheckpoint = sendCheckpoint;
+  }
+
+  @Override
+  public Context storeInto(Context context) {
+    return context.with(CONTEXT_KEY, this);
+  }
+}

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

@@ -1,5 +1,8 @@
 package datadog.trace.core.datastreams;
 
+import static datadog.trace.api.DDTags.PATHWAY_HASH;
+import static datadog.trace.bootstrap.instrumentation.api.PathwayContext.PROPAGATION_KEY_BASE64;
+
 import datadog.context.Context;
 import datadog.context.propagation.CarrierSetter;
 import datadog.context.propagation.CarrierVisitor;
@@ -10,23 +13,27 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import java.io.IOException;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import javax.annotation.ParametersAreNonnullByDefault;
 
 // TODO Javadoc
 @ParametersAreNonnullByDefault
 public class DataStreamPropagator implements Propagator {
+  private final DataStreamsMonitoring dataStreamsMonitoring;
   private final Supplier<TraceConfig> traceConfigSupplier;
   private final TimeSource timeSource;
   private final long hashOfKnownTags;
   private final String serviceNameOverride;
 
   public DataStreamPropagator(
+      DataStreamsMonitoring dataStreamsMonitoring,
       Supplier<TraceConfig> traceConfigSupplier,
       TimeSource timeSource,
       long hashOfKnownTags,
       String serviceNameOverride) {
+    this.dataStreamsMonitoring = dataStreamsMonitoring;
     this.traceConfigSupplier = traceConfigSupplier;
     this.timeSource = timeSource;
     this.hashOfKnownTags = hashOfKnownTags;
@@ -35,12 +42,48 @@ public DataStreamPropagator(
 
   @Override
   public <C> void inject(Context context, C carrier, CarrierSetter<C> setter) {
-    // TODO Still in CorePropagation, not migrated yet
+    // TODO Pathway context needs to be stored into its own context element instead of span context
+    AgentSpan span = AgentSpan.fromContext(context);
+    DataStreamContext dsmContext = DataStreamContext.fromContext(context);
+    PathwayContext pathwayContext;
+    if (span == null
+        || (pathwayContext = span.context().getPathwayContext()) == null
+        || dsmContext == null) {
+      return;
+    }
+
+    // TODO Allow set checkpoint to use DsmContext as parameter?
+    pathwayContext.setCheckpoint(
+        dsmContext.sortedTags,
+        dsmContext.sendCheckpoint ? dataStreamsMonitoring::add : pathwayContext::saveStats,
+        dsmContext.defaultTimestamp,
+        dsmContext.payloadSizeBytes);
+
+    boolean injected = injectPathwayContext(pathwayContext, carrier, setter);
+
+    if (injected && pathwayContext.getHash() != 0) {
+      span.setTag(PATHWAY_HASH, Long.toUnsignedString(pathwayContext.getHash()));
+    }
+  }
+
+  private <C> boolean injectPathwayContext(
+      PathwayContext pathwayContext, C carrier, CarrierSetter<C> setter) {
+    try {
+      String encodedContext = pathwayContext.encode();
+      if (encodedContext != null) {
+        //        LOGGER.debug("Injecting pathway context {}", pathwayContext);
+        setter.set(carrier, PROPAGATION_KEY_BASE64, encodedContext);
+        return true;
+      }
+    } catch (IOException e) {
+      //      LOGGER.debug("Unable to set encode pathway context", e);
+    }
+    return false;
   }
 
   @Override
   public <C> Context extract(Context context, C carrier, CarrierVisitor<C> visitor) {
-    // TODO Pathway context needs to be stored into its own context element
+    // TODO Pathway context needs to be stored into its own context element instead of span context
     // Get span context to store pathway context into
     TagContext spanContext = getSpanContextOrNull(context);
     PathwayContext pathwayContext;

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

@@ -204,7 +204,11 @@ public PathwayContext newPathwayContext() {
   @Override
   public Propagator propagator() {
     return new DataStreamPropagator(
-        this.traceConfigSupplier, this.timeSource, this.hashOfKnownTags, getThreadServiceName());
+        this,
+        this.traceConfigSupplier,
+        this.timeSource,
+        this.hashOfKnownTags,
+        getThreadServiceName());
   }
 
   @Override