Make kafka library and javaagent instrumentations more similar (#9738)

Author: laurit

instrumentation/kafka/kafka-clients/kafka-clients-0.11/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/kafka/KafkaClientsConsumerProcessTracing.java

@@ -5,6 +5,8 @@
 
 package io.opentelemetry.javaagent.bootstrap.kafka;
 
+import java.util.function.BooleanSupplier;
+
 // Classes used by multiple instrumentations should be in a bootstrap module to ensure that all
 // instrumentations see the same class. Helper classes are injected into each class loader that
 // contains an instrumentation that uses them, so instrumentations in different class loaders will
@@ -23,4 +25,8 @@ public static boolean setEnabled(boolean enabled) {
   public static boolean wrappingEnabled() {
     return wrappingEnabled.get();
   }
+
+  public static BooleanSupplier wrappingEnabledSupplier() {
+    return KafkaClientsConsumerProcessTracing::wrappingEnabled;
+  }
 }

instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/ConsumerRecordsInstrumentation.java

@@ -5,6 +5,8 @@
 
 package io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11;
 
+import static io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing.wrappingEnabledSupplier;
+import static io.opentelemetry.javaagent.instrumentation.kafkaclients.v0_11.KafkaSingletons.consumerProcessInstrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -14,6 +16,9 @@
 
 import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
+import io.opentelemetry.instrumentation.kafka.internal.TracingIterable;
+import io.opentelemetry.instrumentation.kafka.internal.TracingIterator;
+import io.opentelemetry.instrumentation.kafka.internal.TracingList;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.util.Iterator;
@@ -70,7 +75,9 @@ public static <K, V> void wrap(
       // case it's important to overwrite the leaked span instead of suppressing the correct span
       // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
       KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
-      iterable = TracingIterable.wrap(iterable, consumerContext);
+      iterable =
+          TracingIterable.wrap(
+              iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
@@ -88,7 +95,9 @@ public static <K, V> void wrap(
       // case it's important to overwrite the leaked span instead of suppressing the correct span
       // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
       KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
-      list = TracingList.wrap(list, consumerContext);
+      list =
+          TracingList.wrap(
+              list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
@@ -106,7 +115,9 @@ public static <K, V> void wrap(
       // case it's important to overwrite the leaked span instead of suppressing the correct span
       // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
       KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
-      iterator = TracingIterator.wrap(iterator, consumerContext);
+      iterator =
+          TracingIterator.wrap(
+              iterator, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 }

instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/TracingIterable.java

@@ -20,4 +20,24 @@ tasks {
   withType<Test>().configureEach {
     usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
   }
+
+  val testReceiveSpansDisabled by registering(Test::class) {
+    filter {
+      includeTestsMatching("InterceptorsSuppressReceiveSpansTest")
+      includeTestsMatching("WrapperSuppressReceiveSpansTest")
+    }
+    include("**/InterceptorsSuppressReceiveSpansTest.*", "**/WrapperSuppressReceiveSpansTest.*")
+  }
+
+  test {
+    filter {
+      excludeTestsMatching("InterceptorsSuppressReceiveSpansTest")
+      excludeTestsMatching("WrapperSuppressReceiveSpansTest")
+    }
+    jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
+  }
+
+  check {
+    dependsOn(testReceiveSpansDisabled)
+  }
 }

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

@@ -13,16 +13,24 @@
 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.kafka.internal.KafkaConsumerContext;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaHeadersSetter;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaProducerRequest;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaUtil;
 import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetryMetricsReporter;
 import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetrySupplier;
+import io.opentelemetry.instrumentation.kafka.internal.TracingList;
 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;
@@ -37,6 +45,7 @@
 import org.apache.kafka.clients.producer.Producer;
 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;
 
@@ -47,16 +56,19 @@ public final class KafkaTelemetry {
 
   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;
 
   KafkaTelemetry(
       OpenTelemetry openTelemetry,
       Instrumenter<KafkaProducerRequest, RecordMetadata> producerInstrumenter,
+      Instrumenter<KafkaReceiveRequest, Void> consumerReceiveInstrumenter,
       Instrumenter<KafkaProcessRequest, Void> consumerProcessInstrumenter,
       boolean producerPropagationEnabled) {
     this.openTelemetry = openTelemetry;
     this.producerInstrumenter = producerInstrumenter;
+    this.consumerReceiveInstrumenter = consumerReceiveInstrumenter;
     this.consumerProcessInstrumenter = consumerProcessInstrumenter;
     this.producerPropagationEnabled = producerPropagationEnabled;
   }
@@ -115,6 +127,7 @@ public <K, V> Consumer<K, V> wrap(Consumer<K, V> consumer) {
             new Class<?>[] {Consumer.class},
             (proxy, method, args) -> {
               Object result;
+              Timer timer = "poll".equals(method.getName()) ? Timer.start() : null;
               try {
                 result = method.invoke(consumer, args);
               } catch (InvocationTargetException exception) {
@@ -123,12 +136,36 @@ public <K, V> Consumer<K, V> wrap(Consumer<K, V> consumer) {
               // ConsumerRecords<K, V> poll(long timeout)
               // ConsumerRecords<K, V> poll(Duration duration)
               if ("poll".equals(method.getName()) && result instanceof ConsumerRecords) {
-                buildAndFinishSpan((ConsumerRecords) result, consumer);
+                ConsumerRecords<K, V> consumerRecords = (ConsumerRecords<K, V>) result;
+                Context receiveContext = buildAndFinishSpan(consumerRecords, consumer, timer);
+                if (receiveContext == null) {
+                  receiveContext = Context.current();
+                }
+                KafkaConsumerContext consumerContext =
+                    KafkaConsumerContextUtil.create(receiveContext, consumer);
+                result = 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
@@ -221,23 +258,37 @@ <K, V> Future<RecordMetadata> buildAndInjectSpan(
     }
   }
 
-  private <K, V> void buildAndFinishSpan(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
-    buildAndFinishSpan(
-        records, KafkaUtil.getConsumerGroup(consumer), KafkaUtil.getClientId(consumer));
+  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> void buildAndFinishSpan(
-      ConsumerRecords<K, V> records, String consumerGroup, String clientId) {
+  <K, V> Context buildAndFinishSpan(
+      ConsumerRecords<K, V> records, String consumerGroup, String clientId, Timer timer) {
+    if (records.isEmpty()) {
+      return null;
+    }
     Context parentContext = Context.current();
-    for (ConsumerRecord<K, V> record : records) {
-      KafkaProcessRequest request = KafkaProcessRequest.create(record, consumerGroup, clientId);
-      if (!consumerProcessInstrumenter.shouldStart(parentContext, request)) {
-        continue;
-      }
-
-      Context context = consumerProcessInstrumenter.start(parentContext, request);
-      consumerProcessInstrumenter.end(context, request, null, null);
+    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 {

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

@@ -10,10 +10,10 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
-import io.opentelemetry.instrumentation.api.instrumenter.messaging.MessageOperation;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaProcessRequest;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaProducerRequest;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
@@ -25,8 +25,10 @@ public final class KafkaTelemetryBuilder {
   private final OpenTelemetry openTelemetry;
   private final List<AttributesExtractor<KafkaProducerRequest, RecordMetadata>>
       producerAttributesExtractors = new ArrayList<>();
-  private final List<AttributesExtractor<KafkaProcessRequest, Void>> consumerAttributesExtractors =
-      new ArrayList<>();
+  private final List<AttributesExtractor<KafkaProcessRequest, Void>>
+      consumerProcessAttributesExtractors = new ArrayList<>();
+  private final List<AttributesExtractor<KafkaReceiveRequest, Void>>
+      consumerReceiveAttributesExtractors = new ArrayList<>();
   private List<String> capturedHeaders = emptyList();
   private boolean captureExperimentalSpanAttributes = false;
   private boolean propagationEnabled = true;
@@ -43,10 +45,25 @@ public KafkaTelemetryBuilder addProducerAttributesExtractors(
     return this;
   }
 
+  /** Use {@link #addConsumerProcessAttributesExtractors(AttributesExtractor)} instead. */
+  @Deprecated
   @CanIgnoreReturnValue
   public KafkaTelemetryBuilder addConsumerAttributesExtractors(
       AttributesExtractor<KafkaProcessRequest, Void> extractor) {
-    consumerAttributesExtractors.add(extractor);
+    return addConsumerProcessAttributesExtractors(extractor);
+  }
+
+  @CanIgnoreReturnValue
+  public KafkaTelemetryBuilder addConsumerProcessAttributesExtractors(
+      AttributesExtractor<KafkaProcessRequest, Void> extractor) {
+    consumerProcessAttributesExtractors.add(extractor);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public KafkaTelemetryBuilder addConsumerReceiveAttributesExtractors(
+      AttributesExtractor<KafkaReceiveRequest, Void> extractor) {
+    consumerReceiveAttributesExtractors.add(extractor);
     return this;
   }
 
@@ -109,8 +126,8 @@ public KafkaTelemetry build() {
     return new KafkaTelemetry(
         openTelemetry,
         instrumenterFactory.createProducerInstrumenter(producerAttributesExtractors),
-        instrumenterFactory.createConsumerOperationInstrumenter(
-            MessageOperation.RECEIVE, consumerAttributesExtractors),
+        instrumenterFactory.createConsumerReceiveInstrumenter(consumerReceiveAttributesExtractors),
+        instrumenterFactory.createConsumerProcessInstrumenter(consumerProcessAttributesExtractors),
         propagationEnabled);
   }
 }

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

@@ -7,6 +7,11 @@
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.internal.ConfigPropertiesUtil;
+import io.opentelemetry.instrumentation.api.internal.Timer;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
 import java.util.Map;
 import java.util.Objects;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -22,16 +27,28 @@
  */
 public class TracingConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
 
-  private static final KafkaTelemetry telemetry = KafkaTelemetry.create(GlobalOpenTelemetry.get());
+  private static final KafkaTelemetry telemetry =
+      KafkaTelemetry.builder(GlobalOpenTelemetry.get())
+          .setMessagingReceiveInstrumentationEnabled(
+              ConfigPropertiesUtil.getBoolean(
+                  "otel.instrumentation.messaging.experimental.receive-telemetry.enabled", false))
+          .build();
 
   private String consumerGroup;
   private String clientId;
 
   @Override
   @CanIgnoreReturnValue
   public ConsumerRecords<K, V> onConsume(ConsumerRecords<K, V> records) {
-    telemetry.buildAndFinishSpan(records, consumerGroup, clientId);
-    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