Update muzzleCheck for kafka

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

@@ -2,7 +2,7 @@ muzzle {
   pass {
     group = "org.apache.kafka"
     module = "kafka-clients"
-    versions = "[0.11.0.0,)"
+    versions = "[0.11.0.0,3.7.0)"
     assertInverse = true
   }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/ConsumerCoordinatorInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.core.datastreams.TagsProcessor.CONSUMER_GROUP_TAG;
 import static datadog.trace.core.datastreams.TagsProcessor.KAFKA_CLUSTER_ID_TAG;
@@ -17,11 +18,13 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
 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.RequestFuture;
+import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.TopicPartition;
 
 @AutoService(InstrumenterModule.class)
@@ -32,6 +35,13 @@ public ConsumerCoordinatorInstrumentation() {
     super("kafka");
   }
 
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public Map<String, String> contextStore() {
     Map<String, String> contextStores = new HashMap<>();
@@ -107,10 +117,12 @@ public static void trackCommitOffset(
       }
     }
 
-    public static void muzzleCheck(ConsumerRecord record) {
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
       // KafkaConsumerInstrumentation only applies for kafka versions with headers
       // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
       record.headers();
+      record.checksum();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
     }
   }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
@@ -9,6 +10,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -24,12 +26,14 @@
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.internals.ConsumerCoordinator;
+import org.apache.kafka.clients.producer.Producer;
 
 /**
  * This instrumentation saves additional information from the KafkaConsumer, such as consumer group
@@ -57,6 +61,13 @@ public Map<String, String> contextStore() {
     return contextStores;
   }
 
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public String instrumentedType() {
     return "org.apache.kafka.clients.consumer.KafkaConsumer";
@@ -176,4 +187,11 @@ public static void captureGroup(
       scope.close();
     }
   }
+
+  public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+    // KafkaConsumerInstrumentation only applies for kafka versions with headers
+    // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+    record.headers();
+    producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+  }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInstrumentation.java

@@ -1,10 +1,12 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.CONSUMER_DECORATE;
 import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_CONSUME;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -18,9 +20,11 @@
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
 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.producer.Producer;
 
 @AutoService(InstrumenterModule.class)
 public final class KafkaConsumerInstrumentation extends InstrumenterModule.Tracing
@@ -39,6 +43,12 @@ public Map<String, String> contextStore() {
     return contextStores;
   }
 
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public String instrumentedType() {
     return "org.apache.kafka.clients.consumer.ConsumerRecords";
@@ -108,6 +118,14 @@ public static void wrap(
                 iterable, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
       }
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      record.checksum();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   public static class ListAdvice {
@@ -130,6 +148,14 @@ public static void wrap(
                 iterable, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
       }
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      record.checksum();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   public static class IteratorAdvice {
@@ -152,5 +178,13 @@ public static void wrap(
                 iterator, KAFKA_CONSUME, CONSUMER_DECORATE, group, clusterId, bootstrapServers);
       }
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      record.checksum();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaDeserializerInstrumentation.java

@@ -1,8 +1,10 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasInterface;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
@@ -20,6 +22,8 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.serialization.Deserializer;
 
 @AutoService(InstrumenterModule.class)
@@ -44,6 +48,12 @@ public KafkaDeserializerInstrumentation() {
     super("kafka");
   }
 
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public String hierarchyMarkerType() {
     return DESERIALIZER_CLASS;
@@ -95,6 +105,13 @@ public static void configure(
           InstrumentationContext.get(Deserializer.class, Boolean.class);
       KafkaIastHelper.configure(store, deserializer, isKey);
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   @SuppressWarnings("rawtypes")
@@ -120,6 +137,13 @@ public static void afterDeserialize(
           InstrumentationContext.get(Deserializer.class, Boolean.class);
       KafkaIastHelper.afterDeserialize(ctx, store, deserializer, result);
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   @SuppressWarnings({"rawtypes", "JavaExistingMethodCanBeUsed"})
@@ -145,6 +169,13 @@ public static void afterDeserialize(
           InstrumentationContext.get(Deserializer.class, Boolean.class);
       KafkaIastHelper.afterDeserialize(ctx, store, deserializer, result);
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   @SuppressWarnings({"rawtypes", "JavaExistingMethodCanBeUsed"})
@@ -170,5 +201,12 @@ public static void afterDeserialize(
           InstrumentationContext.get(Deserializer.class, Boolean.class);
       KafkaIastHelper.afterDeserialize(ctx, store, deserializer, result);
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
@@ -18,6 +19,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPrivate;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
@@ -33,9 +35,12 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.kafka.clients.ApiVersions;
 import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.internals.Sender;
@@ -49,6 +54,12 @@ public KafkaProducerInstrumentation() {
     super("kafka");
   }
 
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public String instrumentedType() {
     return "org.apache.kafka.clients.producer.KafkaProducer";
@@ -177,6 +188,13 @@ public static void stopSpan(
       PRODUCER_DECORATE.beforeFinish(scope);
       scope.close();
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 
   public static class PayloadSizeAdvice {
@@ -205,5 +223,12 @@ public static void onEnter(@Advice.Argument(value = 0) int estimatedPayloadSize)
         AgentTracer.get().getDataStreamsMonitoring().add(updated);
       }
     }
+
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
+      // KafkaConsumerInstrumentation only applies for kafka versions with headers
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
+      record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
+    }
   }
 }

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java

@@ -1,9 +1,11 @@
 package datadog.trace.instrumentation.kafka_clients;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
@@ -16,6 +18,7 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.requests.MetadataResponse;
 
@@ -27,6 +30,12 @@ public MetadataInstrumentation() {
     super("kafka");
   }
 
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Avoid matching kafka 3.8 which has its own instrumentation
+    return not(
+        hasClassNamed("org.apache.kafka.clients.consumer.internals.OffsetCommitCallbackInvoker"));
+  }
+
   @Override
   public String hierarchyMarkerType() {
     return "org.apache.kafka.clients.Metadata";
@@ -88,10 +97,11 @@ public static void onEnter(
       }
     }
 
-    public static void muzzleCheck(ConsumerRecord record) {
+    public static void muzzleCheck(ConsumerRecord record, Producer producer) {
       // KafkaConsumerInstrumentation only applies for kafka versions with headers
-      // Make an explicit call so MetadataInstrumentation does the same
+      // Make an explicit call so ConsumerCoordinatorInstrumentation does the same
       record.headers();
+      producer.close(2, java.util.concurrent.TimeUnit.SECONDS);
     }
   }
 }