diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/jmh/java/datadog/trace/instrumentation/kafka_clients/KafkaInjectBenchmark.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/jmh/java/datadog/trace/instrumentation/kafka_clients/KafkaInjectBenchmark.java new file mode 100644 index 00000000000..0892e4548c6 --- /dev/null +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/jmh/java/datadog/trace/instrumentation/kafka_clients/KafkaInjectBenchmark.java @@ -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}. + * + *

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. + * + *

Single-threaded: each call writes its own headers, with no shared state to contend on. + * + * 

+ *   ./gradlew :dd-java-agent:instrumentation:kafka:kafka-clients-0.11:jmh \
+ *       -Pjmh.includes=KafkaInjectBenchmark -Pjmh.profilers=gc -PtestJvm=17
+ *
+ */ +@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; + } +} diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapter.java index d1ce9a000bf..e87a22813ec 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapter.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapter.java @@ -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) { diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapterTest.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapterTest.java new file mode 100644 index 00000000000..88ed3b89dd3 --- /dev/null +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/java/datadog/trace/instrumentation/kafka_clients/TextMapInjectAdapterTest.java @@ -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 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)); + } + } +} diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapter.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapter.java index 2b58ff6bae4..591fdb5b976 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapter.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapter.java @@ -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) { diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/java/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapterTest.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/java/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapterTest.java new file mode 100644 index 00000000000..4f7ad44d5c0 --- /dev/null +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/java/datadog/trace/instrumentation/kafka_clients38/TextMapInjectAdapterTest.java @@ -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 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)); + } + } +}