Skip context extraction for Kafka records without headers

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

@@ -54,6 +54,9 @@
  * case; with no local Datadog Agent, DSM still builds tags, contexts and checkpoints, but drops the
  * stats points instead of sending them.
  *
+ * <p>{@code traceHeaders} {@code false} gives records with no headers at all, as an untraced
+ * producer sends them.
+ *
  * <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 +91,9 @@ public class KafkaConsumeScenariosBenchmark {
   @Param({"false", "true"})
   boolean dataStreams;
 
+  @Param({"true", "false"})
+  boolean traceHeaders;
+
   CoreTracer tracer;
   List<ConsumerRecord<?, ?>> batch;
 
@@ -112,9 +118,11 @@ public void setup(Blackhole blackhole) {
     long timestamp = System.currentTimeMillis();
     for (int i = 0; i < BATCH_SIZE; i++) {
       RecordHeaders headers = new RecordHeaders();
-      AgentSpan producer = tracer.buildSpan("bench", "kafka.produce").start();
-      defaultPropagator().inject(producer, headers, SETTER);
-      producer.finish();
+      if (traceHeaders) {
+        AgentSpan producer = tracer.buildSpan("bench", "kafka.produce").start();
+        defaultPropagator().inject(producer, headers, SETTER);
+        producer.finish();
+      }
 
       batch.add(
           new ConsumerRecord<>(

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

@@ -19,6 +19,7 @@
 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 datadog.trace.instrumentation.kafka_common.Utils.hasHeaders;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -151,9 +152,11 @@ public static ContextScope onEnter(
         ClusterIdHolder.set(clusterId);
       }
 
-      // Try to extract existing trace context from record headers
+      // Try to extract existing trace context from record headers, which most records don't have
       final AgentSpanContext extractedContext =
-          extractContextAndGetSpanContext(record.headers(), TextMapExtractAdapter.GETTER);
+          hasHeaders(record.headers())
+              ? extractContextAndGetSpanContext(record.headers(), TextMapExtractAdapter.GETTER)
+              : null;
 
       final AgentSpan localActiveSpan = activeSpan();
 

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

@@ -18,6 +18,7 @@
 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 datadog.trace.instrumentation.kafka_common.Utils.hasHeaders;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import datadog.context.Context;
@@ -100,8 +101,11 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
       AgentSpan span, queueSpan = null;
       if (val != null) {
         if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) {
+          // A record without headers has no context to extract
           final AgentSpanContext spanContext =
-              extractContextAndGetSpanContext(val.headers(), GETTER);
+              hasHeaders(val.headers())
+                  ? extractContextAndGetSpanContext(val.headers(), GETTER)
+                  : null;
           long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers());
           if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) {
             span = startSpan(JAVA_KAFKA.toString(), operationName, spanContext);

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

@@ -8,6 +8,7 @@
 import static datadog.trace.instrumentation.kafka_clients38.KafkaDecorator.JAVA_KAFKA;
 import static datadog.trace.instrumentation.kafka_clients38.KafkaDecorator.KAFKA_PRODUCE;
 import static datadog.trace.instrumentation.kafka_clients38.KafkaDecorator.PRODUCER_DECORATE;
+import static datadog.trace.instrumentation.kafka_common.Utils.hasHeaders;
 
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.InstrumentationContext;
@@ -41,9 +42,11 @@ public static ContextScope onEnter(
       ClusterIdHolder.set(clusterId);
     }
 
-    // Try to extract existing trace context from record headers
+    // Try to extract existing trace context from record headers, which most records don't have
     final AgentSpanContext extractedContext =
-        extractContextAndGetSpanContext(record.headers(), TextMapExtractAdapter.GETTER);
+        hasHeaders(record.headers())
+            ? extractContextAndGetSpanContext(record.headers(), TextMapExtractAdapter.GETTER)
+            : null;
 
     final AgentSpan localActiveSpan = activeSpan();
 

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

@@ -96,8 +96,11 @@ protected void startNewRecordSpan(ConsumerRecord<?, ?> val) {
       AgentSpan span, queueSpan = null;
       if (val != null) {
         if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) {
+          // A record without headers has no context to extract
           final AgentSpanContext spanContext =
-              extractContextAndGetSpanContext(val.headers(), GETTER);
+              Utils.hasHeaders(val.headers())
+                  ? extractContextAndGetSpanContext(val.headers(), GETTER)
+                  : null;
           long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers());
           if (timeInQueueStart == 0 || !KafkaDecorator.TIME_IN_QUEUE_ENABLED) {
             span = startSpan(JAVA_KAFKA.toString(), operationName, spanContext);

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

@@ -19,6 +19,14 @@ private Utils() {} // prevent instantiation
             }
           };
 
+  /**
+   * Whether {@code headers} holds any header. {@link Headers#toArray()} would copy the headers of
+   * every record that has some, just to check.
+   */
+  public static boolean hasHeaders(Headers headers) {
+    return headers.iterator().hasNext();
+  }
+
   // 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,32 @@
+package datadog.trace.instrumentation.kafka_common;
+
+import static datadog.trace.instrumentation.kafka_common.Utils.hasHeaders;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.junit.jupiter.api.Test;
+
+class UtilsTest {
+
+  @Test
+  void recordWithoutHeadersHasNone() {
+    assertFalse(hasHeaders(new RecordHeaders()));
+  }
+
+  @Test
+  void recordWithAHeaderHasHeaders() {
+    RecordHeaders headers = new RecordHeaders();
+    headers.add("x-datadog-trace-id", new byte[] {'1'});
+
+    assertTrue(hasHeaders(headers));
+  }
+
+  @Test
+  void readOnlyHeadersCanBeChecked() {
+    RecordHeaders headers = new RecordHeaders();
+    headers.setReadOnly();
+
+    assertFalse(hasHeaders(headers));
+  }
+}