Update muzzleCheck for kafka

dd-java-agent/instrumentation/kafka-clients-3.8/build.gradle

@@ -6,7 +6,7 @@ muzzle {
     group = "org.apache.kafka"
     module = "kafka-clients"
     versions = "[3.8.0,)"
-    assertInverse = false
+    assertInverse = true
   }
 }
 

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java

@@ -7,7 +7,6 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.internals.ConsumerDelegate;
 import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
@@ -54,9 +53,8 @@ public static void captureGroup(
     }
   }
 
-  public static void muzzleCheck(ConsumerRecord record) {
-    // KafkaConsumerInstrumentation only applies for kafka versions with headers
-    // Make an explicit call so KafkaConsumerGroupInstrumentation does the same
-    record.headers();
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
   }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConsumerCoordinatorAdvice.java

@@ -12,9 +12,9 @@
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.Metadata;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 import org.apache.kafka.clients.consumer.internals.RequestFuture;
 import org.apache.kafka.common.TopicPartition;
 
@@ -66,9 +66,8 @@ public static void trackCommitOffset(
     }
   }
 
-  public static void muzzleCheck(ConsumerRecord record) {
-    // KafkaConsumerInstrumentation only applies for kafka versions with headers
-    // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
-    record.headers();
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
   }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/IterableAdvice.java

@@ -5,6 +5,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class IterableAdvice {
 
@@ -31,4 +32,9 @@ public static void wrap(
               bootstrapServers);
     }
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/IteratorAdvice.java

@@ -6,6 +6,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class IteratorAdvice {
 
@@ -32,4 +33,9 @@ public static void wrap(
               bootstrapServers);
     }
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/KafkaConsumerInstrumentationHelper.java

@@ -3,6 +3,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.ContextStore;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class KafkaConsumerInstrumentationHelper {
   public static String extractGroup(KafkaConsumerInfo kafkaConsumerInfo) {
@@ -26,4 +27,9 @@ public static String extractClusterId(
   public static String extractBootstrapServers(KafkaConsumerInfo kafkaConsumerInfo) {
     return kafkaConsumerInfo == null ? null : kafkaConsumerInfo.getBootstrapServers().get();
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java

@@ -6,9 +6,9 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
 import org.apache.kafka.clients.consumer.internals.ConsumerDelegate;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class LegacyConstructorAdvice {
   // new - capture the ConsumerDelegate instead of KafkaConsumer
@@ -46,9 +46,8 @@ public static void captureGroup(
     }
   }
 
-  public static void muzzleCheck(ConsumerRecord record) {
-    // KafkaConsumerInstrumentation only applies for kafka versions with headers
-    // Make an explicit call so KafkaConsumerGroupInstrumentation does the same
-    record.headers();
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
   }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ListAdvice.java

@@ -9,6 +9,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class ListAdvice {
 
@@ -30,4 +31,9 @@ public static void wrap(
               iterable, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
     }
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java

@@ -3,7 +3,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.Metadata;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 import org.apache.kafka.common.requests.MetadataResponse;
 
 public class MetadataUpdate22AndAfterAdvice {
@@ -15,9 +15,8 @@ public static void onEnter(
     }
   }
 
-  public static void muzzleCheck(ConsumerRecord record) {
-    // KafkaConsumerInstrumentation only applies for kafka versions with headers
-    // Make an explicit call so MetadataInstrumentation does the same
-    record.headers();
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
   }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java

@@ -3,7 +3,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.Metadata;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 import org.apache.kafka.common.Cluster;
 
 public class MetadataUpdateBefore22Advice {
@@ -16,9 +16,8 @@ public static void onEnter(
     }
   }
 
-  public static void muzzleCheck(ConsumerRecord record) {
-    // KafkaConsumerInstrumentation only applies for kafka versions with headers
-    // Make an explicit call so MetadataInstrumentation does the same
-    record.headers();
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
   }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/OffsetCommitCallbackInvokerAdvice.java

@@ -15,4 +15,9 @@ public static void onEnqueue(
             .get(callbackInvoker);
     callback = new DDOffsetCommitCallback(callback, kafkaConsumerInfo);
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/PayloadSizeAdvice.java

@@ -5,6 +5,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.StatsPoint;
 import net.bytebuddy.asm.Advice;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 public class PayloadSizeAdvice {
 
@@ -32,4 +33,9 @@ public static void onEnter(@Advice.Argument(value = 0) int estimatedPayloadSize)
       AgentTracer.get().getDataStreamsMonitoring().add(updated);
     }
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java

@@ -23,6 +23,7 @@
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
@@ -114,4 +115,9 @@ public static void stopSpan(
     PRODUCER_DECORATE.beforeFinish(scope);
     scope.close();
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }

dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/RecordsAdvice.java

@@ -13,6 +13,7 @@
 import net.bytebuddy.asm.Advice;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.internals.ConsumerDelegate;
+import org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker;
 
 /**
  * this method transfers the consumer group from the KafkaConsumer class key to the ConsumerRecords
@@ -60,4 +61,9 @@ public static void captureGroup(
     span.finish();
     scope.close();
   }
+
+  public static void muzzleCheck(OffsetCommitCallbackInvoker invoker) {
+    // Only applies for kafka versions with OffsetCommitCallbackInvoker
+    invoker.executeCallbacks();
+  }
 }