add instrumentation to allow sending kafka payload size on produce

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

@@ -14,6 +14,7 @@
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.TIME_IN_QUEUE_ENABLED;
 import static datadog.trace.instrumentation.kafka_clients.TextMapInjectAdapter.SETTER;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
@@ -22,7 +23,9 @@
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags;
+import datadog.trace.bootstrap.instrumentation.api.StatsPoint;
 import java.util.LinkedHashMap;
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.ApiVersions;
@@ -62,6 +65,13 @@ public void adviceTransformations(AdviceTransformation transformation) {
             .and(takesArgument(0, named("org.apache.kafka.clients.producer.ProducerRecord")))
             .and(takesArgument(1, named("org.apache.kafka.clients.producer.Callback"))),
         KafkaProducerInstrumentation.class.getName() + "$ProducerAdvice");
+
+    transformation.applyAdvice(
+        isMethod()
+            .and(isPrivate())
+            .and(takesArgument(0, int.class))
+            .and(named("ensureValidRecordSize")),
+        KafkaProducerInstrumentation.class.getName() + "$PayloadSizeAdvice");
   }
 
   public static class ProducerAdvice {
@@ -99,7 +109,8 @@ public static AgentScope onEnter(
         sortedTags.put(TYPE_TAG, "kafka");
         try {
           propagate().inject(span, record.headers(), SETTER);
-          propagate().injectPathwayContext(span, record.headers(), SETTER, sortedTags);
+          propagate()
+              .injectPathwayContextWithoutSendingStats(span, record.headers(), SETTER, sortedTags);
         } catch (final IllegalStateException e) {
           // headers must be read-only from reused record. try again with new one.
           record =
@@ -112,7 +123,8 @@ record =
                   record.headers());
 
           propagate().inject(span, record.headers(), SETTER);
-          propagate().injectPathwayContext(span, record.headers(), SETTER, sortedTags);
+          propagate()
+              .injectPathwayContextWithoutSendingStats(span, record.headers(), SETTER, sortedTags);
         }
         if (TIME_IN_QUEUE_ENABLED) {
           SETTER.injectTimeInQueue(record.headers());
@@ -130,4 +142,31 @@ public static void stopSpan(
       scope.close();
     }
   }
+
+  public static class PayloadSizeAdvice {
+
+    /**
+     * Instrumentation for the method KafkaProducer.ensureValidRecordSize that is called as part of
+     * sending a kafka payload. This gives us access to an estimate of the payload size "for free",
+     * that we send as a metric.
+     */
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(@Advice.Argument(value = 0) int estimatedPayloadSize) {
+      StatsPoint saved = activeSpan().context().getPathwayContext().getSavedStats();
+      if (saved != null) {
+        // create new stats including the payload size
+        StatsPoint updated =
+            new StatsPoint(
+                saved.getEdgeTags(),
+                saved.getHash(),
+                saved.getParentHash(),
+                saved.getTimestampNanos(),
+                saved.getPathwayLatencyNano(),
+                saved.getEdgeLatencyNano(),
+                estimatedPayloadSize);
+        // then send the point
+        AgentTracer.get().getDataStreamsMonitoring().add(updated);
+      }
+    }
+  }
 }

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

@@ -24,14 +24,31 @@ public <C> void injectPathwayContext(
       C carrier,
       AgentPropagation.Setter<C> setter,
       LinkedHashMap<String, String> sortedTags) {
-    PathwayContext pathwayContext = span.context().getPathwayContext();
+    injectPathwayContext(span, carrier, setter, sortedTags, true);
+  }
 
+  /** Same as injectPathwayContext, but the stats collected in the StatsPoint are not sent. */
+  public <C> void injectPathwayContextWithoutSendingStats(
+      AgentSpan span,
+      C carrier,
+      AgentPropagation.Setter<C> setter,
+      LinkedHashMap<String, String> sortedTags) {
+    injectPathwayContext(span, carrier, setter, sortedTags, false);
+  }
+
+  private <C> void injectPathwayContext(
+      AgentSpan span,
+      C carrier,
+      AgentPropagation.Setter<C> setter,
+      LinkedHashMap<String, String> sortedTags,
+      boolean sendCheckpoint) {
+    PathwayContext pathwayContext = span.context().getPathwayContext();
     if (pathwayContext == null
         || (span.traceConfig() != null && !span.traceConfig().isDataStreamsEnabled())) {
       return;
     }
-    pathwayContext.setCheckpoint(sortedTags, dataStreamsMonitoring::add);
-
+    pathwayContext.setCheckpoint(
+        sortedTags, sendCheckpoint ? dataStreamsMonitoring::add : pathwayContext::saveStats);
     boolean injected =
         setter instanceof AgentPropagation.BinarySetter
             ? injectBinaryPathwayContext(

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

@@ -44,6 +44,7 @@ public class DefaultPathwayContext implements PathwayContext {
   private long pathwayStartNanos;
   private long pathwayStartNanoTicks;
   private long edgeStartNanoTicks;
+  private StatsPoint savedStats;
   private long hash;
   private boolean started;
   // state variables used to memoize the pathway hash with
@@ -174,7 +175,7 @@ public void setCheckpoint(
               allTags,
               newHash,
               hash,
-              timeSource.getCurrentTimeNanos(),
+              startNanos,
               pathwayLatencyNano,
               edgeLatencyNano,
               payloadSizeBytes);
@@ -188,6 +189,16 @@ public void setCheckpoint(
     }
   }
 
+  @Override
+  public void saveStats(StatsPoint point) {
+    this.savedStats = point;
+  }
+
+  @Override
+  public StatsPoint getSavedStats() {
+    return this.savedStats;
+  }
+
   @Override
   public byte[] encode() throws IOException {
     lock.lock();

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

@@ -71,6 +71,13 @@ public <C> void injectPathwayContext(
     this.dataStreamContextInjector.injectPathwayContext(span, carrier, setter, sortedTags);
   }
 
+  @Override
+  public <C> void injectPathwayContextWithoutSendingStats(
+      AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags) {
+    this.dataStreamContextInjector.injectPathwayContextWithoutSendingStats(
+        span, carrier, setter, sortedTags);
+  }
+
   @Override
   public <C> AgentSpan.Context.Extracted extract(final C carrier, final ContextVisitor<C> getter) {
     return extractor.extract(carrier, getter);

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

@@ -14,6 +14,9 @@ public interface AgentPropagation {
   <C> void injectPathwayContext(
       AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags);
 
+  <C> void injectPathwayContextWithoutSendingStats(
+      AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags);
+
   interface Setter<C> {
     void set(C carrier, String key, String value);
   }

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

@@ -867,6 +867,10 @@ public <C> void inject(
     public <C> void injectPathwayContext(
         AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags) {}
 
+    @Override
+    public <C> void injectPathwayContextWithoutSendingStats(
+        AgentSpan span, C carrier, Setter<C> setter, LinkedHashMap<String, String> sortedTags) {}
+
     @Override
     public <C> Context.Extracted extract(final C carrier, final ContextVisitor<C> getter) {
       return NoopContext.INSTANCE;
@@ -1082,6 +1086,14 @@ public void setCheckpoint(
     public void setCheckpoint(
         LinkedHashMap<String, String> sortedTags, Consumer<StatsPoint> pointConsumer) {}
 
+    @Override
+    public void saveStats(StatsPoint point) {}
+
+    @Override
+    public StatsPoint getSavedStats() {
+      return null;
+    }
+
     @Override
     public byte[] encode() {
       return null;

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

@@ -27,6 +27,10 @@ void setCheckpoint(
   // The input tags should be sorted.
   void setCheckpoint(LinkedHashMap<String, String> sortedTags, Consumer<StatsPoint> pointConsumer);
 
+  void saveStats(StatsPoint point);
+
+  StatsPoint getSavedStats();
+
   byte[] encode() throws IOException;
 
   String strEncode() throws IOException;