fix instrumentation

Author: laurit

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

@@ -26,6 +26,7 @@
 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;
@@ -73,83 +74,76 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class IterableAdvice {
 
-    @SuppressWarnings("unchecked")
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> void wrap(
+    public static <K, V> Iterable<ConsumerRecord<K, V>> wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return(readOnly = false) Iterable<ConsumerRecord<K, V>> iterable) {
+        @Advice.Return 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);
-      iterable =
-          TracingIterable.wrap(
-              iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+      return TracingIterable.wrap(
+          iterable, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ListAdvice {
 
-    @SuppressWarnings("unchecked")
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> void wrap(
+    public static <K, V> List<ConsumerRecord<K, V>> wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return(readOnly = false) List<ConsumerRecord<K, V>> list) {
+        @Advice.Return 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);
-      list =
-          TracingList.wrap(
-              list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
+      return TracingList.wrap(
+          list, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
   @SuppressWarnings("unused")
   public static class IteratorAdvice {
 
-    @SuppressWarnings("unchecked")
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> void wrap(
+    public static <K, V> Iterator<ConsumerRecord<K, V>> wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return(readOnly = false) Iterator<ConsumerRecord<K, V>> iterator) {
+        @Advice.Return 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);
+      return TracingIterator.wrap(
+          iterator, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ListIteratorAdvice {
 
-    @SuppressWarnings("unchecked")
+    @AssignReturned.ToReturned
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> void wrap(
+    public static <K, V> ListIterator<ConsumerRecord<K, V>> wrap(
         @Advice.This ConsumerRecords<?, ?> records,
-        @Advice.Return(readOnly = false) ListIterator<ConsumerRecord<K, V>> listIterator) {
+        @Advice.Return 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);
-      listIterator =
-          TracingListIterator.wrap(
-              listIterator,
-              consumerProcessInstrumenter(),
-              wrappingEnabledSupplier(),
-              consumerContext);
+      return TracingListIterator.wrap(
+          listIterator, consumerProcessInstrumenter(), wrappingEnabledSupplier(), consumerContext);
     }
   }
 }

instrumentation/kafka/kafka-clients/kafka-clients-0.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/kafkaclients/v0_11/KafkaClientDefaultTest.java

@@ -172,6 +172,7 @@ void testRecordsWithTopicPartitionKafkaConsume(boolean testListIterator)
         consumerRecords.records(KafkaClientBaseTest.topicPartition);
     assertThat(recordsInPartition.size()).isEqualTo(1);
 
+    // iterate over records to generate spans
     if (testListIterator) {
       for (ListIterator<? extends ConsumerRecord<?, ?>> iterator =
               recordsInPartition.listIterator();
@@ -181,7 +182,6 @@ void testRecordsWithTopicPartitionKafkaConsume(boolean testListIterator)
         assertThat(record.key()).isNull();
       }
     } else {
-      // iterate over records to generate spans
       for (ConsumerRecord<?, ?> record : recordsInPartition) {
         assertThat(record.value()).isEqualTo(greeting);
         assertThat(record.key()).isNull();