Fix spring kafka context leak when batch listener is retried

instrumentation/spring/spring-kafka-2.7/javaagent/build.gradle.kts

@@ -69,6 +69,7 @@ tasks {
   test {
     usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
 
+    systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
     jvmArgs("-Dotel.instrumentation.kafka.experimental-span-attributes=true")
     jvmArgs("-Dotel.instrumentation.messaging.experimental.receive-telemetry.enabled=true")
   }

instrumentation/spring/spring-kafka-2.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/kafka/v2_7/AbstractMessageListenerContainerInstrumentation.java

@@ -16,7 +16,6 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.springframework.kafka.listener.BatchInterceptor;
 import org.springframework.kafka.listener.RecordInterceptor;
 
 public class AbstractMessageListenerContainerInstrumentation implements TypeInstrumentation {
@@ -28,16 +27,11 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
   @Override
   public void transform(TypeTransformer transformer) {
-    // getBatchInterceptor() is called internally by AbstractMessageListenerContainer
-    // implementations
-    transformer.applyAdviceToMethod(
-        named("getBatchInterceptor")
-            .and(isProtected())
-            .and(takesArguments(0))
-            .and(returns(named("org.springframework.kafka.listener.BatchInterceptor"))),
-        this.getClass().getName() + "$GetBatchInterceptorAdvice");
     // getRecordInterceptor() is called internally by AbstractMessageListenerContainer
     // implementations
+    // for batch listeners we don't instrument getBatchInterceptor() here but instead instrument
+    // KafkaMessageListenerContainer$ListenerConsumer because spring doesn't always call the success
+    // and failure methods on a batch interceptor
     transformer.applyAdviceToMethod(
         named("getRecordInterceptor")
             .and(isProtected())
@@ -46,24 +40,6 @@ public void transform(TypeTransformer transformer) {
         this.getClass().getName() + "$GetRecordInterceptorAdvice");
   }
 
-  @SuppressWarnings("unused")
-  public static class GetBatchInterceptorAdvice {
-
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static <K, V> void onExit(
-        @Advice.Return(readOnly = false) BatchInterceptor<K, V> interceptor) {
-
-      if (interceptor == null
-          || !interceptor
-              .getClass()
-              .getName()
-              .equals(
-                  "io.opentelemetry.instrumentation.spring.kafka.v2_7.InstrumentedBatchInterceptor")) {
-        interceptor = telemetry().createBatchInterceptor(interceptor);
-      }
-    }
-  }
-
   @SuppressWarnings("unused")
   public static class GetRecordInterceptorAdvice {
 

instrumentation/spring/spring-kafka-2.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/kafka/v2_7/ListenerConsumerInstrumentation.java

@@ -5,16 +5,25 @@
 
 package io.opentelemetry.javaagent.instrumentation.spring.kafka.v2_7;
 
+import static io.opentelemetry.javaagent.instrumentation.spring.kafka.v2_7.SpringKafkaSingletons.batchProcessInstrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContext;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaConsumerContextUtil;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
 import io.opentelemetry.javaagent.bootstrap.kafka.KafkaClientsConsumerProcessTracing;
 import io.opentelemetry.javaagent.bootstrap.spring.SpringSchedulingTaskTracing;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 
 public class ListenerConsumerInstrumentation implements TypeInstrumentation {
 
@@ -29,6 +38,10 @@ public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(named("run"), this.getClass().getName() + "$RunLoopAdvice");
     transformer.applyAdviceToMethod(
         isConstructor(), this.getClass().getName() + "$ConstructorAdvice");
+    transformer.applyAdviceToMethod(
+        named("invokeBatchOnMessageWithRecordsOrList")
+            .and(takesArgument(0, named("org.apache.kafka.clients.consumer.ConsumerRecords"))),
+        this.getClass().getName() + "$InvokeBatchAdvice");
   }
 
   // this advice suppresses the CONSUMER spans created by the kafka-clients instrumentation
@@ -60,4 +73,41 @@ public static void onExit(@Advice.Enter boolean previousValue) {
       SpringSchedulingTaskTracing.setEnabled(previousValue);
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class InvokeBatchAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) ConsumerRecords<?, ?> records,
+        @Advice.FieldValue("consumer") Consumer<?, ?> consumer,
+        @Advice.Local("otelRequest") KafkaReceiveRequest request,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      KafkaConsumerContext consumerContext = KafkaConsumerContextUtil.get(records);
+      Context receiveContext = consumerContext.getContext();
+
+      // use the receive CONSUMER span as parent if it's available
+      Context parentContext = receiveContext != null ? receiveContext : Context.current();
+
+      request = KafkaReceiveRequest.create(records, consumer);
+      if (batchProcessInstrumenter().shouldStart(parentContext, request)) {
+        context = batchProcessInstrumenter().start(parentContext, request);
+        scope = context.makeCurrent();
+      }
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelRequest") KafkaReceiveRequest request,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelScope") Scope scope) {
+      if (scope == null) {
+        return;
+      }
+      scope.close();
+      batchProcessInstrumenter().end(context, request, null, throwable);
+    }
+  }
 }

instrumentation/spring/spring-kafka-2.7/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/kafka/v2_7/SpringKafkaSingletons.java

@@ -6,24 +6,48 @@
 package io.opentelemetry.javaagent.instrumentation.spring.kafka.v2_7;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaInstrumenterFactory;
+import io.opentelemetry.instrumentation.kafka.internal.KafkaReceiveRequest;
 import io.opentelemetry.instrumentation.spring.kafka.v2_7.SpringKafkaTelemetry;
+import io.opentelemetry.instrumentation.spring.kafka.v2_7.internal.SpringKafkaErrorCauseExtractor;
 import io.opentelemetry.javaagent.bootstrap.internal.ExperimentalConfig;
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
 
 public final class SpringKafkaSingletons {
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.spring-kafka-2.7";
 
   private static final SpringKafkaTelemetry TELEMETRY =
       SpringKafkaTelemetry.builder(GlobalOpenTelemetry.get())
+          .setCapturedHeaders(ExperimentalConfig.get().getMessagingHeaders())
           .setCaptureExperimentalSpanAttributes(
               InstrumentationConfig.get()
                   .getBoolean("otel.instrumentation.kafka.experimental-span-attributes", false))
           .setMessagingReceiveInstrumentationEnabled(
               ExperimentalConfig.get().messagingReceiveInstrumentationEnabled())
           .build();
+  private static final Instrumenter<KafkaReceiveRequest, Void> BATCH_PROCESS_INSTRUMENTER;
+
+  static {
+    KafkaInstrumenterFactory factory =
+        new KafkaInstrumenterFactory(GlobalOpenTelemetry.get(), INSTRUMENTATION_NAME)
+            .setCapturedHeaders(ExperimentalConfig.get().getMessagingHeaders())
+            .setCaptureExperimentalSpanAttributes(
+                InstrumentationConfig.get()
+                    .getBoolean("otel.instrumentation.kafka.experimental-span-attributes", false))
+            .setMessagingReceiveInstrumentationEnabled(
+                ExperimentalConfig.get().messagingReceiveInstrumentationEnabled())
+            .setErrorCauseExtractor(SpringKafkaErrorCauseExtractor.INSTANCE);
+    BATCH_PROCESS_INSTRUMENTER = factory.createBatchProcessInstrumenter();
+  }
 
   public static SpringKafkaTelemetry telemetry() {
     return TELEMETRY;
   }
 
+  public static Instrumenter<KafkaReceiveRequest, Void> batchProcessInstrumenter() {
+    return BATCH_PROCESS_INSTRUMENTER;
+  }
+
   private SpringKafkaSingletons() {}
 }