Update ConsumerRecordsInstrumentation.java

Author: tryfixotelpr

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

@@ -19,12 +19,13 @@
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.TracingIterable;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.TracingIterator;
 import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.TracingList;
+import io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal.TracingListIterator;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import java.util.Iterator;
 import java.util.List;
+import java.util.ListIterator;
 import net.bytebuddy.asm.Advice;
-import net.bytebuddy.asm.Advice.AssignReturned;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -60,62 +61,95 @@ public void transform(TypeTransformer transformer) {
             .and(takesArguments(0))
             .and(returns(Iterator.class)),
         ConsumerRecordsInstrumentation.class.getName() + "$IteratorAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("listIterator"))
+            .and(takesArguments(0))
+            .and(returns(ListIterator.class)),
+        ConsumerRecordsInstrumentation.class.getName() + "$ListIteratorAdvice");
   }
 
   @SuppressWarnings("unused")
   public static class IterableAdvice {
 
-    @AssignReturned.ToReturned
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> Iterable<ConsumerRecord<K, V>> wrap(
+    public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return Iterable<ConsumerRecord<K, V>> iterable) {
+        @Advice.Return(readOnly = false) Iterable<ConsumerRecord<K, V>> iterable) {
 
       // it's important not to suppress consumer span creation here because this instrumentation can
       // leak the context and so there may be a leaked consumer span in the context, in which
       // 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);
-      return TracingIterable.wrap(
-          iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+      iterable =
+          TracingIterable.wrap(
+              iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ListAdvice {
 
-    @AssignReturned.ToReturned
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> List<ConsumerRecord<K, V>> wrap(
+    public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return List<ConsumerRecord<K, V>> list) {
+        @Advice.Return(readOnly = false) List<ConsumerRecord<K, V>> list) {
 
       // it's important not to suppress consumer span creation here because this instrumentation can
       // leak the context and so there may be a leaked consumer span in the context, in which
       // 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);
-      return TracingList.wrap(
-          list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+      list =
+          TracingList.wrap(
+              list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
   @SuppressWarnings("unused")
   public static class IteratorAdvice {
 
-    @AssignReturned.ToReturned
+    @SuppressWarnings("unchecked")
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static <K, V> void wrap(
+        @Advice.This ConsumerRecords<?, ?> records,
+        @Advice.Return(readOnly = false) Iterator<ConsumerRecord<K, V>> iterator) {
+
+      // it's important not to suppress consumer span creation here because this instrumentation can
+      // leak the context and so there may be a leaked consumer span in the context, in which
+      // 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, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class ListIteratorAdvice {
+
+    @SuppressWarnings("unchecked")
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> Iterator<ConsumerRecord<K, V>> wrap(
+    public static <K, V> void wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return Iterator<ConsumerRecord<K, V>> iterator) {
+        @Advice.Return(readOnly = false) ListIterator<ConsumerRecord<K, V>> listIterator) {
 
       // it's important not to suppress consumer span creation here because this instrumentation can
       // leak the context and so there may be a leaked consumer span in the context, in which
       // 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);
-      return TracingIterator.wrap(
-          iterator, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+      listIterator =
+          TracingListIterator.wrap(
+              listIterator,
+              consumerProcessInstrumenter(),
+              wrappingEnabledSupplier(),
+              consumerContext);
     }
   }
 }