Fix kafka latest dep tests (#13544)

Author: laurit

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

@@ -65,8 +65,7 @@ void testKafkaProducerAndConsumerSpan(boolean testHeaders) throws Exception {
         });
 
     awaitUntilConsumerIsReady();
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     assertThat(records.count()).isEqualTo(1);
 
     // iterate over records to generate spans
@@ -118,8 +117,7 @@ void testPassThroughTombstone()
       throws ExecutionException, InterruptedException, TimeoutException {
     producer.send(new ProducerRecord<>(SHARED_TOPIC, null)).get(5, TimeUnit.SECONDS);
     awaitUntilConsumerIsReady();
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     assertThat(records.count()).isEqualTo(1);
 
     // iterate over records to generate spans
@@ -168,8 +166,7 @@ void testRecordsWithTopicPartitionKafkaConsume()
     testing.waitForTraces(1);
 
     awaitUntilConsumerIsReady();
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> consumerRecords = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> consumerRecords = poll(Duration.ofSeconds(5));
     List<? extends ConsumerRecord<?, ?>> recordsInPartition =
         consumerRecords.records(KafkaClientBaseTest.topicPartition);
     assertThat(recordsInPartition.size()).isEqualTo(1);

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

@@ -41,8 +41,7 @@ void testReadRemoteContextWhenPropagationIsDisabled() throws InterruptedExceptio
 
     awaitUntilConsumerIsReady();
 
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     assertThat(records.count()).isEqualTo(1);
 
     // iterate over records to generate spans

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

@@ -58,8 +58,7 @@ void testKafkaProduceAndConsume() throws InterruptedException {
 
     awaitUntilConsumerIsReady();
     // check that the message was received
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     for (ConsumerRecord<?, ?> record : records) {
       testing.runWithSpan(
           "processing",
@@ -99,8 +98,7 @@ void testPassThroughTombstone()
       throws ExecutionException, InterruptedException, TimeoutException {
     producer.send(new ProducerRecord<>(SHARED_TOPIC, null)).get(5, TimeUnit.SECONDS);
     awaitUntilConsumerIsReady();
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     assertThat(records.count()).isEqualTo(1);
 
     // iterate over records to generate spans
@@ -136,8 +134,7 @@ void testRecordsWithTopicPartitionKafkaConsume()
     testing.waitForTraces(1);
 
     awaitUntilConsumerIsReady();
-    @SuppressWarnings("PreferJavaTimeOverload")
-    ConsumerRecords<?, ?> consumerRecords = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> consumerRecords = poll(Duration.ofSeconds(5));
     List<? extends ConsumerRecord<?, ?>> recordsInPartition =
         consumerRecords.records(KafkaClientBaseTest.topicPartition);
     assertThat(recordsInPartition.size()).isEqualTo(1);

instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/AbstractOpenTelemetryMetricsReporterTest.java

@@ -268,8 +268,8 @@ void observeMetrics() {
                 "kafka.consumer.io_time_ns_avg",
                 "kafka.consumer.io_wait_ratio",
                 "kafka.consumer.io_wait_time_ns_avg",
-                "kafka.consumer.io_waittime_total",
-                "kafka.consumer.iotime_total",
+                // "kafka.consumer.io_waittime_total",
+                // "kafka.consumer.iotime_total",
                 "kafka.consumer.last_poll_seconds_ago",
                 "kafka.consumer.network_io_rate",
                 "kafka.consumer.network_io_total",
@@ -300,7 +300,7 @@ void observeMetrics() {
                 "kafka.producer.buffer_exhausted_total",
                 "kafka.producer.buffer_total_bytes",
                 "kafka.producer.bufferpool_wait_ratio",
-                "kafka.producer.bufferpool_wait_time_total",
+                // "kafka.producer.bufferpool_wait_time_total",
                 "kafka.producer.compression_rate_avg",
                 "kafka.producer.connection_close_rate",
                 "kafka.producer.connection_close_total",
@@ -317,8 +317,8 @@ void observeMetrics() {
                 "kafka.producer.io_time_ns_avg",
                 "kafka.producer.io_wait_ratio",
                 "kafka.producer.io_wait_time_ns_avg",
-                "kafka.producer.io_waittime_total",
-                "kafka.producer.iotime_total",
+                // "kafka.producer.io_waittime_total",
+                // "kafka.producer.iotime_total",
                 "kafka.producer.metadata_age",
                 "kafka.producer.network_io_rate",
                 "kafka.producer.network_io_total",
@@ -404,7 +404,7 @@ private static void consumeRecords() {
     consumer.subscribe(TOPICS);
     Instant stopTime = Instant.now().plusSeconds(10);
     while (Instant.now().isBefore(stopTime)) {
-      consumer.poll(1_000);
+      KafkaTestUtil.poll(consumer, Duration.ofSeconds(1));
     }
   }
 

instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/KafkaClientBaseTest.java

@@ -37,6 +37,7 @@
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
@@ -147,13 +148,12 @@ void cleanupClass() {
     kafka.stop();
   }
 
-  @SuppressWarnings("PreferJavaTimeOverload")
   public void awaitUntilConsumerIsReady() throws InterruptedException {
     if (consumerReady.await(0, TimeUnit.SECONDS)) {
       return;
     }
     for (int i = 0; i < 10; i++) {
-      consumer.poll(0);
+      poll(Duration.ofMillis(100));
       if (consumerReady.await(1, TimeUnit.SECONDS)) {
         break;
       }
@@ -164,6 +164,10 @@ public void awaitUntilConsumerIsReady() throws InterruptedException {
     consumer.seekToBeginning(Collections.emptyList());
   }
 
+  public ConsumerRecords<Integer, String> poll(Duration duration) {
+    return KafkaTestUtil.poll(consumer, duration);
+  }
+
   @SuppressWarnings("deprecation") // using deprecated semconv
   protected static List<AttributeAssertion> sendAttributes(
       String messageKey, String messageValue, boolean testHeaders) {

instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/KafkaClientPropagationBaseTest.java

@@ -26,7 +26,7 @@ void testClientHeaderPropagationManualConfig() throws InterruptedException {
 
     awaitUntilConsumerIsReady();
     // check that the message was received
-    ConsumerRecords<?, ?> records = consumer.poll(Duration.ofSeconds(5).toMillis());
+    ConsumerRecords<?, ?> records = poll(Duration.ofSeconds(5));
     assertThat(records.count()).isEqualTo(1);
     for (ConsumerRecord<?, ?> record : records) {
       assertThat(record.headers().iterator().hasNext()).isEqualTo(producerPropagationEnabled);

instrumentation/kafka/kafka-clients/kafka-clients-0.11/testing/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/KafkaTestUtil.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.kafkaclients.common.v0_11.internal;
+
+import java.lang.reflect.Method;
+import java.time.Duration;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+
+@SuppressWarnings("OtelInternalJavadoc")
+public class KafkaTestUtil {
+  private static final Method consumerPollDurationMethod = getConsumerPollDurationMethod();
+
+  private static Method getConsumerPollDurationMethod() {
+    try {
+      return Consumer.class.getMethod("poll", Duration.class);
+    } catch (NoSuchMethodException e) {
+      return null;
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  public static <K, V> ConsumerRecords<K, V> poll(Consumer<K, V> consumer, Duration duration) {
+    // not present in early versions
+    if (consumerPollDurationMethod != null) {
+      try {
+        return (ConsumerRecords<K, V>) consumerPollDurationMethod.invoke(consumer, duration);
+      } catch (Exception exception) {
+        throw new IllegalStateException(exception);
+      }
+    }
+    // not present in 4.x
+    return consumer.poll(duration.toMillis());
+  }
+
+  private KafkaTestUtil() {}
+}

instrumentation/kafka/kafka-clients/kafka-clients-common-0.11/library/build.gradle.kts

@@ -3,6 +3,7 @@ plugins {
 }
 
 dependencies {
+  compileOnly(project(":muzzle"))
   compileOnly("org.apache.kafka:kafka-clients:0.11.0.0")
   compileOnly("com.google.auto.value:auto-value-annotations")
   annotationProcessor("com.google.auto.value:auto-value")

instrumentation/kafka/kafka-clients/kafka-clients-common-0.11/library/src/main/java/io/opentelemetry/instrumentation/kafkaclients/common/v0_11/internal/KafkaPropagation.java

@@ -7,6 +7,7 @@
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.record.RecordBatch;
@@ -18,6 +19,7 @@
 public final class KafkaPropagation {
 
   private static final KafkaHeadersSetter SETTER = KafkaHeadersSetter.INSTANCE;
+  private static final boolean hasMaxUsableProduceMagic = hasMaxUsableProduceMagic();
 
   // Do not inject headers for batch versions below 2
   // This is how similar check is being done in Kafka client itself:
@@ -27,7 +29,23 @@ public final class KafkaPropagation {
   // headers attempt to read messages that were produced by clients > 0.11 and the magic
   // value of the broker(s) is >= 2
   public static boolean shouldPropagate(ApiVersions apiVersions) {
-    return apiVersions.maxUsableProduceMagic() >= RecordBatch.MAGIC_VALUE_V2;
+    return !hasMaxUsableProduceMagic
+        || maxUsableProduceMagic(apiVersions) >= RecordBatch.MAGIC_VALUE_V2;
+  }
+
+  @NoMuzzle
+  private static byte maxUsableProduceMagic(ApiVersions apiVersions) {
+    return apiVersions.maxUsableProduceMagic();
+  }
+
+  private static boolean hasMaxUsableProduceMagic() {
+    try {
+      // missing in kafka 4.x
+      ApiVersions.class.getMethod("maxUsableProduceMagic");
+      return true;
+    } catch (NoSuchMethodException e) {
+      return false;
+    }
   }
 
   public static <K, V> ProducerRecord<K, V> propagateContext(