fix instrumented type & depth count

DataDog · Nov 16, 2023 · d7b4f5c · d7b4f5c
1 parent 619cb0f
commit d7b4f5c
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/...c/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java
@@ -37,15 +37,18 @@
 
 @AutoService(Instrumenter.class)
 public final class KafkaProducerInstrumentation extends Instrumenter.Tracing
-    implements Instrumenter.ForSingleType {
+    implements Instrumenter.ForKnownTypes {
 
   public KafkaProducerInstrumentation() {
     super("kafka");
   }
 
   @Override
-  public String instrumentedType() {
-    return "org.apache.kafka.clients.producer.KafkaProducer";
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.kafka.clients.producer.KafkaProducer",
+      "org.apache.kafka.common.record.AbstractRecords"
+    };
   }
 
   @Override
@@ -131,6 +134,9 @@ record =
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+      // reset depth count for EstimateSizeAdvice
+      CallDepthThreadLocalMap.reset(AbstractRecords.class);
+
       PRODUCER_DECORATE.onError(scope, throwable);
       PRODUCER_DECORATE.beforeFinish(scope);
       scope.close();