From 0e7a9c6af6e85bc6e900cc78a33948af3563faaa Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@pivotal.io>
Date: Wed, 13 Nov 2019 13:43:39 -0500
Subject: [PATCH] Fix BatchErrorHandler container argument

When passing in the container to batch error handlers, if there is a
parent container, it should be passed in, e.g. if the error handler wants
to stop the entire container
---
 .../KafkaMessageListenerContainer.java        |  2 +-
 .../SeekToCurrentBatchErrorHandlerTests.java  | 41 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java
index 87a4ccf1..9af3bc3f 100644
--- a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java
+++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java
@@ -877,7 +877,7 @@ public class KafkaMessageListenerContainer<K, V> // NOSONAR comment density
 				}
 				else if (this.isBatchListener && this.batchErrorHandler != null) {
 					this.batchErrorHandler.handle(e, new ConsumerRecords<K, V>(Collections.emptyMap()), this.consumer,
-							KafkaMessageListenerContainer.this);
+							KafkaMessageListenerContainer.this.container);
 				}
 				else {
 					this.logger.error("Consumer exception", e);
diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/SeekToCurrentBatchErrorHandlerTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/SeekToCurrentBatchErrorHandlerTests.java
index 66a2cb49..db76bf7d 100644
--- a/spring-kafka/src/test/java/org/springframework/kafka/listener/SeekToCurrentBatchErrorHandlerTests.java
+++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/SeekToCurrentBatchErrorHandlerTests.java
@@ -18,6 +18,7 @@ package org.springframework.kafka.listener;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.BDDMockito.willAnswer;
 import static org.mockito.Mockito.inOrder;
@@ -34,6 +35,7 @@ import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
@@ -116,6 +118,45 @@ public class SeekToCurrentBatchErrorHandlerTests {
 		assertThat(((ListenerExecutionFailedException) this.config.ehException).getGroupId()).isEqualTo(CONTAINER_ID);
 	}
 
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	@Test
+	public void verifyCorrectContainer() throws InterruptedException {
+		ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
+		final Consumer consumer = mock(Consumer.class);
+		AtomicBoolean first = new AtomicBoolean(true);
+		willAnswer(invocation -> {
+			if (first.getAndSet(false)) {
+				throw new IllegalStateException("intentional");
+			}
+			Thread.sleep(50);
+			return new ConsumerRecords(Collections.emptyMap());
+		}).given(consumer).poll(any());
+		given(consumerFactory.createConsumer(anyString(), anyString(), anyString(), any()))
+						.willReturn(consumer);
+		ContainerProperties containerProperties = new ContainerProperties("foo");
+		containerProperties.setGroupId("grp");
+		containerProperties.setMessageListener((BatchMessageListener) record -> { });
+		containerProperties.setMissingTopicsFatal(false);
+		ConcurrentMessageListenerContainer container = new ConcurrentMessageListenerContainer<>(consumerFactory,
+				containerProperties);
+		AtomicReference<MessageListenerContainer> parent = new AtomicReference<>();
+		CountDownLatch latch = new CountDownLatch(1);
+		container.setBatchErrorHandler(new ContainerAwareBatchErrorHandler() {
+
+			@Override
+			public void handle(Exception thrownException, ConsumerRecords<?, ?> data, Consumer<?, ?> consumer,
+					MessageListenerContainer container) {
+
+				parent.set(container);
+				latch.countDown();
+			}
+		});
+		container.start();
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		container.stop();
+		assertThat(parent.get()).isSameAs(container);
+	}
+
 	@Configuration
 	@EnableKafka
 	public static class Config {