Add configurable OpenTelemetry kafka-clients interceptors

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/README.md

@@ -31,20 +31,31 @@ There are two options for capturing traces, either using interceptors or wrappin
 
 #### Using interceptors
 
-The Kafka clients API provides a way to "intercept" messages before they are sent to the brokers as well as messages received from the broker before being passed to the application.
-The OpenTelemetry instrumented Kafka library provides two interceptors to be configured to add tracing information automatically.
-The interceptor class has to be set in the properties bag used to create the Kafka client.
+The Kafka clients API provides a way to intercept messages before they are sent to the brokers as well as messages received from the broker before being passed to the application.
 
-Use the `TracingProducerInterceptor` for the producer in order to create a "send" span automatically, each time a message is sent.
+To intercept messages and emit telemetry for a `KafkaProducer`:
 
 ```java
-props.setProperty(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingProducerInterceptor.class.getName());
+KafkaTelemetry telemetry = KafkaTelemetry.create(openTelemetry);
+
+Map<String, Object> props = new HashMap<>();
+props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+props.putAll(telemetry.producerInterceptorConfigProperties());
+
+Producer<String, String> producer = new KafkaProducer<>(props);
 ```
 
-Use the `TracingConsumerInterceptor` for the consumer in order to create a "receive" span automatically, each time a message is received.
+To intercept messages and emit telemetry for a `KafkaConsumer`:
 
 ```java
-props.setProperty(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, TracingConsumerInterceptor.class.getName());
+KafkaTelemetry telemetry = KafkaTelemetry.create(openTelemetry);
+
+Map<String, Object> props = new HashMap<>();
+props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+props.put(ConsumerConfig.GROUP_ID_CONFIG, "my-group");
+props.putAll(telemetry.consumerInterceptorConfigProperties());
+
+Consumer<String, String> consumer = new KafkaConsumer<>(props);
 ```
 
 #### Wrapping clients

instrumentation/kafka/kafka-clients/kafka-clients-2.6/library/build.gradle.kts

@@ -22,27 +22,33 @@ tasks {
     systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
   }
 
-  val testReceiveSpansDisabled by registering(Test::class) {
+  test {
+    filter {
+      excludeTestsMatching("*Deprecated*")
+    }
+  }
+
+  val testDeprecated by registering(Test::class) {
     testClassesDirs = sourceSets.test.get().output.classesDirs
     classpath = sourceSets.test.get().runtimeClasspath
     filter {
-      includeTestsMatching("InterceptorsSuppressReceiveSpansTest")
-      includeTestsMatching("WrapperSuppressReceiveSpansTest")
+      includeTestsMatching("*DeprecatedInterceptorsTest")
     }
-    include("**/InterceptorsSuppressReceiveSpansTest.*", "**/WrapperSuppressReceiveSpansTest.*")
+    forkEvery = 1 // to avoid system properties polluting other tests
+    systemProperty("otel.instrumentation.messaging.experimental.receive-telemetry.enabled", "true")
+    systemProperty("otel.instrumentation.messaging.experimental.capture-headers", "Test-Message-Header")
   }
 
-  test {
+  val testDeprecatedSuppressReceiveSpans by registering(Test::class) {
+    testClassesDirs = sourceSets.test.get().output.classesDirs
+    classpath = sourceSets.test.get().runtimeClasspath
     filter {
-      excludeTestsMatching("InterceptorsSuppressReceiveSpansTest")
-      excludeTestsMatching("WrapperSuppressReceiveSpansTest")
+      includeTestsMatching("*DeprecatedInterceptorsSuppressReceiveSpansTest")
     }
-    jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
-    systemProperty("otel.instrumentation.messaging.experimental.capture-headers", "Test-Message-Header")
   }
 
   check {
-    dependsOn(testReceiveSpansDisabled)
+    dependsOn(testDeprecated, testDeprecatedSuppressReceiveSpans)
   }
 }
 

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

@@ -26,6 +26,9 @@
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.OpenTelemetryMetricsReporter;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.OpenTelemetrySupplier;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.TracingList;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.KafkaTelemetrySupplier;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.OpenTelemetryConsumerInterceptor;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.OpenTelemetryProducerInterceptor;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Proxy;
 import java.util.Collections;
@@ -38,12 +41,14 @@
 import java.util.logging.Logger;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.TopicPartition;
@@ -207,6 +212,62 @@ <K, V> ConsumerRecords<K, V> addTracing(
     return Collections.unmodifiableMap(config);
   }
 
+  /**
+   * Returns configuration properties that can be used to enable OpenTelemetry instrumentation via
+   * {@code OpenTelemetryProducerInterceptor}. Add these resulting properties to the configuration
+   * map used to initialize a {@link org.apache.kafka.clients.producer.KafkaProducer}.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * //    KafkaTelemetry telemetry = KafkaTelemetry.create(openTelemetry);
+   * //    Map<String, Object> config = new HashMap<>();
+   * //    config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, ...);
+   * //    config.putAll(telemetry.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(
+        ProducerConfig.INTERCEPTOR_CLASSES_CONFIG,
+        OpenTelemetryProducerInterceptor.class.getName());
+    config.put(
+        OpenTelemetryProducerInterceptor.CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER,
+        new KafkaTelemetrySupplier(this));
+    return Collections.unmodifiableMap(config);
+  }
+
+  /**
+   * Returns configuration properties that can be used to enable OpenTelemetry instrumentation via
+   * {@code OpenTelemetryConsumerInterceptor}. Add these resulting properties to the configuration
+   * map used to initialize a {@link org.apache.kafka.clients.consumer.KafkaConsumer}.
+   *
+   * <p>Example usage:
+   *
+   * <pre>{@code
+   * //    KafkaTelemetry telemetry = KafkaTelemetry.create(openTelemetry);
+   * //    Map<String, Object> config = new HashMap<>();
+   * //    config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, ...);
+   * //    config.putAll(telemetry.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(
+        ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG,
+        OpenTelemetryConsumerInterceptor.class.getName());
+    config.put(
+        OpenTelemetryConsumerInterceptor.CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER,
+        new KafkaTelemetrySupplier(this));
+    return Collections.unmodifiableMap(config);
+  }
+
   /**
    * Build and inject span into record.
    *

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

@@ -26,7 +26,10 @@
  * A ConsumerInterceptor that adds tracing capability. Add this interceptor's class name or class
  * via ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG property to your Consumer's properties to get it
  * instantiated and used. See more details on ConsumerInterceptor usage in its Javadoc.
+ *
+ * @deprecated Use {@link OpenTelemetryConsumerInterceptor} instead.
  */
+@Deprecated
 public class TracingConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
 
   private static final KafkaTelemetry telemetry =

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

@@ -22,7 +22,10 @@
  * A ProducerInterceptor that adds tracing capability. Add this interceptor's class name or class
  * via ProducerConfig.INTERCEPTOR_CLASSES_CONFIG property to your Producer's properties to get it
  * instantiated and used. See more details on ProducerInterceptor usage in its Javadoc.
+ *
+ * @deprecated Use {@link OpenTelemetryProducerInterceptor} instead.
  */
+@Deprecated
 public class TracingProducerInterceptor<K, V> implements ProducerInterceptor<K, V> {
 
   private static final KafkaTelemetry telemetry =

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

@@ -0,0 +1,38 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.kafkaclients.v2_6.internal;
+
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.KafkaTelemetry;
+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.
+ *
+ * <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<KafkaTelemetry>, Serializable {
+  private static final long serialVersionUID = 1L;
+  private final transient KafkaTelemetry kafkaTelemetry;
+
+  public KafkaTelemetrySupplier(KafkaTelemetry kafkaTelemetry) {
+    Objects.requireNonNull(kafkaTelemetry);
+    this.kafkaTelemetry = kafkaTelemetry;
+  }
+
+  @Override
+  public KafkaTelemetry get() {
+    return kafkaTelemetry;
+  }
+
+  private Object writeReplace() {
+    // serialize this object to null
+    return null;
+  }
+}

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

@@ -0,0 +1,88 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.kafkaclients.v2_6.internal;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.Timer;
+import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaConsumerContext;
+import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.KafkaConsumerContextUtil;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.KafkaTelemetry;
+import java.util.Map;
+import java.util.Objects;
+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;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+
+/**
+ * A ConsumerInterceptor that adds OpenTelemetry instrumentation.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public class OpenTelemetryConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
+
+  public static final String CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER =
+      "opentelemetry.kafka-telemetry.supplier";
+
+  @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);
+    if (receiveContext == null) {
+      receiveContext = Context.current();
+    }
+    KafkaConsumerContext consumerContext =
+        KafkaConsumerContextUtil.create(receiveContext, consumerGroup, clientId);
+    return telemetry.addTracing(records, consumerContext);
+  }
+
+  @Override
+  public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {}
+
+  @Override
+  public void close() {}
+
+  @Override
+  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);
+
+    Object telemetrySupplier = configs.get(CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER);
+    if (telemetrySupplier == null) {
+      return;
+    }
+
+    KafkaTelemetrySupplier supplier =
+        getProperty(configs, CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER, KafkaTelemetrySupplier.class);
+    this.telemetry = supplier.get();
+  }
+
+  @SuppressWarnings("unchecked")
+  private static <T> T getProperty(Map<String, ?> configs, String key, Class<T> requiredType) {
+    Object value = configs.get(key);
+    if (value == null) {
+      throw new IllegalStateException("Missing required configuration property: " + key);
+    }
+    if (!requiredType.isInstance(value)) {
+      throw new IllegalStateException(
+          "Configuration property " + key + " is not instance of " + requiredType.getSimpleName());
+    }
+    return (T) value;
+  }
+}

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

@@ -0,0 +1,73 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.kafkaclients.v2_6.internal;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.KafkaTelemetry;
+import java.util.Map;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerInterceptor;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+
+/**
+ * A ProducerInterceptor that adds OpenTelemetry instrumentation.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public class OpenTelemetryProducerInterceptor<K, V> implements ProducerInterceptor<K, V> {
+
+  public static final String CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER =
+      "opentelemetry.kafka-telemetry.supplier";
+
+  @Nullable private KafkaTelemetry telemetry;
+  @Nullable private String clientId;
+
+  @Override
+  @CanIgnoreReturnValue
+  public ProducerRecord<K, V> onSend(ProducerRecord<K, V> producerRecord) {
+    if (telemetry != null) {
+      telemetry.buildAndInjectSpan(producerRecord, clientId);
+    }
+    return producerRecord;
+  }
+
+  @Override
+  public void onAcknowledgement(RecordMetadata recordMetadata, Exception e) {}
+
+  @Override
+  public void close() {}
+
+  @Override
+  public void configure(Map<String, ?> configs) {
+    clientId = Objects.toString(configs.get(ProducerConfig.CLIENT_ID_CONFIG), null);
+
+    Object telemetrySupplier = configs.get(CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER);
+    if (telemetrySupplier == null) {
+      return;
+    }
+
+    KafkaTelemetrySupplier supplier =
+        getProperty(configs, CONFIG_KEY_KAFKA_TELEMETRY_SUPPLIER, KafkaTelemetrySupplier.class);
+    this.telemetry = supplier.get();
+  }
+
+  @SuppressWarnings("unchecked")
+  private static <T> T getProperty(Map<String, ?> configs, String key, Class<T> requiredType) {
+    Object value = configs.get(key);
+    if (value == null) {
+      throw new IllegalStateException("Missing required configuration property: " + key);
+    }
+    if (!requiredType.isInstance(value)) {
+      throw new IllegalStateException(
+          "Configuration property " + key + " is not instance of " + requiredType.getSimpleName());
+    }
+    return (T) value;
+  }
+}