Implement ExtendedTextMapGetter in kafka-clients instrumentation (#13068)

xiepuhuan · web-flow · commit ab09fcee98c3 · 2025-01-20T17:43:48.000-08:00
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientDefaultTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientDefaultTest.java
@@ -108,7 +108,7 @@ void testKafkaProducerAndConsumerSpan(boolean testHeaders) throws Exception {
                         .hasLinks(LinkData.create(producerSpan.get().getSpanContext()))
                         .hasParent(trace.getSpan(0))
                         .hasAttributesSatisfyingExactly(
-                            processAttributes("10", greeting, testHeaders)),
+                            processAttributes("10", greeting, testHeaders, false)),
                 span -> span.hasName("processing").hasParent(trace.getSpan(1))));
   }
 
@@ -152,7 +152,8 @@ void testPassThroughTombstone()
                         .hasKind(SpanKind.CONSUMER)
                         .hasLinks(LinkData.create(producerSpan.get().getSpanContext()))
                         .hasParent(trace.getSpan(0))
-                        .hasAttributesSatisfyingExactly(processAttributes(null, null, false))));
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes(null, null, false, false))));
   }
 
   @DisplayName("test records(TopicPartition) kafka consume")
@@ -203,6 +204,7 @@ void testRecordsWithTopicPartitionKafkaConsume()
                         .hasKind(SpanKind.CONSUMER)
                         .hasLinks(LinkData.create(producerSpan.get().getSpanContext()))
                         .hasParent(trace.getSpan(0))
-                        .hasAttributesSatisfyingExactly(processAttributes(null, greeting, false))));
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes(null, greeting, false, false))));
   }
 }
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientPropagationDisabledTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientPropagationDisabledTest.java
@@ -64,7 +64,8 @@ void testReadRemoteContextWhenPropagationIsDisabled() throws InterruptedExceptio
                     span.hasName(SHARED_TOPIC + " process")
                         .hasKind(SpanKind.CONSUMER)
                         .hasLinks(Collections.emptyList())
-                        .hasAttributesSatisfyingExactly(processAttributes(null, message, false)),
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes(null, message, false, false)),
                 span -> span.hasName("processing").hasParent(trace.getSpan(0))));
   }
 }
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientSuppressReceiveSpansTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientSuppressReceiveSpansTest.java
@@ -12,6 +12,7 @@
 import io.opentelemetry.instrumentation.kafka.internal.KafkaClientPropagationBaseTest;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
@@ -33,8 +34,19 @@ void testKafkaProduceAndConsume() throws InterruptedException {
     testing.runWithSpan(
         "parent",
         () -> {
+          ProducerRecord<Integer, String> producerRecord =
+              new ProducerRecord<>(SHARED_TOPIC, 10, greeting);
+          producerRecord
+              .headers()
+              // adding baggage header in w3c baggage format
+              .add(
+                  "baggage",
+                  "test-baggage-key-1=test-baggage-value-1".getBytes(StandardCharsets.UTF_8))
+              .add(
+                  "baggage",
+                  "test-baggage-key-2=test-baggage-value-2".getBytes(StandardCharsets.UTF_8));
           producer.send(
-              new ProducerRecord<>(SHARED_TOPIC, 10, greeting),
+              producerRecord,
               (meta, ex) -> {
                 if (ex == null) {
                   testing.runWithSpan("producer callback", () -> {});
@@ -70,7 +82,8 @@ void testKafkaProduceAndConsume() throws InterruptedException {
                     span.hasName(SHARED_TOPIC + " process")
                         .hasKind(SpanKind.CONSUMER)
                         .hasParent(trace.getSpan(1))
-                        .hasAttributesSatisfyingExactly(processAttributes("10", greeting, false)),
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes("10", greeting, false, true)),
                 span ->
                     span.hasName("processing")
                         .hasKind(SpanKind.INTERNAL)
@@ -108,7 +121,8 @@ void testPassThroughTombstone()
                     span.hasName(SHARED_TOPIC + " process")
                         .hasKind(SpanKind.CONSUMER)
                         .hasParent(trace.getSpan(0))
-                        .hasAttributesSatisfyingExactly(processAttributes(null, null, false))));
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes(null, null, false, false))));
   }
 
   @Test
@@ -146,6 +160,7 @@ void testRecordsWithTopicPartitionKafkaConsume()
                     span.hasName(SHARED_TOPIC + " process")
                         .hasKind(SpanKind.CONSUMER)
                         .hasParent(trace.getSpan(0))
-                        .hasAttributesSatisfyingExactly(processAttributes(null, greeting, false))));
+                        .hasAttributesSatisfyingExactly(
+                            processAttributes(null, greeting, false, false))));
   }
 }
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafka/internal/KafkaClientBaseTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafka/internal/KafkaClientBaseTest.java
@@ -216,7 +216,7 @@ protected static List<AttributeAssertion> receiveAttributes(boolean testHeaders)
 
   @SuppressWarnings("deprecation") // using deprecated semconv
   protected static List<AttributeAssertion> processAttributes(
-      String messageKey, String messageValue, boolean testHeaders) {
+      String messageKey, String messageValue, boolean testHeaders, boolean testMultiBaggage) {
     List<AttributeAssertion> assertions =
         new ArrayList<>(
             Arrays.asList(
@@ -249,6 +249,11 @@ protected static List<AttributeAssertion> processAttributes(
               AttributeKey.stringArrayKey("messaging.header.test_message_header"),
               Collections.singletonList("test")));
     }
+
+    if (testMultiBaggage) {
+      assertions.add(equalTo(AttributeKey.stringKey("test-baggage-key-1"), "test-baggage-value-1"));
+      assertions.add(equalTo(AttributeKey.stringKey("test-baggage-key-2"), "test-baggage-value-2"));
+    }
     return assertions;
   }
 }
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/test/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/AbstractInterceptorsTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/test/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/AbstractInterceptorsTest.java
@@ -10,6 +10,7 @@
 import io.opentelemetry.instrumentation.kafka.internal.KafkaClientBaseTest;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Map;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -48,8 +49,19 @@ void testInterceptors() throws InterruptedException {
     testing.runWithSpan(
         "parent",
         () -> {
+          ProducerRecord<Integer, String> producerRecord =
+              new ProducerRecord<>(SHARED_TOPIC, greeting);
+          producerRecord
+              .headers()
+              // adding baggage header in w3c baggage format
+              .add(
+                  "baggage",
+                  "test-baggage-key-1=test-baggage-value-1".getBytes(StandardCharsets.UTF_8))
+              .add(
+                  "baggage",
+                  "test-baggage-key-2=test-baggage-value-2".getBytes(StandardCharsets.UTF_8));
           producer.send(
-              new ProducerRecord<>(SHARED_TOPIC, greeting),
+              producerRecord,
               (meta, ex) -> {
                 if (ex == null) {
                   testing.runWithSpan("producer callback", () -> {});
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/test/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/InterceptorsSuppressReceiveSpansTest.java b/instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/test/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/InterceptorsSuppressReceiveSpansTest.java
@@ -15,6 +15,7 @@
 import static io.opentelemetry.semconv.incubating.MessagingIncubatingAttributes.MESSAGING_OPERATION;
 import static io.opentelemetry.semconv.incubating.MessagingIncubatingAttributes.MESSAGING_SYSTEM;
 
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.SpanKind;
 import java.nio.charset.StandardCharsets;
 import org.assertj.core.api.AbstractLongAssert;
@@ -59,7 +60,13 @@ void assertTraces() {
                             equalTo(MESSAGING_KAFKA_CONSUMER_GROUP, "test"),
                             satisfies(
                                 MESSAGING_CLIENT_ID,
-                                stringAssert -> stringAssert.startsWith("consumer"))),
+                                stringAssert -> stringAssert.startsWith("consumer")),
+                            equalTo(
+                                AttributeKey.stringKey("test-baggage-key-1"),
+                                "test-baggage-value-1"),
+                            equalTo(
+                                AttributeKey.stringKey("test-baggage-key-2"),
+                                "test-baggage-value-2")),
                 span ->
                     span.hasName("process child")
                         .hasKind(SpanKind.INTERNAL)
diff --git a/instrumentation/kafka/kafka-clients/kafka-clients-common/library/src/main/java/io/opentelemetry/instrumentation/kafka/internal/KafkaConsumerRecordGetter.java b/instrumentation/kafka/kafka-clients/kafka-clients-common/library/src/main/java/io/opentelemetry/instrumentation/kafka/internal/KafkaConsumerRecordGetter.java
@@ -5,14 +5,15 @@
 
 package io.opentelemetry.instrumentation.kafka.internal;
 
-import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.internal.ExtendedTextMapGetter;
 import java.nio.charset.StandardCharsets;
+import java.util.Iterator;
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 import javax.annotation.Nullable;
 import org.apache.kafka.common.header.Header;
 
-enum KafkaConsumerRecordGetter implements TextMapGetter<KafkaProcessRequest> {
+enum KafkaConsumerRecordGetter implements ExtendedTextMapGetter<KafkaProcessRequest> {
   INSTANCE;
 
   @Override
@@ -35,4 +36,11 @@ public String get(@Nullable KafkaProcessRequest carrier, String key) {
     }
     return new String(value, StandardCharsets.UTF_8);
   }
+
+  @Override
+  public Iterator<String> getAll(@Nullable KafkaProcessRequest carrier, String key) {
+    return StreamSupport.stream(carrier.getRecord().headers().headers(key).spliterator(), false)
+        .map(header -> new String(header.value(), StandardCharsets.UTF_8))
+        .iterator();
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ void testKafkaProducerAndConsumerSpan(boolean testHeaders) throws Exception {`
`108`	`108`	`.hasLinks(LinkData.create(producerSpan.get().getSpanContext()))`
`109`	`109`	`.hasParent(trace.getSpan(0))`
`110`	`110`	`.hasAttributesSatisfyingExactly(`
`111`		`- processAttributes("10", greeting, testHeaders)),`
	`111`	`+ processAttributes("10", greeting, testHeaders, false)),`
`112`	`112`	`span -> span.hasName("processing").hasParent(trace.getSpan(1))));`
`113`	`113`	`}`
`114`	`114`
`@@ -152,7 +152,8 @@ void testPassThroughTombstone()`
`152`	`152`	`.hasKind(SpanKind.CONSUMER)`
`153`	`153`	`.hasLinks(LinkData.create(producerSpan.get().getSpanContext()))`
`154`	`154`	`.hasParent(trace.getSpan(0))`
`155`		`- .hasAttributesSatisfyingExactly(processAttributes(null, null, false))));`
	`155`	`+ .hasAttributesSatisfyingExactly(`
	`156`	`+ processAttributes(null, null, false, false))));`
`156`	`157`	`}`
`157`	`158`
`158`	`159`	`@DisplayName("test records(TopicPartition) kafka consume")`
`@@ -203,6 +204,7 @@ void testRecordsWithTopicPartitionKafkaConsume()`
`203`	`204`	`.hasKind(SpanKind.CONSUMER)`
`204`	`205`	`.hasLinks(LinkData.create(producerSpan.get().getSpanContext()))`
`205`	`206`	`.hasParent(trace.getSpan(0))`
`206`		`- .hasAttributesSatisfyingExactly(processAttributes(null, greeting, false))));`
	`207`	`+ .hasAttributesSatisfyingExactly(`
	`208`	`+ processAttributes(null, greeting, false, false))));`
`207`	`209`	`}`
`208`	`210`	`}`
Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,8 @@ void testReadRemoteContextWhenPropagationIsDisabled() throws InterruptedExceptio`
`64`	`64`	`span.hasName(SHARED_TOPIC + " process")`
`65`	`65`	`.hasKind(SpanKind.CONSUMER)`
`66`	`66`	`.hasLinks(Collections.emptyList())`
`67`		`- .hasAttributesSatisfyingExactly(processAttributes(null, message, false)),`
	`67`	`+ .hasAttributesSatisfyingExactly(`
	`68`	`+ processAttributes(null, message, false, false)),`
`68`	`69`	`span -> span.hasName("processing").hasParent(trace.getSpan(0))));`
`69`	`70`	`}`
`70`	`71`	`}`
Original file line number	Diff line number	Diff line change
`@@ -216,7 +216,7 @@ protected static List<AttributeAssertion> receiveAttributes(boolean testHeaders)`
`216`	`216`
`217`	`217`	`@SuppressWarnings("deprecation") // using deprecated semconv`
`218`	`218`	`protected static List<AttributeAssertion> processAttributes(`
`219`		`- String messageKey, String messageValue, boolean testHeaders) {`
	`219`	`+ String messageKey, String messageValue, boolean testHeaders, boolean testMultiBaggage) {`
`220`	`220`	`List<AttributeAssertion> assertions =`
`221`	`221`	`new ArrayList<>(`
`222`	`222`	`Arrays.asList(`
`@@ -249,6 +249,11 @@ protected static List<AttributeAssertion> processAttributes(`
`249`	`249`	`AttributeKey.stringArrayKey("messaging.header.test_message_header"),`
`250`	`250`	`Collections.singletonList("test")));`
`251`	`251`	`}`
	`252`	`+`
	`253`	`+ if (testMultiBaggage) {`
	`254`	`+ assertions.add(equalTo(AttributeKey.stringKey("test-baggage-key-1"), "test-baggage-value-1"));`
	`255`	`+ assertions.add(equalTo(AttributeKey.stringKey("test-baggage-key-2"), "test-baggage-value-2"));`
	`256`	`+ }`
`252`	`257`	`return assertions;`
`253`	`258`	`}`
`254`	`259`	`}`