Add configurable OpenTelemetry kafka-clients interceptors (#14929)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: trask <218610+trask@users.noreply.github.com>
Co-authored-by: otelbot <197425009+otelbot@users.noreply.github.com>

Author: 3 people

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,32 @@ 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.*")
+    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

@@ -5,61 +5,47 @@
 
 package io.opentelemetry.instrumentation.kafkaclients.v2_6;
 
-import static java.util.logging.Level.WARNING;
-
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
-import io.opentelemetry.context.propagation.TextMapPropagator;
-import io.opentelemetry.context.propagation.TextMapSetter;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
-import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
 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.common.v0_11.internal.KafkaHeadersSetter;
 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.KafkaUtil;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.MetricsReporterList;
 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.KafkaConsumerTelemetry;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.KafkaConsumerTelemetrySupplier;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.KafkaProducerTelemetry;
+import io.opentelemetry.instrumentation.kafkaclients.v2_6.internal.KafkaProducerTelemetrySupplier;
+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;
 import java.util.HashMap;
-import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Map;
-import java.util.concurrent.Future;
-import java.util.function.BiFunction;
-import java.util.logging.Logger;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
 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;
-import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.metrics.MetricsReporter;
 
 public final class KafkaTelemetry {
-  private static final Logger logger = Logger.getLogger(KafkaTelemetry.class.getName());
-
-  private static final TextMapSetter<Headers> SETTER = KafkaHeadersSetter.INSTANCE;
 
   private final OpenTelemetry openTelemetry;
-  private final Instrumenter<KafkaProducerRequest, RecordMetadata> producerInstrumenter;
-  private final Instrumenter<KafkaReceiveRequest, Void> consumerReceiveInstrumenter;
-  private final Instrumenter<KafkaProcessRequest, Void> consumerProcessInstrumenter;
-  private final boolean producerPropagationEnabled;
+  private final KafkaProducerTelemetry producerTelemetry;
+  private final KafkaConsumerTelemetry consumerTelemetry;
 
   KafkaTelemetry(
       OpenTelemetry openTelemetry,
@@ -68,10 +54,13 @@ public final class KafkaTelemetry {
       Instrumenter<KafkaProcessRequest, Void> consumerProcessInstrumenter,
       boolean producerPropagationEnabled) {
     this.openTelemetry = openTelemetry;
-    this.producerInstrumenter = producerInstrumenter;
-    this.consumerReceiveInstrumenter = consumerReceiveInstrumenter;
-    this.consumerProcessInstrumenter = consumerProcessInstrumenter;
-    this.producerPropagationEnabled = producerPropagationEnabled;
+    this.producerTelemetry =
+        new KafkaProducerTelemetry(
+            openTelemetry.getPropagators().getTextMapPropagator(),
+            producerInstrumenter,
+            producerPropagationEnabled);
+    this.consumerTelemetry =
+        new KafkaConsumerTelemetry(consumerReceiveInstrumenter, consumerProcessInstrumenter);
   }
 
   /** Returns a new {@link KafkaTelemetry} configured with the given {@link OpenTelemetry}. */
@@ -86,8 +75,14 @@ public static KafkaTelemetryBuilder builder(OpenTelemetry openTelemetry) {
     return new KafkaTelemetryBuilder(openTelemetry);
   }
 
-  private TextMapPropagator propagator() {
-    return openTelemetry.getPropagators().getTextMapPropagator();
+  // this method can be removed when the deprecated TracingProducerInterceptor is removed
+  KafkaProducerTelemetry getProducerTelemetry() {
+    return producerTelemetry;
+  }
+
+  // this method can be removed when the deprecated TracingProducerInterceptor is removed
+  KafkaConsumerTelemetry getConsumerTelemetry() {
+    return consumerTelemetry;
   }
 
   /** Returns a decorated {@link Producer} that emits spans for each sent message. */
@@ -109,7 +104,8 @@ public <K, V> Producer<K, V> wrap(Producer<K, V> producer) {
                             && method.getParameterTypes()[1] == Callback.class
                         ? (Callback) args[1]
                         : null;
-                return buildAndInjectSpan(record, producer, callback, producer::send);
+                return producerTelemetry.buildAndInjectSpan(
+                    record, producer, callback, producer::send);
               }
               try {
                 return method.invoke(producer, args);
@@ -138,35 +134,19 @@ public <K, V> Consumer<K, V> wrap(Consumer<K, V> consumer) {
               // ConsumerRecords<K, V> poll(Duration duration)
               if ("poll".equals(method.getName()) && result instanceof ConsumerRecords) {
                 ConsumerRecords<K, V> consumerRecords = (ConsumerRecords<K, V>) result;
-                Context receiveContext = buildAndFinishSpan(consumerRecords, consumer, timer);
+                Context receiveContext =
+                    consumerTelemetry.buildAndFinishSpan(consumerRecords, consumer, timer);
                 if (receiveContext == null) {
                   receiveContext = Context.current();
                 }
                 KafkaConsumerContext consumerContext =
                     KafkaConsumerContextUtil.create(receiveContext, consumer);
-                result = addTracing(consumerRecords, consumerContext);
+                result = consumerTelemetry.addTracing(consumerRecords, consumerContext);
               }
               return result;
             });
   }
 
-  <K, V> ConsumerRecords<K, V> addTracing(
-      ConsumerRecords<K, V> consumerRecords, KafkaConsumerContext consumerContext) {
-    if (consumerRecords.isEmpty()) {
-      return consumerRecords;
-    }
-
-    Map<TopicPartition, List<ConsumerRecord<K, V>>> records = new LinkedHashMap<>();
-    for (TopicPartition partition : consumerRecords.partitions()) {
-      List<ConsumerRecord<K, V>> list = consumerRecords.records(partition);
-      if (list != null && !list.isEmpty()) {
-        list = TracingList.wrap(list, consumerProcessInstrumenter, () -> true, consumerContext);
-      }
-      records.put(partition, list);
-    }
-    return new ConsumerRecords<>(records);
-  }
-
   /**
    * Produces a set of kafka client config properties (consumer or producer) to register a {@link
    * MetricsReporter} that records metrics to an {@code openTelemetry} instance. Add these resulting
@@ -208,113 +188,58 @@ <K, V> ConsumerRecords<K, V> addTracing(
   }
 
   /**
-   * Build and inject span into record.
+   * 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>
    *
-   * @param record the producer record to inject span info.
+   * @return the kafka producer interceptor config properties
    */
-  <K, V> void buildAndInjectSpan(ProducerRecord<K, V> record, String clientId) {
-    Context parentContext = Context.current();
-
-    KafkaProducerRequest request = KafkaProducerRequest.create(record, clientId);
-    if (!producerInstrumenter.shouldStart(parentContext, request)) {
-      return;
-    }
-
-    Context context = producerInstrumenter.start(parentContext, request);
-    if (producerPropagationEnabled) {
-      try {
-        propagator().inject(context, record.headers(), SETTER);
-      } catch (Throwable t) {
-        // it can happen if headers are read only (when record is sent second time)
-        logger.log(WARNING, "failed to inject span context. sending record second time?", t);
-      }
-    }
-    producerInstrumenter.end(context, request, null, null);
+  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_PRODUCER_TELEMETRY_SUPPLIER,
+        new KafkaProducerTelemetrySupplier(producerTelemetry));
+    return Collections.unmodifiableMap(config);
   }
 
   /**
-   * Build and inject span into record.
+   * 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>
    *
-   * @param record the producer record to inject span info.
-   * @param callback the producer send callback
-   * @return send function's result
+   * @return the kafka consumer interceptor config properties
    */
-  <K, V> Future<RecordMetadata> buildAndInjectSpan(
-      ProducerRecord<K, V> record,
-      Producer<K, V> producer,
-      Callback callback,
-      BiFunction<ProducerRecord<K, V>, Callback, Future<RecordMetadata>> sendFn) {
-    Context parentContext = Context.current();
-
-    KafkaProducerRequest request = KafkaProducerRequest.create(record, producer);
-    if (!producerInstrumenter.shouldStart(parentContext, request)) {
-      return sendFn.apply(record, callback);
-    }
-
-    Context context = producerInstrumenter.start(parentContext, request);
-    propagator().inject(context, record.headers(), SETTER);
-
-    try (Scope ignored = context.makeCurrent()) {
-      return sendFn.apply(record, new ProducerCallback(callback, parentContext, context, request));
-    }
-  }
-
-  private <K, V> Context buildAndFinishSpan(
-      ConsumerRecords<K, V> records, Consumer<K, V> consumer, Timer timer) {
-    return buildAndFinishSpan(
-        records, KafkaUtil.getConsumerGroup(consumer), KafkaUtil.getClientId(consumer), timer);
-  }
-
-  <K, V> Context buildAndFinishSpan(
-      ConsumerRecords<K, V> records, String consumerGroup, String clientId, Timer timer) {
-    if (records.isEmpty()) {
-      return null;
-    }
-    Context parentContext = Context.current();
-    KafkaReceiveRequest request = KafkaReceiveRequest.create(records, consumerGroup, clientId);
-    Context context = null;
-    if (consumerReceiveInstrumenter.shouldStart(parentContext, request)) {
-      context =
-          InstrumenterUtil.startAndEnd(
-              consumerReceiveInstrumenter,
-              parentContext,
-              request,
-              null,
-              null,
-              timer.startTime(),
-              timer.now());
-    }
-
-    // we're returning the context of the receive span so that process spans can use it as
-    // parent context even though the span has ended
-    // this is the suggested behavior according to the spec batch receive scenario:
-    // https://github.com/open-telemetry/semantic-conventions/blob/main/docs/messaging/messaging-spans.md#batch-receiving
-    return context;
-  }
-
-  private class ProducerCallback implements Callback {
-    private final Callback callback;
-    private final Context parentContext;
-    private final Context context;
-    private final KafkaProducerRequest request;
-
-    ProducerCallback(
-        Callback callback, Context parentContext, Context context, KafkaProducerRequest request) {
-      this.callback = callback;
-      this.parentContext = parentContext;
-      this.context = context;
-      this.request = request;
-    }
-
-    @Override
-    public void onCompletion(RecordMetadata metadata, Exception exception) {
-      producerInstrumenter.end(context, request, metadata, exception);
-
-      if (callback != null) {
-        try (Scope ignored = parentContext.makeCurrent()) {
-          callback.onCompletion(metadata, exception);
-        }
-      }
-    }
+  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_CONSUMER_TELEMETRY_SUPPLIER,
+        new KafkaConsumerTelemetrySupplier(consumerTelemetry));
+    return Collections.unmodifiableMap(config);
   }
 }