dougqh · Nov 3, 2026 · Nov 3, 2026 · Nov 3, 2026
@@ -0,0 +1,81 @@
+package datadog.trace.instrumentation.kafka_clients;
+
+import static datadog.context.propagation.Propagators.defaultPropagator;
+import static datadog.trace.instrumentation.kafka_clients.TextMapInjectAdapter.SETTER;
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.common.writer.ListWriter;
+import datadog.trace.core.CoreTracer;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * Injecting a span's context into a new record's headers, as the producer does on each send, with a
+ * default-configured {@link CoreTracer}.
+ *
+ * <p>The same span is injected every time, so the header values themselves may come from caches;
+ * this measures what writing them to Kafka's {@link RecordHeaders} costs. {@code spanBaggage} gives
+ * the span a baggage item, which both the Datadog and W3C codecs write as the same {@code
+ * ot-baggage-} header, so the second write replaces the first.
+ *
+ * <p>Single-threaded: each call writes its own headers, with no shared state to contend on.
+ *
+ * <pre>
+ *   ./gradlew :dd-java-agent:instrumentation:kafka:kafka-clients-0.11:jmh \
+ *       -Pjmh.includes=KafkaInjectBenchmark -Pjmh.profilers=gc -PtestJvm=17
+ * </pre>
+ */
+@State(Scope.Benchmark)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(NANOSECONDS)
+@Threads(1)
+@Fork(3)
+public class KafkaInjectBenchmark {
+  @Param({"false", "true"})
+  boolean spanBaggage;
+
+  CoreTracer tracer;
+  AgentSpan span;
+
+  @Setup
+  public void setup() {
+    // Building the tracer registers its propagators as the default ones. The span only finishes at
+    // teardown, so the writer never runs while measuring.
+    tracer = CoreTracer.builder().writer(new ListWriter()).build();
+    AgentTracer.forceRegister(tracer);
+    span = tracer.buildSpan("bench", "kafka.produce").start();
+    if (spanBaggage) {
+      span.setBaggageItem("user.id", "42");
+    }
+  }
+
+  @TearDown
+  public void tearDown() {
+    span.finish();
+    AgentTracer.forceRegister(AgentTracer.NOOP_TRACER);
+    tracer.close();
+  }
+
+  @Benchmark
+  public RecordHeaders inject() {
+    RecordHeaders headers = new RecordHeaders();
+    defaultPropagator().inject(span, headers, SETTER);
+    return headers;
+  }
+}
@@ -13,7 +13,12 @@ public class TextMapInjectAdapter implements TextMapInjectAdapterInterface {
 
   @Override
   public void set(final Headers headers, final String key, final String value) {
-    headers.remove(key).add(key, value.getBytes(UTF_8));
+    // Codecs can write the same key, so replace it, but probe first: remove allocates an iterator,
+    // lastHeader doesn't, and the key is usually not there yet
+    if (headers.lastHeader(key) != null) {
+      headers.remove(key);
+    }
+    headers.add(key, value.getBytes(UTF_8));
   }
 
   public void injectTimeInQueue(Headers headers) {

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.kafka_clients;
+
+import static datadog.trace.instrumentation.kafka_clients.TextMapInjectAdapter.SETTER;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.junit.jupiter.api.Test;
+
+class TextMapInjectAdapterTest {
+
+  @Test
+  void addsAHeaderThatIsNotThereYet() {
+    Headers headers = new RecordHeaders();
+
+    SETTER.set(headers, "traceparent", "00-abc-def-01");
+
+    assertValues(headers, "traceparent", "00-abc-def-01");
+  }
+
+  @Test
+  void replacesAHeaderAlreadyThere() {
+    Headers headers = new RecordHeaders();
+    headers.add("x-datadog-trace-id", "1".getBytes(UTF_8));
+    headers.add("x-datadog-trace-id", "2".getBytes(UTF_8));
+
+    SETTER.set(headers, "x-datadog-trace-id", "3");
+
+    assertValues(headers, "x-datadog-trace-id", "3");
+  }
+
+  /** The Datadog and W3C codecs both write ot-baggage- keys, and the record must carry one. */
+  @Test
+  void keepsOneHeaderWhenTheSameKeyIsSetTwice() {
+    Headers headers = new RecordHeaders();
+
+    SETTER.set(headers, "ot-baggage-t0", "100");
+    SETTER.set(headers, "ot-baggage-t0", "200");
+
+    assertValues(headers, "ot-baggage-t0", "200");
+  }
+
+  /** The producer advice relies on this to retry with a writable copy of the record. */
+  @Test
+  void throwsOnReadOnlyHeaders() {
+    RecordHeaders empty = new RecordHeaders();
+    empty.setReadOnly();
+    RecordHeaders withKey = new RecordHeaders();
+    withKey.add("traceparent", "old".getBytes(UTF_8));
+    withKey.setReadOnly();
+
+    assertThrows(IllegalStateException.class, () -> SETTER.set(empty, "traceparent", "new"));
+    assertThrows(IllegalStateException.class, () -> SETTER.set(withKey, "traceparent", "new"));
+  }
+
+  private static void assertValues(Headers headers, String key, String... expected) {
+    List<byte[]> values = new ArrayList<>();
+    for (Header header : headers.headers(key)) {
+      values.add(header.value());
+    }
+    assertEquals(expected.length, values.size());
+    for (int i = 0; i < expected.length; i++) {
+      assertArrayEquals(expected[i].getBytes(UTF_8), values.get(i));
+    }
+  }
+}
@@ -10,7 +10,12 @@ public class TextMapInjectAdapter implements TextMapInjectAdapterInterface {
 
   @Override
   public void set(final Headers headers, final String key, final String value) {
-    headers.remove(key).add(key, value.getBytes(UTF_8));
+    // Codecs can write the same key, so replace it, but probe first: remove allocates an iterator,
+    // lastHeader doesn't, and the key is usually not there yet
+    if (headers.lastHeader(key) != null) {
+      headers.remove(key);
+    }
+    headers.add(key, value.getBytes(UTF_8));
   }
 
   public void injectTimeInQueue(Headers headers) {

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.kafka_clients38;
+
+import static datadog.trace.instrumentation.kafka_clients38.TextMapInjectAdapter.SETTER;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.junit.jupiter.api.Test;
+
+class TextMapInjectAdapterTest {
+
+  @Test
+  void addsAHeaderThatIsNotThereYet() {
+    Headers headers = new RecordHeaders();
+
+    SETTER.set(headers, "traceparent", "00-abc-def-01");
+
+    assertValues(headers, "traceparent", "00-abc-def-01");
+  }
+
+  @Test
+  void replacesAHeaderAlreadyThere() {
+    Headers headers = new RecordHeaders();
+    headers.add("x-datadog-trace-id", "1".getBytes(UTF_8));
+    headers.add("x-datadog-trace-id", "2".getBytes(UTF_8));
+
+    SETTER.set(headers, "x-datadog-trace-id", "3");
+
+    assertValues(headers, "x-datadog-trace-id", "3");
+  }
+
+  /** The Datadog and W3C codecs both write ot-baggage- keys, and the record must carry one. */
+  @Test
+  void keepsOneHeaderWhenTheSameKeyIsSetTwice() {
+    Headers headers = new RecordHeaders();
+
+    SETTER.set(headers, "ot-baggage-t0", "100");
+    SETTER.set(headers, "ot-baggage-t0", "200");
+
+    assertValues(headers, "ot-baggage-t0", "200");
+  }
+
+  /** The producer advice relies on this to retry with a writable copy of the record. */
+  @Test
+  void throwsOnReadOnlyHeaders() {
+    RecordHeaders empty = new RecordHeaders();
+    empty.setReadOnly();
+    RecordHeaders withKey = new RecordHeaders();
+    withKey.add("traceparent", "old".getBytes(UTF_8));
+    withKey.setReadOnly();
+
+    assertThrows(IllegalStateException.class, () -> SETTER.set(empty, "traceparent", "new"));
+    assertThrows(IllegalStateException.class, () -> SETTER.set(withKey, "traceparent", "new"));
+  }
+
+  private static void assertValues(Headers headers, String key, String... expected) {
+    List<byte[]> values = new ArrayList<>();
+    for (Header header : headers.headers(key)) {
+      values.add(header.value());
+    }
+    assertEquals(expected.length, values.size());
+    for (int i = 0; i < expected.length; i++) {
+      assertArrayEquals(expected[i].getBytes(UTF_8), values.get(i));
+    }
+  }
+}