Implement messaging.kafka.* attributes spec (part 1) (#7824)

Part 1 of #7771

It's a lot more than I initially expected it to be; I'll introduce the
clientId attribute in part 2

Author: Mateusz Rzeszutek

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

@@ -21,6 +21,7 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 
@@ -59,6 +60,7 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class IterableAdvice {
 
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
@@ -69,13 +71,16 @@ 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)
       Context receiveContext = VirtualField.find(ConsumerRecords.class, Context.class).get(records);
-      iterable = TracingIterable.wrap(iterable, receiveContext);
+      Consumer<K, V> consumer =
+          VirtualField.find(ConsumerRecords.class, Consumer.class).get(records);
+      iterable = TracingIterable.wrap(iterable, receiveContext, consumer);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ListAdvice {
 
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
@@ -86,13 +91,16 @@ 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)
       Context receiveContext = VirtualField.find(ConsumerRecords.class, Context.class).get(records);
-      list = TracingList.wrap(list, receiveContext);
+      Consumer<K, V> consumer =
+          VirtualField.find(ConsumerRecords.class, Consumer.class).get(records);
+      list = TracingList.wrap(list, receiveContext, consumer);
     }
   }
 
   @SuppressWarnings("unused")
   public static class IteratorAdvice {
 
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
@@ -103,7 +111,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)
       Context receiveContext = VirtualField.find(ConsumerRecords.class, Context.class).get(records);
-      iterator = TracingIterator.wrap(iterator, receiveContext);
+      Consumer<K, V> consumer =
+          VirtualField.find(ConsumerRecords.class, Consumer.class).get(records);
+      iterator = TracingIterator.wrap(iterator, receiveContext, consumer);
     }
   }
 }

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

@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.internal.InstrumenterUtil;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.instrumentation.kafka.internal.ConsumerAndRecord;
 import io.opentelemetry.instrumentation.kafka.internal.Timer;
 import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
@@ -28,6 +29,7 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 
@@ -83,6 +85,7 @@ public static Timer onEnter() {
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void onExit(
         @Advice.Enter Timer timer,
+        @Advice.This Consumer<?, ?> consumer,
         @Advice.Return ConsumerRecords<?, ?> records,
         @Advice.Thrown Throwable error) {
 
@@ -91,24 +94,32 @@ public static void onExit(
         return;
       }
 
+      // we're attaching the consumer to the records to be able to retrieve things like consumer
+      // group or clientId later
+      VirtualField<ConsumerRecords<?, ?>, Consumer<?, ?>> consumerRecordsConsumer =
+          VirtualField.find(ConsumerRecords.class, Consumer.class);
+      consumerRecordsConsumer.set(records, consumer);
+
       Context parentContext = currentContext();
-      if (consumerReceiveInstrumenter().shouldStart(parentContext, records)) {
+      ConsumerAndRecord<ConsumerRecords<?, ?>> request =
+          ConsumerAndRecord.create(consumer, records);
+
+      if (consumerReceiveInstrumenter().shouldStart(parentContext, request)) {
         // disable process tracing and store the receive span for each individual record too
         boolean previousValue = KafkaClientsConsumerProcessTracing.setEnabled(false);
         try {
           Context context =
               InstrumenterUtil.startAndEnd(
                   consumerReceiveInstrumenter(),
                   parentContext,
-                  records,
+                  request,
                   null,
                   error,
                   timer.startTime(),
                   timer.now());
 
           // we're storing the context of the receive span so that process spans can use it as
-          // parent
-          // context even though the span has ended
+          // 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/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/messaging.md#batch-receiving
           VirtualField<ConsumerRecords<?, ?>, Context> consumerRecordsContext =

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

@@ -7,6 +7,7 @@
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.kafka.internal.ConsumerAndRecord;
 import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
 import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetryMetricsReporter;
 import io.opentelemetry.instrumentation.kafka.internal.OpenTelemetrySupplier;
@@ -34,8 +35,10 @@ public final class KafkaSingletons {
           .getBoolean("otel.instrumentation.kafka.metric-reporter.enabled", true);
 
   private static final Instrumenter<ProducerRecord<?, ?>, RecordMetadata> PRODUCER_INSTRUMENTER;
-  private static final Instrumenter<ConsumerRecords<?, ?>, Void> CONSUMER_RECEIVE_INSTRUMENTER;
-  private static final Instrumenter<ConsumerRecord<?, ?>, Void> CONSUMER_PROCESS_INSTRUMENTER;
+  private static final Instrumenter<ConsumerAndRecord<ConsumerRecords<?, ?>>, Void>
+      CONSUMER_RECEIVE_INSTRUMENTER;
+  private static final Instrumenter<ConsumerAndRecord<ConsumerRecord<?, ?>>, Void>
+      CONSUMER_PROCESS_INSTRUMENTER;
 
   static {
     KafkaInstrumenterFactory instrumenterFactory =
@@ -59,11 +62,13 @@ public static boolean isProducerPropagationEnabled() {
     return PRODUCER_INSTRUMENTER;
   }
 
-  public static Instrumenter<ConsumerRecords<?, ?>, Void> consumerReceiveInstrumenter() {
+  public static Instrumenter<ConsumerAndRecord<ConsumerRecords<?, ?>>, Void>
+      consumerReceiveInstrumenter() {
     return CONSUMER_RECEIVE_INSTRUMENTER;
   }
 
-  public static Instrumenter<ConsumerRecord<?, ?>, Void> consumerProcessInstrumenter() {
+  public static Instrumenter<ConsumerAndRecord<ConsumerRecord<?, ?>>, Void>
+      consumerProcessInstrumenter() {
     return CONSUMER_PROCESS_INSTRUMENTER;
   }
 

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

@@ -9,23 +9,30 @@
 import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
 import java.util.Iterator;
 import javax.annotation.Nullable;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 public class TracingIterable<K, V> implements Iterable<ConsumerRecord<K, V>> {
   private final Iterable<ConsumerRecord<K, V>> delegate;
   @Nullable private final Context receiveContext;
+  private final Consumer<K, V> consumer;
   private boolean firstIterator = true;
 
   protected TracingIterable(
-      Iterable<ConsumerRecord<K, V>> delegate, @Nullable Context receiveContext) {
+      Iterable<ConsumerRecord<K, V>> delegate,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
     this.delegate = delegate;
     this.receiveContext = receiveContext;
+    this.consumer = consumer;
   }
 
   public static <K, V> Iterable<ConsumerRecord<K, V>> wrap(
-      Iterable<ConsumerRecord<K, V>> delegate, @Nullable Context receiveContext) {
+      Iterable<ConsumerRecord<K, V>> delegate,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
     if (KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
-      return new TracingIterable<>(delegate, receiveContext);
+      return new TracingIterable<>(delegate, receiveContext, consumer);
     }
     return delegate;
   }
@@ -37,7 +44,7 @@ public Iterator<ConsumerRecord<K, V>> iterator() {
     // However, this is not thread-safe, but usually the first (hopefully only) traversal of
     // ConsumerRecords is performed in the same thread that called poll()
     if (firstIterator) {
-      it = TracingIterator.wrap(delegate.iterator(), receiveContext);
+      it = TracingIterator.wrap(delegate.iterator(), receiveContext, consumer);
       firstIterator = false;
     } else {
       it = delegate.iterator();

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

@@ -9,35 +9,44 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.kafka.internal.ConsumerAndRecord;
 import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
 import java.util.Iterator;
 import javax.annotation.Nullable;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 public class TracingIterator<K, V> implements Iterator<ConsumerRecord<K, V>> {
+
   private final Iterator<ConsumerRecord<K, V>> delegateIterator;
   private final Context parentContext;
+  private final Consumer<K, V> consumer;
 
   /*
    * Note: this may potentially create problems if this iterator is used from different threads. But
    * at the moment we cannot do much about this.
    */
-  @Nullable private ConsumerRecord<?, ?> currentRequest;
+  @Nullable private ConsumerAndRecord<ConsumerRecord<?, ?>> currentRequest;
   @Nullable private Context currentContext;
   @Nullable private Scope currentScope;
 
   private TracingIterator(
-      Iterator<ConsumerRecord<K, V>> delegateIterator, @Nullable Context receiveContext) {
+      Iterator<ConsumerRecord<K, V>> delegateIterator,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
     this.delegateIterator = delegateIterator;
 
     // use the receive CONSUMER as parent if it's available
     this.parentContext = receiveContext != null ? receiveContext : Context.current();
+    this.consumer = consumer;
   }
 
   public static <K, V> Iterator<ConsumerRecord<K, V>> wrap(
-      Iterator<ConsumerRecord<K, V>> delegateIterator, @Nullable Context receiveContext) {
+      Iterator<ConsumerRecord<K, V>> delegateIterator,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
     if (KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
-      return new TracingIterator<>(delegateIterator, receiveContext);
+      return new TracingIterator<>(delegateIterator, receiveContext, consumer);
     }
     return delegateIterator;
   }
@@ -60,7 +69,7 @@ public ConsumerRecord<K, V> next() {
     // (https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1947)
     ConsumerRecord<K, V> next = delegateIterator.next();
     if (next != null && KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
-      currentRequest = next;
+      currentRequest = ConsumerAndRecord.create(consumer, next);
       currentContext = consumerProcessInstrumenter().start(parentContext, currentRequest);
       currentScope = currentContext.makeCurrent();
     }

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

@@ -11,20 +11,26 @@
 import java.util.List;
 import java.util.ListIterator;
 import javax.annotation.Nullable;
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 public class TracingList<K, V> extends TracingIterable<K, V> implements List<ConsumerRecord<K, V>> {
   private final List<ConsumerRecord<K, V>> delegate;
 
-  private TracingList(List<ConsumerRecord<K, V>> delegate, @Nullable Context receiveContext) {
-    super(delegate, receiveContext);
+  private TracingList(
+      List<ConsumerRecord<K, V>> delegate,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
+    super(delegate, receiveContext, consumer);
     this.delegate = delegate;
   }
 
   public static <K, V> List<ConsumerRecord<K, V>> wrap(
-      List<ConsumerRecord<K, V>> delegate, @Nullable Context receiveContext) {
+      List<ConsumerRecord<K, V>> delegate,
+      @Nullable Context receiveContext,
+      Consumer<K, V> consumer) {
     if (KafkaClientsConsumerProcessTracing.wrappingEnabled()) {
-      return new TracingList<>(delegate, receiveContext);
+      return new TracingList<>(delegate, receiveContext, consumer);
     }
     return delegate;
   }