Initial version of transactions tracking implementation for DSM

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpClientDecorator.java

@@ -11,6 +11,8 @@
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.ProductActivation;
 import datadog.trace.api.appsec.HttpClientRequest;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
+import datadog.trace.api.datastreams.DataStreamsTransactionTracker;
 import datadog.trace.api.gateway.BlockResponseFunction;
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.api.gateway.RequestContext;
@@ -68,8 +70,19 @@ protected boolean shouldSetResourceName() {
     return true;
   }
 
+  private final DataStreamsTransactionTracker.TransactionSourceReader
+      DSM_TRANSACTION_SOURCE_READER =
+          (source, headerName) -> getRequestHeader((REQUEST) source, headerName);
+
   public AgentSpan onRequest(final AgentSpan span, final REQUEST request) {
     if (request != null) {
+      AgentTracer.get()
+          .getDataStreamsMonitoring()
+          .trackTransaction(
+              span,
+              DataStreamsTransactionExtractor.Type.HTTP_OUT_HEADERS,
+              request,
+              DSM_TRANSACTION_SOURCE_READER);
 
       String method = method(request);
       span.setTag(Tags.HTTP_METHOD, method);

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -13,6 +13,8 @@
 import datadog.context.propagation.Propagators;
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
+import datadog.trace.api.datastreams.DataStreamsTransactionTracker;
 import datadog.trace.api.function.TriConsumer;
 import datadog.trace.api.function.TriFunction;
 import datadog.trace.api.gateway.BlockResponseFunction;
@@ -95,6 +97,14 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
 
   protected abstract int status(RESPONSE response);
 
+  protected String getRequestHeader(REQUEST request, String key) {
+    // This method was not marked as abstract in order to avoid changing all server instrumentation
+    // at once.
+    // Instead, only ones required (by DSM specifically) have it implemented.
+    // This can change in the future.
+    return null;
+  }
+
   protected String requestedSessionId(REQUEST request) {
     return null;
   }
@@ -174,6 +184,10 @@ protected AgentSpanContext startInferredProxySpan(Context context, AgentSpanCont
     return span.start(extracted);
   }
 
+  private final DataStreamsTransactionTracker.TransactionSourceReader
+      DSM_TRANSACTION_SOURCE_READER =
+          (source, headerName) -> getRequestHeader((REQUEST) source, headerName);
+
   public AgentSpan onRequest(
       final AgentSpan span,
       final CONNECTION connection,
@@ -326,6 +340,13 @@ public AgentSpan onRequest(
       span.setRequestBlockingAction((RequestBlockingAction) flow.getAction());
     }
 
+    AgentTracer.get()
+        .getDataStreamsMonitoring()
+        .trackTransaction(
+            span,
+            DataStreamsTransactionExtractor.Type.HTTP_IN_HEADERS,
+            request,
+            DSM_TRANSACTION_SOURCE_READER);
     return span;
   }
 

dd-java-agent/agent-bootstrap/src/main/resources/META-INF/native-image/com.datadoghq/dd-java-agent/reflect-config.json

@@ -70,6 +70,20 @@
       {"name": "toJson"}
     ]
   },
+  {
+    "name" : "datadog.trace.agent.core.datastreams.DataStreamsTransactionExtractors$DataStreamsTransactionExtractorAdapter",
+    "methods": [
+      {"name": "fromJson"},
+      {"name": "toJson"}
+    ]
+  },
+  {
+    "name" : "datadog.trace.agent.core.datastreams.DataStreamsTransactionExtractors$DataStreamsTransactionExtractorsAdapter",
+    "methods": [
+      {"name": "fromJson"},
+      {"name": "toJson"}
+    ]
+  },
   {
     "name" : "datadog.trace.agent.core.DDSpanLink$SpanLinkAdapter",
     "methods": [

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecoratorTest.groovy

@@ -468,6 +468,11 @@ class HttpServerDecoratorTest extends ServerDecoratorTest {
         protected int status(Map m) {
           return m.status == null ? 0 : m.status
         }
+
+        @Override
+        protected String getRequestHeader(Map map, String key) {
+          return map.getOrDefault(key, null)
+        }
       }
   }
 

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
@@ -33,6 +43,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
@@ -56,6 +67,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
@@ -69,24 +87,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.
@@ -250,6 +250,10 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
     List<? extends SamplingRule.TraceSamplingRule> getTraceSamplingRules() {
       return null
     }
+
+    List<DataStreamsTransactionExtractor> getDataStreamsTransactionExtractors() {
+      return null
+    }
   }
 
   @SuppressFBWarnings(value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", justification = "The variable is accessed only by the test thread in setup and cleanup.")

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerDecorator.java

@@ -88,6 +88,12 @@ protected int status(final HttpResponse httpResponse) {
     return httpResponse.status().intValue();
   }
 
+  @Override
+  protected String getRequestHeader(HttpRequest httpRequest, String key) {
+    Optional<akka.http.javadsl.model.HttpHeader> header = httpRequest.getHeader(key);
+    return header.map(HttpHeader::value).orElse(null);
+  }
+
   @Override
   protected boolean isAppSecOnResponseSeparate() {
     return true;

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-10.0/src/main/java11/datadog/trace/instrumentation/jetty10/JettyDecorator.java

@@ -78,6 +78,11 @@ protected boolean isAppSecOnResponseSeparate() {
     return true;
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeader(key);
+  }
+
   @Override
   protected BlockResponseFunction createBlockResponseFunction(Request request, Request connection) {
     return new JettyBlockResponseFunction(request);

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-11.0/src/main/java11/datadog/trace/instrumentation/jetty11/JettyDecorator.java

@@ -78,6 +78,11 @@ protected boolean isAppSecOnResponseSeparate() {
     return true;
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeader(key);
+  }
+
   public AgentSpan onResponse(AgentSpan span, HttpChannel channel) {
     Request request = channel.getRequest();
     Response response = channel.getResponse();

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-12.0/src/main/java17/datadog/trace/instrumentation/jetty12/JettyDecorator.java

@@ -83,6 +83,11 @@ protected int status(final Response response) {
     return response.getStatus();
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeaders().get(key);
+  }
+
   public AgentSpan onResponse(AgentSpan span, HttpChannelState channel) {
     Request request = channel.getRequest();
     Response response = channel.getResponse();

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyDecorator.java

@@ -77,6 +77,11 @@ protected boolean isAppSecOnResponseSeparate() {
     return true;
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeader(key);
+  }
+
   public AgentSpan onResponse(AgentSpan span, HttpConnection channel) {
     Request request = channel.getRequest();
     Response response = channel.getResponse();

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyDecorator.java

@@ -77,6 +77,11 @@ protected boolean isAppSecOnResponseSeparate() {
     return true;
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeader(key);
+  }
+
   public AgentSpan onResponse(AgentSpan span, AbstractHttpConnection connection) {
     Request request = connection.getRequest();
     Response response = connection.getResponse();

dd-java-agent/instrumentation/jetty/jetty-server/jetty-server-9.0/src/main/java/datadog/trace/instrumentation/jetty9/JettyDecorator.java

@@ -77,6 +77,11 @@ protected boolean isAppSecOnResponseSeparate() {
     return true;
   }
 
+  @Override
+  protected String getRequestHeader(final Request request, String key) {
+    return request.getHeader(key);
+  }
+
   public AgentSpan onResponse(AgentSpan span, HttpChannel<?> channel) {
     Request request = channel.getRequest();
     Response response = channel.getResponse();

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java

@@ -15,6 +15,7 @@
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.PRODUCER_DECORATE;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.TIME_IN_QUEUE_ENABLED;
 import static datadog.trace.instrumentation.kafka_common.StreamingContext.STREAMING_CONTEXT;
+import static datadog.trace.instrumentation.kafka_common.Utils.DSM_TRANSACTION_SOURCE_READER;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
@@ -30,6 +31,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.api.datastreams.StatsPoint;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -83,6 +85,7 @@ public String[] helperClassNames() {
       packageName + ".KafkaProducerCallback",
       "datadog.trace.instrumentation.kafka_common.StreamingContext",
       "datadog.trace.instrumentation.kafka_common.ClusterIdHolder",
+      "datadog.trace.instrumentation.kafka_common.Utils",
       packageName + ".AvroSchemaExtractor",
     };
   }
@@ -202,6 +205,14 @@ record =
       if (TIME_IN_QUEUE_ENABLED) {
         setter.injectTimeInQueue(record.headers());
       }
+
+      AgentTracer.get()
+          .getDataStreamsMonitoring()
+          .trackTransaction(
+              span,
+              DataStreamsTransactionExtractor.Type.KAFKA_PRODUCE_HEADERS,
+              record,
+              DSM_TRANSACTION_SOURCE_READER);
       return activateSpan(span);
     }
 

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java

@@ -15,6 +15,7 @@
 import static datadog.trace.instrumentation.kafka_clients.TextMapExtractAdapter.GETTER;
 import static datadog.trace.instrumentation.kafka_clients.TextMapInjectAdapter.SETTER;
 import static datadog.trace.instrumentation.kafka_common.StreamingContext.STREAMING_CONTEXT;
+import static datadog.trace.instrumentation.kafka_common.Utils.DSM_TRANSACTION_SOURCE_READER;
 import static datadog.trace.instrumentation.kafka_common.Utils.computePayloadSizeBytes;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
@@ -23,6 +24,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -128,6 +130,14 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
         if (null != queueSpan) {
           queueSpan.finish();
         }
+
+        AgentTracer.get()
+            .getDataStreamsMonitoring()
+            .trackTransaction(
+                span,
+                DataStreamsTransactionExtractor.Type.KAFKA_CONSUME_HEADERS,
+                val.headers(),
+                DSM_TRANSACTION_SOURCE_READER);
       }
     } catch (final Exception e) {
       log.debug("Error starting new record span", e);

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java/datadog/trace/instrumentation/kafka_clients38/KafkaProducerInstrumentation.java

@@ -43,6 +43,7 @@ public String[] helperClassNames() {
       packageName + ".KafkaProducerCallback",
       "datadog.trace.instrumentation.kafka_common.StreamingContext",
       "datadog.trace.instrumentation.kafka_common.ClusterIdHolder",
+      "datadog.trace.instrumentation.kafka_common.Utils",
       packageName + ".AvroSchemaExtractor",
     };
   }

dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java

@@ -18,6 +18,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.datastreams.DataStreamsContext;
 import datadog.trace.api.datastreams.DataStreamsTags;
+import datadog.trace.api.datastreams.DataStreamsTransactionExtractor;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -128,6 +129,14 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
         if (null != queueSpan) {
           queueSpan.finish();
         }
+
+        AgentTracer.get()
+            .getDataStreamsMonitoring()
+            .trackTransaction(
+                span,
+                DataStreamsTransactionExtractor.Type.KAFKA_CONSUME_HEADERS,
+                val.headers(),
+                Utils.DSM_TRANSACTION_SOURCE_READER);
       }
     } catch (final Exception e) {
       log.debug("Error starting new record span", e);

dd-java-agent/instrumentation/kafka/kafka-common/src/main/java/datadog/trace/instrumentation/kafka_common/Utils.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.kafka_common;
 
+import datadog.trace.api.datastreams.DataStreamsTransactionTracker;
 import java.nio.charset.StandardCharsets;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.header.Header;
@@ -8,6 +9,10 @@
 public final class Utils {
   private Utils() {} // prevent instantiation
 
+  public static DataStreamsTransactionTracker.TransactionSourceReader
+      DSM_TRANSACTION_SOURCE_READER =
+          (source, headerName) -> new String(((Headers) source).lastHeader(headerName).value());
+
   // this method is used in kafka-clients and kafka-streams instrumentations
   public static long computePayloadSizeBytes(ConsumerRecord<?, ?> val) {
     long headersSize = 0;

dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletDecorator.java

@@ -67,4 +67,9 @@ protected int peerPort(final HttpExchange exchange) {
   protected int status(final HttpExchange exchange) {
     return exchange.getResponseCode();
   }
+
+  @Override
+  protected String getRequestHeader(final HttpExchange exchange, String key) {
+    return exchange.getRequestHeaders().getFirst(key);
+  }
 }