diff --git a/spring-kafka/src/main/java/org/springframework/kafka/support/micrometer/KafkaRecordSenderContext.java b/spring-kafka/src/main/java/org/springframework/kafka/support/micrometer/KafkaRecordSenderContext.java index 0d186993..50e9e310 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/support/micrometer/KafkaRecordSenderContext.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/support/micrometer/KafkaRecordSenderContext.java @@ -1,5 +1,5 @@ /* - * Copyright 2022-2024 the original author or authors. + * Copyright 2022-2025 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -21,6 +21,7 @@ import java.util.function.Supplier; import io.micrometer.observation.transport.SenderContext; import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.header.Headers; /** * {@link SenderContext} for {@link ProducerRecord}s. @@ -28,6 +29,7 @@ import org.apache.kafka.clients.producer.ProducerRecord; * @author Gary Russell * @author Christian Mergenthaler * @author Wang Zhiyang + * @author Soby Chacko * * @since 3.0 * @@ -39,8 +41,12 @@ public class KafkaRecordSenderContext extends SenderContext private final ProducerRecord record; public KafkaRecordSenderContext(ProducerRecord record, String beanName, Supplier clusterId) { - super((carrier, key, value) -> record.headers().add(key, - value == null ? null : value.getBytes(StandardCharsets.UTF_8))); + super((carrier, key, value) -> { + Headers headers = record.headers(); + headers.remove(key); + headers.add(key, value == null ? null : value.getBytes(StandardCharsets.UTF_8)); + }); + setCarrier(record); this.beanName = beanName; this.record = record; diff --git a/spring-kafka/src/test/java/org/springframework/kafka/support/micrometer/ObservationTests.java b/spring-kafka/src/test/java/org/springframework/kafka/support/micrometer/ObservationTests.java index aaacdbdb..f23335c0 100644 --- a/spring-kafka/src/test/java/org/springframework/kafka/support/micrometer/ObservationTests.java +++ b/spring-kafka/src/test/java/org/springframework/kafka/support/micrometer/ObservationTests.java @@ -16,6 +16,7 @@ package org.springframework.kafka.support.micrometer; +import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Deque; import java.util.List; @@ -25,6 +26,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.StreamSupport; import io.micrometer.common.KeyValues; import io.micrometer.core.instrument.MeterRegistry; @@ -53,6 +55,7 @@ import org.apache.kafka.clients.producer.RecordMetadata; import org.apache.kafka.common.errors.InvalidTopicException; import org.apache.kafka.common.header.Header; import org.apache.kafka.common.header.Headers; +import org.apache.kafka.common.header.internals.RecordHeader; import org.junit.jupiter.api.Test; import org.springframework.beans.factory.annotation.Autowired; @@ -72,6 +75,7 @@ import org.springframework.kafka.core.KafkaAdmin; import org.springframework.kafka.core.KafkaTemplate; import org.springframework.kafka.core.ProducerFactory; import org.springframework.kafka.listener.MessageListenerContainer; +import org.springframework.kafka.support.ProducerListener; import org.springframework.kafka.support.micrometer.KafkaListenerObservation.DefaultKafkaListenerObservationConvention; import org.springframework.kafka.support.micrometer.KafkaTemplateObservation.DefaultKafkaTemplateObservationConvention; import org.springframework.kafka.test.EmbeddedKafkaBroker; @@ -97,9 +101,9 @@ import static org.mockito.Mockito.mock; * @since 3.0 */ @SpringJUnitConfig -@EmbeddedKafka(topics = { ObservationTests.OBSERVATION_TEST_1, ObservationTests.OBSERVATION_TEST_2, +@EmbeddedKafka(topics = {ObservationTests.OBSERVATION_TEST_1, ObservationTests.OBSERVATION_TEST_2, ObservationTests.OBSERVATION_TEST_3, ObservationTests.OBSERVATION_RUNTIME_EXCEPTION, - ObservationTests.OBSERVATION_ERROR }, partitions = 1) + ObservationTests.OBSERVATION_ERROR, ObservationTests.OBSERVATION_TRACEPARENT_DUPLICATE}, partitions = 1) @DirtiesContext public class ObservationTests { @@ -113,6 +117,8 @@ public class ObservationTests { public final static String OBSERVATION_ERROR = "observation.error"; + public final static String OBSERVATION_TRACEPARENT_DUPLICATE = "observation.traceparent.duplicate"; + @Test void endToEnd(@Autowired Listener listener, @Autowired KafkaTemplate template, @Autowired SimpleTracer tracer, @Autowired KafkaListenerEndpointRegistry rler, @@ -120,11 +126,12 @@ public class ObservationTests { @Autowired KafkaListenerEndpointRegistry endpointRegistry, @Autowired KafkaAdmin admin, @Autowired @Qualifier("customTemplate") KafkaTemplate customTemplate, @Autowired Config config) - throws InterruptedException, ExecutionException, TimeoutException { + throws InterruptedException, ExecutionException, TimeoutException { AtomicReference spanFromCallback = new AtomicReference<>(); template.setProducerInterceptor(new ProducerInterceptor<>() { + @Override public ProducerRecord onSend(ProducerRecord record) { tracer.currentSpanCustomizer().tag("key", "value"); @@ -309,10 +316,10 @@ public class ObservationTests { meterRegistryAssert.hasTimerWithNameAndTags("spring.kafka.template", KeyValues.of("spring.kafka.template.name", "template", - "messaging.operation", "publish", - "messaging.system", "kafka", - "messaging.destination.kind", "topic", - "messaging.destination.name", destName) + "messaging.operation", "publish", + "messaging.system", "kafka", + "messaging.destination.kind", "topic", + "messaging.destination.name", destName) .and(keyValues)); } @@ -321,12 +328,12 @@ public class ObservationTests { meterRegistryAssert.hasTimerWithNameAndTags("spring.kafka.listener", KeyValues.of( - "messaging.kafka.consumer.group", consumerGroup, - "messaging.operation", "receive", - "messaging.source.kind", "topic", - "messaging.source.name", destName, - "messaging.system", "kafka", - "spring.kafka.listener.id", listenerId) + "messaging.kafka.consumer.group", consumerGroup, + "messaging.operation", "receive", + "messaging.source.kind", "topic", + "messaging.source.name", destName, + "messaging.system", "kafka", + "spring.kafka.listener.id", listenerId) .and(keyValues)); } @@ -369,7 +376,7 @@ public class ObservationTests { void observationErrorException(@Autowired ExceptionListener listener, @Autowired SimpleTracer tracer, @Autowired @Qualifier("throwableTemplate") KafkaTemplate errorTemplate, @Autowired KafkaListenerEndpointRegistry endpointRegistry) - throws ExecutionException, InterruptedException, TimeoutException { + throws ExecutionException, InterruptedException, TimeoutException { errorTemplate.send(OBSERVATION_ERROR, "testError").get(10, TimeUnit.SECONDS); assertThat(listener.latch5.await(10, TimeUnit.SECONDS)).isTrue(); @@ -394,6 +401,63 @@ public class ObservationTests { assertThat(template.getKafkaAdmin()).isSameAs(kafkaAdmin); } + @Test + void verifyKafkaRecordSenderContextTraceParentHandling() { + String initialTraceParent = "traceparent-from-previous"; + String updatedTraceParent = "traceparent-current"; + ProducerRecord record = new ProducerRecord<>("test-topic", "test-value"); + record.headers().add("traceparent", initialTraceParent.getBytes(StandardCharsets.UTF_8)); + + // Create the context and update the traceparent + KafkaRecordSenderContext context = new KafkaRecordSenderContext( + record, + "test-bean", + () -> "test-cluster" + ); + context.getSetter().set(record, "traceparent", updatedTraceParent); + + Iterable
traceparentHeaders = record.headers().headers("traceparent"); + + List headerValues = StreamSupport.stream(traceparentHeaders.spliterator(), false) + .map(header -> new String(header.value(), StandardCharsets.UTF_8)) + .toList(); + + // Verify there's only one traceparent header and it contains the updated value + assertThat(headerValues).containsExactly(updatedTraceParent); + } + + @Test + void verifyTraceParentHeader(@Autowired KafkaTemplate template, + @Autowired SimpleTracer tracer) throws Exception { + CompletableFuture> producerRecordFuture = new CompletableFuture<>(); + template.setProducerListener(new ProducerListener<>() { + + @Override + public void onSuccess(ProducerRecord producerRecord, RecordMetadata recordMetadata) { + producerRecordFuture.complete(producerRecord); + } + }); + String initialTraceParent = "traceparent-from-previous"; + Header header = new RecordHeader("traceparent", initialTraceParent.getBytes(StandardCharsets.UTF_8)); + ProducerRecord producerRecord = new ProducerRecord<>( + OBSERVATION_TRACEPARENT_DUPLICATE, + null, null, null, + "test-value", + List.of(header) + ); + + template.send(producerRecord).get(10, TimeUnit.SECONDS); + ProducerRecord recordResult = producerRecordFuture.get(10, TimeUnit.SECONDS); + + Iterable
traceparentHeaders = recordResult.headers().headers("traceparent"); + assertThat(traceparentHeaders).hasSize(1); + + String traceparentValue = new String(traceparentHeaders.iterator().next().value(), StandardCharsets.UTF_8); + assertThat(traceparentValue).isEqualTo("traceparent-from-propagator"); + + tracer.getSpans().clear(); + } + @Configuration @EnableKafka public static class Config { @@ -523,6 +587,9 @@ public class ObservationTests { public void inject(TraceContext context, @Nullable C carrier, Setter setter) { setter.set(carrier, "foo", "some foo value"); setter.set(carrier, "bar", "some bar value"); + + // Add a traceparent header to simulate W3C trace context + setter.set(carrier, "traceparent", "traceparent-from-propagator"); } // This is called on the consumer side when the message is consumed @@ -531,7 +598,9 @@ public class ObservationTests { public Span.Builder extract(C carrier, Getter getter) { String foo = getter.get(carrier, "foo"); String bar = getter.get(carrier, "bar"); - return tracer.spanBuilder().tag("foo", foo).tag("bar", bar); + return tracer.spanBuilder() + .tag("foo", foo) + .tag("bar", bar); } }; }