Skip finding a Kafka record's produce time when time in queue is off

dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/jmh/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumeScenariosBenchmark.java

@@ -5,6 +5,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.CONSUMER_DECORATE;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_CONSUME;
+import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.TIME_IN_QUEUE_ENABLED;
 import static datadog.trace.instrumentation.kafka_clients.TextMapInjectAdapter.SETTER;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
@@ -54,6 +55,10 @@
  * case; with no local Datadog Agent, DSM still builds tags, contexts and checkpoints, but drops the
  * stats points instead of sending them.
  *
+ * <p>{@code timeInQueue} sets {@code dd.kafka.time-in-queue.enabled}, which is off by default with
+ * legacy Kafka tracing. With it on, producers also write their produce time to a header, and the
+ * consumer starts a {@code kafka.deliver} span from it.
+ *
  * <p><b>Read ns/op and {@code gc.alloc.rate.norm}</b> (B/op); both are per record. Expect the
  * relative difference between builds to be larger than in a real consumer: the records here are
  * pre-built, so Kafka's fetch and lazy header parsing (which the agent triggers by reading headers)
@@ -88,6 +93,9 @@ public class KafkaConsumeScenariosBenchmark {
   @Param({"false", "true"})
   boolean dataStreams;
 
+  @Param({"false", "true"})
+  boolean timeInQueue;
+
   CoreTracer tracer;
   List<ConsumerRecord<?, ?>> batch;
 
@@ -96,9 +104,13 @@ public void setup(Blackhole blackhole) {
     // Config is read once per JVM, so this must run before anything touches it. JMH forks a fresh
     // JVM per parameter value; running in-process (-f 0) would silently measure one setting.
     System.setProperty("dd." + DATA_STREAMS_ENABLED, String.valueOf(dataStreams));
+    System.setProperty("dd.kafka.time-in-queue.enabled", String.valueOf(timeInQueue));
     if (Config.get().isDataStreamsEnabled() != dataStreams) {
       throw new IllegalStateException("Config was read before setup; run with forks (not -f 0)");
     }
+    if (TIME_IN_QUEUE_ENABLED != timeInQueue) {
+      throw new IllegalStateException("time in queue doesn't match timeInQueue=" + timeInQueue);
+    }
 
     // Building the tracer registers its propagators as the default ones.
     tracer = CoreTracer.builder().writer(new DropWriter(blackhole)).build();
@@ -114,6 +126,9 @@ public void setup(Blackhole blackhole) {
       RecordHeaders headers = new RecordHeaders();
       AgentSpan producer = tracer.buildSpan("bench", "kafka.produce").start();
       defaultPropagator().inject(producer, headers, SETTER);
+      if (timeInQueue) {
+        SETTER.injectTimeInQueue(headers);
+      }
       producer.finish();
 
       batch.add(

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

@@ -4,11 +4,11 @@
 import static datadog.trace.api.Functions.UTF8_BYTES_TO_STRING;
 import static datadog.trace.api.telemetry.LogCollector.EXCLUDE_TELEMETRY;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_PRODUCED_KEY;
+import static datadog.trace.instrumentation.kafka_common.Utils.readProducedTime;
 
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
-import java.nio.ByteBuffer;
 import java.util.Base64;
 import java.util.function.Function;
 import org.apache.kafka.common.header.Header;
@@ -58,10 +58,7 @@ public long extractTimeInQueueStart(Headers carrier) {
     Header header = carrier.lastHeader(KAFKA_PRODUCED_KEY);
     if (null != header) {
       try {
-        ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(decoder != null ? decoder.decode(header.value()) : header.value());
-        buf.flip();
-        return buf.getLong();
+        return readProducedTime(decoder != null ? decoder.decode(header.value()) : header.value());
       } catch (Exception e) {
         log.debug(EXCLUDE_TELEMETRY, "Unable to get kafka produced time", e);
       }

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

@@ -102,8 +102,10 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
         if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) {
           final AgentSpanContext spanContext =
               extractContextAndGetSpanContext(val.headers(), GETTER);
-          long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers());
-          if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) {
+          // Check the flag first, as finding the produce time scans the record's headers
+          long timeInQueueStart =
+              TIME_IN_QUEUE_ENABLED ? GETTER.extractTimeInQueueStart(val.headers()) : 0;
+          if (timeInQueueStart == 0) {
             span = startSpan(JAVA_KAFKA.toString(), operationName, spanContext);
           } else {
             queueSpan =

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

@@ -3,11 +3,11 @@
 import static datadog.trace.api.Functions.BASE64_DECODE;
 import static datadog.trace.api.Functions.UTF8_BYTES_TO_STRING;
 import static datadog.trace.api.telemetry.LogCollector.EXCLUDE_TELEMETRY;
+import static datadog.trace.instrumentation.kafka_common.Utils.readProducedTime;
 
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
-import java.nio.ByteBuffer;
 import java.util.Base64;
 import java.util.function.Function;
 import org.apache.kafka.common.header.Header;
@@ -57,10 +57,7 @@ public long extractTimeInQueueStart(Headers carrier) {
     Header header = carrier.lastHeader(KafkaDecorator.KAFKA_PRODUCED_KEY);
     if (null != header) {
       try {
-        ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(decoder != null ? decoder.decode(header.value()) : header.value());
-        buf.flip();
-        return buf.getLong();
+        return readProducedTime(decoder != null ? decoder.decode(header.value()) : header.value());
       } catch (Exception e) {
         log.debug(EXCLUDE_TELEMETRY, "Unable to get kafka produced time", e);
       }

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

@@ -98,8 +98,12 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
         if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) {
           final AgentSpanContext spanContext =
               extractContextAndGetSpanContext(val.headers(), GETTER);
-          long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers());
-          if (timeInQueueStart == 0 || !KafkaDecorator.TIME_IN_QUEUE_ENABLED) {
+          // Check the flag first, as finding the produce time scans the record's headers
+          long timeInQueueStart =
+              KafkaDecorator.TIME_IN_QUEUE_ENABLED
+                  ? GETTER.extractTimeInQueueStart(val.headers())
+                  : 0;
+          if (timeInQueueStart == 0) {
             span = startSpan(JAVA_KAFKA.toString(), operationName, spanContext);
           } else {
             queueSpan =

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

@@ -19,6 +19,22 @@ private Utils() {} // prevent instantiation
             }
           };
 
+  /**
+   * Reads the produce time a producer wrote to its header, as 8 big-endian bytes.
+   *
+   * @throws IllegalArgumentException if {@code bytes} isn't 8 bytes long
+   */
+  public static long readProducedTime(byte[] bytes) {
+    if (bytes.length != Long.BYTES) {
+      throw new IllegalArgumentException("Expected 8 bytes, got " + bytes.length);
+    }
+    long time = 0;
+    for (byte b : bytes) {
+      time = (time << 8) | (b & 0xff);
+    }
+    return time;
+  }
+
   // this method is used in kafka-clients and kafka-streams instrumentations
   public static long computePayloadSizeBytes(ConsumerRecord<?, ?> val) {
     long headersSize = 0;

dd-java-agent/instrumentation/kafka/kafka-common/src/test/java/datadog/trace/instrumentation/kafka_common/UtilsTest.java

@@ -0,0 +1,27 @@
+package datadog.trace.instrumentation.kafka_common;
+
+import static datadog.trace.instrumentation.kafka_common.Utils.readProducedTime;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.nio.ByteBuffer;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class UtilsTest {
+
+  @ParameterizedTest
+  @ValueSource(longs = {0L, 1L, 1_793_652_306_477L, -1L, Long.MAX_VALUE, Long.MIN_VALUE, 0x80L})
+  void readsTheTimeAProducerWrote(long time) {
+    byte[] bytes = ByteBuffer.allocate(Long.BYTES).putLong(time).array();
+
+    assertEquals(time, readProducedTime(bytes));
+  }
+
+  @Test
+  void rejectsAValueThatIsNotEightBytes() {
+    assertThrows(IllegalArgumentException.class, () -> readProducedTime(new byte[7]));
+    assertThrows(IllegalArgumentException.class, () -> readProducedTime(new byte[9]));
+  }
+}

dd-java-agent/instrumentation/kafka/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java

@@ -278,8 +278,9 @@ public static void start(
       AgentSpan span, queueSpan = null;
       StreamTaskContext streamTaskContext =
           InstrumentationContext.get(StreamTask.class, StreamTaskContext.class).get(task);
-      long timeInQueueStart = SR_GETTER.extractTimeInQueueStart(record);
-      if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) {
+      // Check the flag first, as finding the produce time scans the record's headers
+      long timeInQueueStart = TIME_IN_QUEUE_ENABLED ? SR_GETTER.extractTimeInQueueStart(record) : 0;
+      if (timeInQueueStart == 0) {
         span = startSpan(JAVA_KAFKA.toString(), KAFKA_CONSUME);
       } else {
         queueSpan =
@@ -344,8 +345,9 @@ public static void start(
       AgentSpan span, queueSpan = null;
       StreamTaskContext streamTaskContext =
           InstrumentationContext.get(StreamTask.class, StreamTaskContext.class).get(task);
-      long timeInQueueStart = PR_GETTER.extractTimeInQueueStart(record);
-      if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) {
+      // Check the flag first, as finding the produce time scans the record's headers
+      long timeInQueueStart = TIME_IN_QUEUE_ENABLED ? PR_GETTER.extractTimeInQueueStart(record) : 0;
+      if (timeInQueueStart == 0) {
         span = startSpan(JAVA_KAFKA.toString(), KAFKA_CONSUME);
       } else {
         queueSpan =

dd-java-agent/instrumentation/kafka/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/ProcessorRecordContextVisitor.java

@@ -1,12 +1,12 @@
 package datadog.trace.instrumentation.kafka_streams;
 
+import static datadog.trace.instrumentation.kafka_common.Utils.readProducedTime;
 import static datadog.trace.instrumentation.kafka_streams.KafkaStreamsDecorator.KAFKA_PRODUCED_KEY;
 import static datadog.trace.instrumentation.kafka_streams.ProcessorRecordContextHeadersAccess.HEADERS_METHOD;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
-import java.nio.ByteBuffer;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.streams.processor.internals.ProcessorRecordContext;
@@ -48,10 +48,7 @@ public long extractTimeInQueueStart(ProcessorRecordContext carrier) {
       Headers headers = (Headers) HEADERS_METHOD.invokeExact(carrier);
       Header header = headers.lastHeader(KAFKA_PRODUCED_KEY);
       if (null != header) {
-        ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(header.value());
-        buf.flip();
-        return buf.getLong();
+        return readProducedTime(header.value());
       }
     } catch (Throwable e) {
       log.debug("Unable to get kafka produced time", e);

dd-java-agent/instrumentation/kafka/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/StampedRecordContextVisitor.java

@@ -1,11 +1,11 @@
 package datadog.trace.instrumentation.kafka_streams;
 
+import static datadog.trace.instrumentation.kafka_common.Utils.readProducedTime;
 import static datadog.trace.instrumentation.kafka_streams.KafkaStreamsDecorator.KAFKA_PRODUCED_KEY;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation.ContextVisitor;
-import java.nio.ByteBuffer;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.streams.processor.internals.StampedRecord;
 import org.slf4j.Logger;
@@ -34,10 +34,7 @@ public long extractTimeInQueueStart(StampedRecord carrier) {
     try {
       Header header = carrier.value.headers().lastHeader(KAFKA_PRODUCED_KEY);
       if (null != header) {
-        ByteBuffer buf = ByteBuffer.allocate(8);
-        buf.put(header.value());
-        buf.flip();
-        return buf.getLong();
+        return readProducedTime(header.value());
       }
     } catch (Exception e) {
       log.debug("Unable to get kafka produced time", e);