1616
1717package org .springframework .kafka .support .micrometer ;
1818
19+ import java .nio .charset .StandardCharsets ;
1920import java .util .Arrays ;
2021import java .util .Deque ;
2122import java .util .List ;
2526import java .util .concurrent .TimeUnit ;
2627import java .util .concurrent .TimeoutException ;
2728import java .util .concurrent .atomic .AtomicReference ;
29+ import java .util .stream .StreamSupport ;
2830
2931import io .micrometer .common .KeyValues ;
3032import io .micrometer .core .instrument .MeterRegistry ;
5355import org .apache .kafka .common .errors .InvalidTopicException ;
5456import org .apache .kafka .common .header .Header ;
5557import org .apache .kafka .common .header .Headers ;
58+ import org .apache .kafka .common .header .internals .RecordHeader ;
5659import org .junit .jupiter .api .Test ;
5760
5861import org .springframework .beans .factory .annotation .Autowired ;
7275import org .springframework .kafka .core .KafkaTemplate ;
7376import org .springframework .kafka .core .ProducerFactory ;
7477import org .springframework .kafka .listener .MessageListenerContainer ;
78+ import org .springframework .kafka .support .ProducerListener ;
7579import org .springframework .kafka .support .micrometer .KafkaListenerObservation .DefaultKafkaListenerObservationConvention ;
7680import org .springframework .kafka .support .micrometer .KafkaTemplateObservation .DefaultKafkaTemplateObservationConvention ;
7781import org .springframework .kafka .test .EmbeddedKafkaBroker ;
97101 * @since 3.0
98102 */
99103@ SpringJUnitConfig
100- @ EmbeddedKafka (topics = { ObservationTests .OBSERVATION_TEST_1 , ObservationTests .OBSERVATION_TEST_2 ,
104+ @ EmbeddedKafka (topics = {ObservationTests .OBSERVATION_TEST_1 , ObservationTests .OBSERVATION_TEST_2 ,
101105 ObservationTests .OBSERVATION_TEST_3 , ObservationTests .OBSERVATION_RUNTIME_EXCEPTION ,
102- ObservationTests .OBSERVATION_ERROR }, partitions = 1 )
106+ ObservationTests .OBSERVATION_ERROR , ObservationTests . OBSERVATION_TRACEPARENT_DUPLICATE }, partitions = 1 )
103107@ DirtiesContext
104108public class ObservationTests {
105109
@@ -113,18 +117,21 @@ public class ObservationTests {
113117
114118 public final static String OBSERVATION_ERROR = "observation.error" ;
115119
120+ public final static String OBSERVATION_TRACEPARENT_DUPLICATE = "observation.traceparent.duplicate" ;
121+
116122 @ Test
117123 void endToEnd (@ Autowired Listener listener , @ Autowired KafkaTemplate <Integer , String > template ,
118124 @ Autowired SimpleTracer tracer , @ Autowired KafkaListenerEndpointRegistry rler ,
119125 @ Autowired MeterRegistry meterRegistry , @ Autowired EmbeddedKafkaBroker broker ,
120126 @ Autowired KafkaListenerEndpointRegistry endpointRegistry , @ Autowired KafkaAdmin admin ,
121127 @ Autowired @ Qualifier ("customTemplate" ) KafkaTemplate <Integer , String > customTemplate ,
122128 @ Autowired Config config )
123- throws InterruptedException , ExecutionException , TimeoutException {
129+ throws InterruptedException , ExecutionException , TimeoutException {
124130
125131 AtomicReference <SimpleSpan > spanFromCallback = new AtomicReference <>();
126132
127133 template .setProducerInterceptor (new ProducerInterceptor <>() {
134+
128135 @ Override
129136 public ProducerRecord <Integer , String > onSend (ProducerRecord <Integer , String > record ) {
130137 tracer .currentSpanCustomizer ().tag ("key" , "value" );
@@ -309,10 +316,10 @@ private void assertThatTemplateHasTimerWithNameAndTags(MeterRegistryAssert meter
309316
310317 meterRegistryAssert .hasTimerWithNameAndTags ("spring.kafka.template" ,
311318 KeyValues .of ("spring.kafka.template.name" , "template" ,
312- "messaging.operation" , "publish" ,
313- "messaging.system" , "kafka" ,
314- "messaging.destination.kind" , "topic" ,
315- "messaging.destination.name" , destName )
319+ "messaging.operation" , "publish" ,
320+ "messaging.system" , "kafka" ,
321+ "messaging.destination.kind" , "topic" ,
322+ "messaging.destination.name" , destName )
316323 .and (keyValues ));
317324 }
318325
@@ -321,12 +328,12 @@ private void assertThatListenerHasTimerWithNameAndTags(MeterRegistryAssert meter
321328
322329 meterRegistryAssert .hasTimerWithNameAndTags ("spring.kafka.listener" ,
323330 KeyValues .of (
324- "messaging.kafka.consumer.group" , consumerGroup ,
325- "messaging.operation" , "receive" ,
326- "messaging.source.kind" , "topic" ,
327- "messaging.source.name" , destName ,
328- "messaging.system" , "kafka" ,
329- "spring.kafka.listener.id" , listenerId )
331+ "messaging.kafka.consumer.group" , consumerGroup ,
332+ "messaging.operation" , "receive" ,
333+ "messaging.source.kind" , "topic" ,
334+ "messaging.source.name" , destName ,
335+ "messaging.system" , "kafka" ,
336+ "spring.kafka.listener.id" , listenerId )
330337 .and (keyValues ));
331338 }
332339
@@ -369,7 +376,7 @@ void observationRuntimeException(@Autowired ExceptionListener listener, @Autowir
369376 void observationErrorException (@ Autowired ExceptionListener listener , @ Autowired SimpleTracer tracer ,
370377 @ Autowired @ Qualifier ("throwableTemplate" ) KafkaTemplate <Integer , String > errorTemplate ,
371378 @ Autowired KafkaListenerEndpointRegistry endpointRegistry )
372- throws ExecutionException , InterruptedException , TimeoutException {
379+ throws ExecutionException , InterruptedException , TimeoutException {
373380
374381 errorTemplate .send (OBSERVATION_ERROR , "testError" ).get (10 , TimeUnit .SECONDS );
375382 assertThat (listener .latch5 .await (10 , TimeUnit .SECONDS )).isTrue ();
@@ -394,6 +401,63 @@ void kafkaAdminNotRecreatedIfBootstrapServersSameInProducerAndAdminConfig(
394401 assertThat (template .getKafkaAdmin ()).isSameAs (kafkaAdmin );
395402 }
396403
404+ @ Test
405+ void verifyKafkaRecordSenderContextTraceParentHandling () {
406+ String initialTraceParent = "traceparent-from-previous" ;
407+ String updatedTraceParent = "traceparent-current" ;
408+ ProducerRecord <Integer , String > record = new ProducerRecord <>("test-topic" , "test-value" );
409+ record .headers ().add ("traceparent" , initialTraceParent .getBytes (StandardCharsets .UTF_8 ));
410+
411+ // Create the context and update the traceparent
412+ KafkaRecordSenderContext context = new KafkaRecordSenderContext (
413+ record ,
414+ "test-bean" ,
415+ () -> "test-cluster"
416+ );
417+ context .getSetter ().set (record , "traceparent" , updatedTraceParent );
418+
419+ Iterable <Header > traceparentHeaders = record .headers ().headers ("traceparent" );
420+
421+ List <String > headerValues = StreamSupport .stream (traceparentHeaders .spliterator (), false )
422+ .map (header -> new String (header .value (), StandardCharsets .UTF_8 ))
423+ .toList ();
424+
425+ // Verify there's only one traceparent header and it contains the updated value
426+ assertThat (headerValues ).containsExactly (updatedTraceParent );
427+ }
428+
429+ @ Test
430+ void verifyTraceParentHeader (@ Autowired KafkaTemplate <Integer , String > template ,
431+ @ Autowired SimpleTracer tracer ) throws Exception {
432+ CompletableFuture <ProducerRecord <Integer , String >> producerRecordFuture = new CompletableFuture <>();
433+ template .setProducerListener (new ProducerListener <>() {
434+
435+ @ Override
436+ public void onSuccess (ProducerRecord <Integer , String > producerRecord , RecordMetadata recordMetadata ) {
437+ producerRecordFuture .complete (producerRecord );
438+ }
439+ });
440+ String initialTraceParent = "traceparent-from-previous" ;
441+ Header header = new RecordHeader ("traceparent" , initialTraceParent .getBytes (StandardCharsets .UTF_8 ));
442+ ProducerRecord <Integer , String > producerRecord = new ProducerRecord <>(
443+ OBSERVATION_TRACEPARENT_DUPLICATE ,
444+ null , null , null ,
445+ "test-value" ,
446+ List .of (header )
447+ );
448+
449+ template .send (producerRecord ).get (10 , TimeUnit .SECONDS );
450+ ProducerRecord <Integer , String > recordResult = producerRecordFuture .get (10 , TimeUnit .SECONDS );
451+
452+ Iterable <Header > traceparentHeaders = recordResult .headers ().headers ("traceparent" );
453+ assertThat (traceparentHeaders ).hasSize (1 );
454+
455+ String traceparentValue = new String (traceparentHeaders .iterator ().next ().value (), StandardCharsets .UTF_8 );
456+ assertThat (traceparentValue ).isEqualTo ("traceparent-from-propagator" );
457+
458+ tracer .getSpans ().clear ();
459+ }
460+
397461 @ Configuration
398462 @ EnableKafka
399463 public static class Config {
@@ -523,6 +587,9 @@ public List<String> fields() {
523587 public <C > void inject (TraceContext context , @ Nullable C carrier , Setter <C > setter ) {
524588 setter .set (carrier , "foo" , "some foo value" );
525589 setter .set (carrier , "bar" , "some bar value" );
590+
591+ // Add a traceparent header to simulate W3C trace context
592+ setter .set (carrier , "traceparent" , "traceparent-from-propagator" );
526593 }
527594
528595 // This is called on the consumer side when the message is consumed
@@ -531,7 +598,9 @@ public <C> void inject(TraceContext context, @Nullable C carrier, Setter<C> sett
531598 public <C > Span .Builder extract (C carrier , Getter <C > getter ) {
532599 String foo = getter .get (carrier , "foo" );
533600 String bar = getter .get (carrier , "bar" );
534- return tracer .spanBuilder ().tag ("foo" , foo ).tag ("bar" , bar );
601+ return tracer .spanBuilder ()
602+ .tag ("foo" , foo )
603+ .tag ("bar" , bar );
535604 }
536605 };
537606 }
0 commit comments