Refactor kafka interceptors to be configurable

Co-authored-by: trask <[email protected]>

Author: Copilot

conventions/.kotlin/sessions/kotlin-compiler-16318826606158418015.salive

@@ -21,6 +21,7 @@
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaProcessRequest;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaProducerRequest;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaReceiveRequest;
+import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaTelemetrySupplier;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaUtil;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.MetricsReporterList;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.OpenTelemetryMetricsReporter;
@@ -207,6 +208,61 @@ <K, V> ConsumerRecords<K, V> addTracing(
     return Collections.unmodifiableMap(config);
   }
 
+  /**
+   * Returns a config property key that can be used to pass this {@link KafkaTelemetry} instance to
+   * interceptors.
+   *
+   * <p>This is an internal config key used by the library instrumentation interceptors.
+   */
+  static final String CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER =
+      "opentelemetry.experimental.kafka-telemetry.supplier";
+
+  /**
+   * Produces a set of kafka producer config properties to configure tracing in {@code
+   * TracingProducerInterceptor}. Add these resulting properties to the configuration map used to
+   * initialize a {@link KafkaProducer}.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * //    Map<String, Object> config = new HashMap<>();
+   * //    config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, ...);
+   * //    config.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingProducerInterceptor.class.getName());
+   * //    config.putAll(kafkaTelemetry.producerInterceptorConfigProperties());
+   * //    try (KafkaProducer<?, ?> producer = new KafkaProducer<>(config)) { ... }
+   * }</pre>
+   *
+   * @return the kafka producer interceptor config properties
+   */
+  public Map<String, ?> producerInterceptorConfigProperties() {
+    Map<String, Object> config = new HashMap<>();
+    config.put(CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER, new KafkaTelemetrySupplier(this));
+    return Collections.unmodifiableMap(config);
+  }
+
+  /**
+   * Produces a set of kafka consumer config properties to configure tracing in {@code
+   * TracingConsumerInterceptor}. Add these resulting properties to the configuration map used to
+   * initialize a {@link KafkaConsumer}.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * //    Map<String, Object> config = new HashMap<>();
+   * //    config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, ...);
+   * //    config.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingConsumerInterceptor.class.getName());
+   * //    config.putAll(kafkaTelemetry.consumerInterceptorConfigProperties());
+   * //    try (KafkaConsumer<?, ?> consumer = new KafkaConsumer<>(config)) { ... }
+   * }</pre>
+   *
+   * @return the kafka consumer interceptor config properties
+   */
+  public Map<String, ?> consumerInterceptorConfigProperties() {
+    Map<String, Object> config = new HashMap<>();
+    config.put(CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER, new KafkaTelemetrySupplier(this));
+    return Collections.unmodifiableMap(config);
+  }
+
   /**
    * Build and inject span into record.
    *

gradle-plugins/.kotlin/sessions/kotlin-compiler-16068091306083868015.salive

@@ -16,6 +16,8 @@
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaConsumerContextUtil;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerInterceptor;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
@@ -29,22 +31,16 @@
  */
 public class TracingConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
 
-  private static final KafkaTelemetry telemetry =
-      KafkaTelemetry.builder(GlobalOpenTelemetry.get())
-          .setMessagingReceiveInstrumentationEnabled(
-              ConfigPropertiesUtil.getBoolean(
-                  "otel.instrumentation.messaging.experimental.receive-telemetry.enabled", false))
-          .setCapturedHeaders(
-              ConfigPropertiesUtil.getList(
-                  "otel.instrumentation.messaging.experimental.capture-headers", emptyList()))
-          .build();
-
+  @Nullable private KafkaTelemetry telemetry;
   private String consumerGroup;
   private String clientId;
 
   @Override
   @CanIgnoreReturnValue
   public ConsumerRecords<K, V> onConsume(ConsumerRecords<K, V> records) {
+    if (telemetry == null) {
+      return records;
+    }
     // timer should be started before fetching ConsumerRecords, but there is no callback for that
     Timer timer = Timer.start();
     Context receiveContext = telemetry.buildAndFinishSpan(records, consumerGroup, clientId, timer);
@@ -67,6 +63,25 @@ public void configure(Map<String, ?> configs) {
     consumerGroup = Objects.toString(configs.get(ConsumerConfig.GROUP_ID_CONFIG), null);
     clientId = Objects.toString(configs.get(ConsumerConfig.CLIENT_ID_CONFIG), null);
 
-    // TODO: support experimental attributes config
+    // Try to get KafkaTelemetry from config
+    Object telemetrySupplier = configs.get(KafkaTelemetry.CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER);
+    if (telemetrySupplier instanceof Supplier) {
+      Object kafkaTelemetry = ((Supplier<?>) telemetrySupplier).get();
+      if (kafkaTelemetry instanceof KafkaTelemetry) {
+        this.telemetry = (KafkaTelemetry) kafkaTelemetry;
+        return;
+      }
+    }
+
+    // Fallback to GlobalOpenTelemetry if not configured
+    this.telemetry =
+        KafkaTelemetry.builder(GlobalOpenTelemetry.get())
+            .setMessagingReceiveInstrumentationEnabled(
+                ConfigPropertiesUtil.getBoolean(
+                    "otel.instrumentation.messaging.experimental.receive-telemetry.enabled", false))
+            .setCapturedHeaders(
+                ConfigPropertiesUtil.getList(
+                    "otel.instrumentation.messaging.experimental.capture-headers", emptyList()))
+            .build();
   }
 }

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/KafkaTelemetry.java

@@ -12,6 +12,7 @@
 import io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerInterceptor;
@@ -25,19 +26,15 @@
  */
 public class TracingProducerInterceptor<K, V> implements ProducerInterceptor<K, V> {
 
-  private static final KafkaTelemetry telemetry =
-      KafkaTelemetry.builder(GlobalOpenTelemetry.get())
-          .setCapturedHeaders(
-              ConfigPropertiesUtil.getList(
-                  "otel.instrumentation.messaging.experimental.capture-headers", emptyList()))
-          .build();
-
+  @Nullable private KafkaTelemetry telemetry;
   @Nullable private String clientId;
 
   @Override
   @CanIgnoreReturnValue
   public ProducerRecord<K, V> onSend(ProducerRecord<K, V> producerRecord) {
-    telemetry.buildAndInjectSpan(producerRecord, clientId);
+    if (telemetry != null) {
+      telemetry.buildAndInjectSpan(producerRecord, clientId);
+    }
     return producerRecord;
   }
 
@@ -48,9 +45,25 @@ public void onAcknowledgement(RecordMetadata recordMetadata, Exception e) {}
   public void close() {}
 
   @Override
-  public void configure(Map<String, ?> map) {
-    clientId = Objects.toString(map.get(ProducerConfig.CLIENT_ID_CONFIG), null);
+  public void configure(Map<String, ?> configs) {
+    clientId = Objects.toString(configs.get(ProducerConfig.CLIENT_ID_CONFIG), null);
+
+    // Try to get KafkaTelemetry from config
+    Object telemetrySupplier = configs.get(KafkaTelemetry.CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER);
+    if (telemetrySupplier instanceof Supplier) {
+      Object kafkaTelemetry = ((Supplier<?>) telemetrySupplier).get();
+      if (kafkaTelemetry instanceof KafkaTelemetry) {
+        this.telemetry = (KafkaTelemetry) kafkaTelemetry;
+        return;
+      }
+    }
 
-    // TODO: support experimental attributes config
+    // Fallback to GlobalOpenTelemetry if not configured
+    this.telemetry =
+        KafkaTelemetry.builder(GlobalOpenTelemetry.get())
+            .setCapturedHeaders(
+                ConfigPropertiesUtil.getList(
+                    "otel.instrumentation.messaging.experimental.capture-headers", emptyList()))
+            .build();
   }
 }

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/TracingConsumerInterceptor.java

@@ -28,19 +28,36 @@ abstract class AbstractInterceptorsTest extends KafkaClientBaseTest {
 
   static final String greeting = "Hello Kafka!";
 
+  private static KafkaTelemetry kafkaTelemetry;
+
+  protected static KafkaTelemetry createKafkaTelemetry() {
+    return KafkaTelemetry.builder(testing.getOpenTelemetry())
+        .setCapturedHeaders(java.util.Collections.singletonList("Test-Message-Header"))
+        .setMessagingReceiveInstrumentationEnabled(true)
+        .build();
+  }
+
   @Override
   public Map<String, Object> producerProps() {
+    if (kafkaTelemetry == null) {
+      kafkaTelemetry = createKafkaTelemetry();
+    }
     Map<String, Object> props = super.producerProps();
     props.put(
         ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingProducerInterceptor.class.getName());
+    props.putAll(kafkaTelemetry.producerInterceptorConfigProperties());
     return props;
   }
 
   @Override
   public Map<String, Object> consumerProps() {
+    if (kafkaTelemetry == null) {
+      kafkaTelemetry = createKafkaTelemetry();
+    }
     Map<String, Object> props = super.consumerProps();
     props.put(
         ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingConsumerInterceptor.class.getName());
+    props.putAll(kafkaTelemetry.consumerInterceptorConfigProperties());
     return props;
   }
 

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/TracingProducerInterceptor.java

@@ -23,6 +23,12 @@
 
 class InterceptorsSuppressReceiveSpansTest extends AbstractInterceptorsTest {
 
+  protected static KafkaTelemetry createKafkaTelemetry() {
+    return KafkaTelemetry.builder(testing.getOpenTelemetry())
+        .setMessagingReceiveInstrumentationEnabled(false)
+        .build();
+  }
+
   @SuppressWarnings("deprecation") // using deprecated semconv
   @Override
   void assertTraces() {

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/src/test/java/io/opentelemetry/instrumentation/kafkaclients/v2_6/AbstractInterceptorsTest.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal;
+
+import java.io.Serializable;
+import java.util.Objects;
+import java.util.function.Supplier;
+
+/**
+ * Wrapper for KafkaTelemetry that can be injected into kafka configuration without breaking
+ * serialization. Used by kafka interceptors to get a configured KafkaTelemetry instance.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public final class KafkaTelemetrySupplier implements Supplier<Object>, Serializable {
+  private static final long serialVersionUID = 1L;
+  private final transient Object kafkaTelemetry;
+
+  public KafkaTelemetrySupplier(Object kafkaTelemetry) {
+    Objects.requireNonNull(kafkaTelemetry);
+    this.kafkaTelemetry = kafkaTelemetry;
+  }
+
+  @Override
+  public Object get() {
+    return kafkaTelemetry;
+  }
+
+  private Object writeReplace() {
+    // serialize this object to null
+    return null;
+  }
+}