Initial dsm extractors configuration

Author: kr-igor

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -1,5 +1,15 @@
 package datadog.trace.agent.test
 
+import static datadog.communication.http.OkHttpUtils.buildHttpClient
+import static datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_HOST
+import static datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_TIMEOUT
+import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_AGENT_PORT
+import static datadog.trace.api.config.DebuggerConfig.DYNAMIC_INSTRUMENTATION_SNAPSHOT_URL
+import static datadog.trace.api.config.DebuggerConfig.DYNAMIC_INSTRUMENTATION_VERIFY_BYTECODE
+import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious
+import static datadog.trace.util.AgentThreadFactory.AgentThread.TASK_SCHEDULER
+
 import ch.qos.logback.classic.Level
 import ch.qos.logback.classic.util.ContextInitializer
 import com.datadog.debugger.agent.ClassesToRetransformFinder
@@ -32,6 +42,7 @@ import datadog.trace.api.TraceConfig
 import datadog.trace.api.config.GeneralConfig
 import datadog.trace.api.config.TracerConfig
 import datadog.trace.api.datastreams.AgentDataStreamsMonitoring
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor
 import datadog.trace.api.sampling.SamplingRule
 import datadog.trace.api.time.SystemTimeSource
 import datadog.trace.bootstrap.ActiveSubsystems
@@ -55,6 +66,13 @@ import de.thetaphi.forbiddenapis.SuppressForbidden
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings
 import groovy.transform.stc.ClosureParams
 import groovy.transform.stc.SimpleType
+import java.lang.instrument.ClassFileTransformer
+import java.lang.instrument.Instrumentation
+import java.nio.ByteBuffer
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeoutException
+import java.util.concurrent.atomic.AtomicInteger
 import net.bytebuddy.agent.ByteBuddyAgent
 import net.bytebuddy.agent.builder.AgentBuilder
 import net.bytebuddy.description.type.TypeDescription
@@ -68,24 +86,6 @@ import org.slf4j.LoggerFactory
 import org.spockframework.mock.MockUtil
 import spock.lang.Shared
 
-import java.lang.instrument.ClassFileTransformer
-import java.lang.instrument.Instrumentation
-import java.nio.ByteBuffer
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.TimeoutException
-import java.util.concurrent.atomic.AtomicInteger
-
-import static datadog.communication.http.OkHttpUtils.buildHttpClient
-import static datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_HOST
-import static datadog.trace.api.ConfigDefaults.DEFAULT_AGENT_TIMEOUT
-import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_AGENT_PORT
-import static datadog.trace.api.config.DebuggerConfig.DYNAMIC_INSTRUMENTATION_SNAPSHOT_URL
-import static datadog.trace.api.config.DebuggerConfig.DYNAMIC_INSTRUMENTATION_VERIFY_BYTECODE
-import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious
-import static datadog.trace.util.AgentThreadFactory.AgentThread.TASK_SCHEDULER
-
 /**
  * A specification that automatically applies instrumentation and exposes a global trace
  * writer.
@@ -249,6 +249,10 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
     List<? extends SamplingRule.TraceSamplingRule> getTraceSamplingRules() {
       return null
     }
+
+    List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors() {
+      return null
+    }
   }
 
   boolean originalAppSecRuntimeValue

dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java

@@ -88,6 +88,8 @@ public final class GeneralConfig {
   public static final String DATA_STREAMS_ENABLED = "data.streams.enabled";
   public static final String DATA_STREAMS_BUCKET_DURATION_SECONDS =
       "data.streams.bucket_duration.seconds";
+  public static final String DATA_STREAMS_TRANSACTION_EXTRACTORS =
+      "data.streams.transaction_extractors";
 
   public static final String TELEMETRY_ENABLED = "instrumentation.telemetry.enabled";
   public static final String TELEMETRY_HEARTBEAT_INTERVAL = "telemetry.heartbeat.interval";

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

@@ -86,6 +86,7 @@
 import datadog.trace.context.TraceScope;
 import datadog.trace.core.baggage.BaggagePropagator;
 import datadog.trace.core.datastreams.DataStreamsMonitoring;
+import datadog.trace.core.datastreams.DataStreamsTransactionExtractors;
 import datadog.trace.core.datastreams.DefaultDataStreamsMonitoring;
 import datadog.trace.core.histogram.Histograms;
 import datadog.trace.core.monitor.HealthMetrics;
@@ -679,6 +680,16 @@ private CoreTracer(
     } else {
       traceSamplingRules = TraceSamplingRules.deserialize(traceSamplingRulesJson);
     }
+
+    DataStreamsTransactionExtractors dataStreamsTransactionExtractors;
+    String dataStreamsTransactionExtractorsJson = config.getDataStreamsTransactionExtractors();
+    if (dataStreamsTransactionExtractorsJson == null) {
+      dataStreamsTransactionExtractors = DataStreamsTransactionExtractors.EMPTY;
+    } else {
+      dataStreamsTransactionExtractors =
+          DataStreamsTransactionExtractors.deserialize(dataStreamsTransactionExtractorsJson);
+    }
+
     // Get initial Span Sampling Rules from config
     String spanSamplingRulesJson = config.getSpanSamplingRules();
     String spanSamplingRulesFile = config.getSpanSamplingRulesFile();
@@ -708,6 +719,9 @@ private CoreTracer(
             .setSpanSamplingRules(spanSamplingRules.getRules())
             .setTraceSamplingRules(traceSamplingRules.getRules(), traceSamplingRulesJson)
             .setTracingTags(config.getMergedSpanTags())
+            .setDataStreamsTransactionExtractors(
+                dataStreamsTransactionExtractors.getExtractors(),
+                dataStreamsTransactionExtractorsJson)
             .apply();
 
     this.logs128bTraceIdEnabled = Config.get().isLogs128bitTraceIdEnabled();

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

@@ -0,0 +1,78 @@
+package datadog.trace.core.datastreams;
+
+import com.squareup.moshi.FromJson;
+import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.Moshi;
+import com.squareup.moshi.ToJson;
+import com.squareup.moshi.Types;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
+import java.lang.reflect.ParameterizedType;
+import java.util.Collections;
+import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class DataStreamsTransactionExtractors {
+  public static final DataStreamsTransactionExtractors EMPTY =
+      new DataStreamsTransactionExtractors(Collections.emptyList());
+  private static final Logger log = LoggerFactory.getLogger(DataStreamsTransactionExtractors.class);
+  private static final Moshi MOSHI =
+      new Moshi.Builder()
+          .add(new DataStreamsTransactionExtractors.DataStreamsTransactionExtractorAdapter())
+          .build();
+  private static final ParameterizedType LIST_OF_RULES =
+      Types.newParameterizedType(List.class, DataStreamsTransactionExtractor.class);
+  private static final JsonAdapter<List<DataStreamsTransactionExtractor>> LIST_OF_RULES_ADAPTER =
+      MOSHI.adapter(LIST_OF_RULES);
+
+  private final List<DataStreamsTransactionExtractor> extractors;
+
+  public DataStreamsTransactionExtractors(List<DataStreamsTransactionExtractor> extractors) {
+    this.extractors = Collections.unmodifiableList(extractors);
+  }
+
+  public static DataStreamsTransactionExtractors deserialize(String json) {
+    try {
+      return new DataStreamsTransactionExtractors(LIST_OF_RULES_ADAPTER.fromJson(json));
+    } catch (Throwable ex) {
+      log.error("Couldn't parse Data Streams Extractors from JSON: {}", json, ex);
+    }
+
+    return EMPTY;
+  }
+
+  public List<DataStreamsTransactionExtractor> getExtractors() {
+    return extractors;
+  }
+
+  private static final class JsonDataStreamsTransactionExtractor {
+    private static final JsonAdapter<JsonDataStreamsTransactionExtractor> jsonAdapter =
+        MOSHI.adapter(JsonDataStreamsTransactionExtractor.class);
+    String name;
+    String type;
+    String value;
+
+    @Override
+    public String toString() {
+      return jsonAdapter.toJson(this);
+    }
+  }
+
+  private static final class DataStreamsTransactionExtractorAdapter {
+    private static DataStreamsTransactionExtractor create(
+        JsonDataStreamsTransactionExtractor jsonExtractor) {
+      return new DataStreamsTransactionExtractor(
+          jsonExtractor.name, jsonExtractor.type, jsonExtractor.value);
+    }
+
+    @FromJson
+    DataStreamsTransactionExtractor fromJson(JsonDataStreamsTransactionExtractor jsonExtractor) {
+      return create(jsonExtractor);
+    }
+
+    @ToJson
+    JsonDataStreamsTransactionExtractor toJson(DataStreamsTransactionExtractor extractor) {
+      throw new UnsupportedOperationException();
+    }
+  }
+}

dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/DataStreamsTransactionExtractorsTest.groovy

@@ -0,0 +1,22 @@
+package datadog.trace.core.datastreams
+
+import datadog.trace.core.test.DDCoreSpecification
+
+class DataStreamsTransactionExtractorsTest extends DDCoreSpecification {
+  def "Deserialize from json"() {
+    when:
+    def list = DataStreamsTransactionExtractors.deserialize("""[
+      {"name": "extractor", "type": "http_request_header", "value": "transaction_id"},
+      {"name": "second_extractor", "type": "http_response_header", "value": "transaction_id"}
+    ]""")
+    def extractors = list.getExtractors()
+    then:
+    extractors.size() == 2
+    extractors[0].getName() == "extractor"
+    extractors[0].getType() == "http_request_header"
+    extractors[0].getValue() == "transaction_id"
+    extractors[1].getName() == "second_extractor"
+    extractors[1].getType() == "http_response_header"
+    extractors[1].getValue() == "transaction_id"
+  }
+}

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -351,6 +351,7 @@
 import static datadog.trace.api.config.GeneralConfig.DATA_JOBS_PARSE_SPARK_PLAN_ENABLED;
 import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_BUCKET_DURATION_SECONDS;
 import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_ENABLED;
+import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_TRANSACTION_EXTRACTORS;
 import static datadog.trace.api.config.GeneralConfig.DOGSTATSD_ARGS;
 import static datadog.trace.api.config.GeneralConfig.DOGSTATSD_HOST;
 import static datadog.trace.api.config.GeneralConfig.DOGSTATSD_NAMED_PIPE;
@@ -1200,6 +1201,7 @@ public static String getHostName() {
 
   private final boolean dataStreamsEnabled;
   private final float dataStreamsBucketDurationSeconds;
+  private final String dataStreamsTransactionExtractors;
 
   private final boolean serviceDiscoveryEnabled;
 
@@ -2645,6 +2647,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     dataStreamsBucketDurationSeconds =
         configProvider.getFloat(
             DATA_STREAMS_BUCKET_DURATION_SECONDS, DEFAULT_DATA_STREAMS_BUCKET_DURATION);
+    dataStreamsTransactionExtractors =
+        configProvider.getString(DATA_STREAMS_TRANSACTION_EXTRACTORS);
 
     azureAppServices = configProvider.getBoolean(AZURE_APP_SERVICES, false);
     traceAgentPath = configProvider.getString(TRACE_AGENT_PATH);
@@ -4447,6 +4451,10 @@ public float getDataStreamsBucketDurationSeconds() {
     return dataStreamsBucketDurationSeconds;
   }
 
+  public String getDataStreamsTransactionExtractors() {
+    return dataStreamsTransactionExtractors;
+  }
+
   public long getDataStreamsBucketDurationNanoseconds() {
     // Rounds to the nearest millisecond before converting to nanos
     int milliseconds = Math.round(dataStreamsBucketDurationSeconds * 1000);

internal-api/src/main/java/datadog/trace/api/DynamicConfig.java

@@ -14,6 +14,7 @@
 import static datadog.trace.util.ConfigStrings.normalizedHeaderTag;
 import static datadog.trace.util.ConfigStrings.trim;
 
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.api.sampling.SamplingRule.SpanSamplingRule;
 import datadog.trace.api.sampling.SamplingRule.TraceSamplingRule;
 import java.util.Collection;
@@ -111,6 +112,8 @@ public final class Builder {
     Double traceSampleRate;
 
     String preferredServiceName;
+    List<DataStreamsTransactionExtractor> dataStreamsTransactionExtractors;
+    String dataStreamsTransactionExtractorsJson;
 
     Builder() {}
 
@@ -134,6 +137,8 @@ public final class Builder {
       this.tracingTags = snapshot.tracingTags;
 
       this.preferredServiceName = snapshot.preferredServiceName;
+      this.dataStreamsTransactionExtractors = snapshot.dataStreamsTransactionExtractors;
+      this.dataStreamsTransactionExtractorsJson = snapshot.dataStreamsTransactionExtractorsJson;
     }
 
     public Builder setRuntimeMetricsEnabled(boolean runtimeMetricsEnabled) {
@@ -151,6 +156,14 @@ public Builder setDataStreamsEnabled(boolean dataStreamsEnabled) {
       return this;
     }
 
+    public Builder setDataStreamsTransactionExtractors(
+        List<DataStreamsTransactionExtractor> dataStreamsTransactionExtractors,
+        String dataStreamsTransactionExtractorsJson) {
+      this.dataStreamsTransactionExtractors = dataStreamsTransactionExtractors;
+      this.dataStreamsTransactionExtractorsJson = dataStreamsTransactionExtractorsJson;
+      return this;
+    }
+
     public Builder setServiceMapping(Map<String, String> serviceMapping) {
       return setServiceMapping(serviceMapping.entrySet());
     }
@@ -324,6 +337,8 @@ public static class Snapshot implements TraceConfig {
     final Map<String, String> tracingTags;
 
     final String preferredServiceName;
+    final List<DataStreamsTransactionExtractor> dataStreamsTransactionExtractors;
+    final String dataStreamsTransactionExtractorsJson;
 
     protected Snapshot(DynamicConfig<?>.Builder builder, Snapshot oldSnapshot) {
 
@@ -345,6 +360,8 @@ protected Snapshot(DynamicConfig<?>.Builder builder, Snapshot oldSnapshot) {
       this.tracingTags = nullToEmpty(builder.tracingTags);
 
       this.preferredServiceName = builder.preferredServiceName;
+      this.dataStreamsTransactionExtractors = builder.dataStreamsTransactionExtractors;
+      this.dataStreamsTransactionExtractorsJson = builder.dataStreamsTransactionExtractorsJson;
     }
 
     private static <K, V> Map<K, V> nullToEmpty(Map<K, V> mapping) {
@@ -415,6 +432,11 @@ public List<? extends TraceSamplingRule> getTraceSamplingRules() {
       return traceSamplingRules;
     }
 
+    @Override
+    public List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors() {
+      return dataStreamsTransactionExtractors;
+    }
+
     @Override
     public Map<String, String> getTracingTags() {
       return tracingTags;

internal-api/src/main/java/datadog/trace/api/TraceConfig.java

@@ -1,5 +1,6 @@
 package datadog.trace.api;
 
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.api.sampling.SamplingRule.SpanSamplingRule;
 import datadog.trace.api.sampling.SamplingRule.TraceSamplingRule;
 import java.util.List;
@@ -46,4 +47,11 @@ public interface TraceConfig {
    * @return The tracer sampler Trace Sampling Rules, or an empty collection if no rule is defined.
    */
   List<? extends TraceSamplingRule> getTraceSamplingRules();
+
+  /**
+   * Get DSM transaction extractors.
+   *
+   * @return List of Data Streams Transactions extractors.
+   */
+  List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors();
 }

internal-api/src/main/java/datadog/trace/api/datastreams/DataStreamsTransactionExtractor.java

@@ -0,0 +1,25 @@
+package datadog.trace.api.datastreams;
+
+public final class DataStreamsTransactionExtractor {
+  private final String name;
+  private final String type;
+  private final String value;
+
+  public DataStreamsTransactionExtractor(final String name, final String type, final String value) {
+    this.name = name;
+    this.type = type;
+    this.value = value;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public String getType() {
+    return type;
+  }
+
+  public String getValue() {
+    return value;
+  }
+}

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -6,6 +6,7 @@
 import datadog.trace.api.EndpointTracker;
 import datadog.trace.api.TraceConfig;
 import datadog.trace.api.datastreams.AgentDataStreamsMonitoring;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.api.datastreams.NoopDataStreamsMonitoring;
 import datadog.trace.api.experimental.DataStreamsCheckpointer;
 import datadog.trace.api.gateway.CallbackProvider;
@@ -797,5 +798,10 @@ public List<? extends SamplingRule.SpanSamplingRule> getSpanSamplingRules() {
     public List<? extends SamplingRule.TraceSamplingRule> getTraceSamplingRules() {
       return Collections.emptyList();
     }
+
+    @Override
+    public List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors() {
+      return null;
+    }
   }
 }