GH-2340: Fix Retrying Batch Error Handling

Resolves #2340

The `RetryingBatchErrorHandler` - now called the `FallbackBatchErrorHandler`
pauses and resumes the consumer during retries, to allow it to poll the
consumer to avoid a forced rebalance.

However, if a normal rebalance occurs, for example if a new member joins,
the error handler does not re-pause the consumer and silently consumes
new records.

Add a mechanism to always re-pause the consume when in this retry mode.

**cherry-pick to 2.9.x, 2.8.x**

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/listener/CommonErrorHandler.java

@@ -16,12 +16,14 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.Collection;
 import java.util.List;
 
 import org.apache.commons.logging.LogFactory;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
 
 import org.springframework.kafka.support.TopicPartitionOffset;
 
@@ -150,4 +152,13 @@ default void setAckAfterHandle(boolean ack) {
 		throw new UnsupportedOperationException("This error handler does not support setting this property");
 	}
 
+	/**
+	 * Called when partitions are assigned.
+	 * @param consumer the consumer.
+	 * @param partitions the newly assigned partitions.
+	 * @since 2.8.8
+	 */
+	default void onPartitionsAssigned(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+	}
+
 }

spring-kafka/src/main/java/org/springframework/kafka/listener/DefaultErrorHandler.java

@@ -16,11 +16,13 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.Collection;
 import java.util.List;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.SerializationException;
 
 import org.springframework.lang.Nullable;
@@ -157,4 +159,9 @@ public void handleOtherException(Exception thrownException, Consumer<?, ?> consu
 		}
 	}
 
+	@Override
+	public void onPartitionsAssigned(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+		getFallbackBatchHandler().onPartitionsAssigned(consumer, partitions);
+	}
+
 }

spring-kafka/src/main/java/org/springframework/kafka/listener/ErrorHandlerAdapter.java

@@ -16,12 +16,14 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
 
 import org.springframework.kafka.support.TopicPartitionOffset;
 import org.springframework.util.Assert;
@@ -157,5 +159,13 @@ public void handleBatch(Exception thrownException, ConsumerRecords<?, ?> data, C
 		}
 	}
 
+	@SuppressWarnings("deprecation")
+	@Override
+	public void onPartitionsAssigned(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+		if (this.batchErrorHandler instanceof RetryingBatchErrorHandler) {
+			((FallbackBatchErrorHandler) this.batchErrorHandler).onPartitionsAssigned(consumer, partitions);
+		}
+	}
+
 }
 

spring-kafka/src/main/java/org/springframework/kafka/listener/FailedBatchProcessor.java

@@ -70,6 +70,15 @@ public FailedBatchProcessor(@Nullable BiConsumer<ConsumerRecord<?, ?>, Exception
 		this.fallbackBatchHandler = fallbackHandler;
 	}
 
+	/**
+	 * Return the fallback batch error handler.
+	 * @return the handler.
+	 * @since 2.8.8
+	 */
+	protected CommonErrorHandler getFallbackBatchHandler() {
+		return this.fallbackBatchHandler;
+	}
+
 	protected void doHandle(Exception thrownException, ConsumerRecords<?, ?> data, Consumer<?, ?> consumer,
 			MessageListenerContainer container, Runnable invokeListener) {
 

spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java

@@ -3373,6 +3373,10 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 					ListenerConsumer.this.firstPoll = true;
 					ListenerConsumer.this.consumerSeekAwareListener.onFirstPoll();
 				}
+				if (ListenerConsumer.this.commonErrorHandler != null) {
+					ListenerConsumer.this.commonErrorHandler.onPartitionsAssigned(ListenerConsumer.this.consumer,
+							partitions);
+				}
 			}
 
 			private void repauseIfNeeded(Collection<TopicPartition> partitions) {

spring-kafka/src/main/java/org/springframework/kafka/listener/RetryingBatchErrorHandler.java

@@ -16,11 +16,13 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.Collection;
 import java.util.function.BiConsumer;
 
 import org.apache.commons.logging.LogFactory;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
 
 import org.springframework.core.log.LogAccessor;
 import org.springframework.lang.Nullable;
@@ -55,6 +57,8 @@ public class RetryingBatchErrorHandler extends KafkaExceptionLogLevelAware
 
 	private boolean ackAfterHandle = true;
 
+	private boolean retrying;
+
 	/**
 	 * Construct an instance with a default {@link FixedBackOff} (unlimited attempts with
 	 * a 5 second back off).
@@ -100,8 +104,16 @@ public void handle(Exception thrownException, @Nullable ConsumerRecords<?, ?> re
 			this.logger.error(thrownException, "Called with no records; consumer exception");
 			return;
 		}
+		this.retrying = true;
 		ErrorHandlingUtils.retryBatch(thrownException, records, consumer, container, invokeListener, this.backOff,
 				this.seeker, this.recoverer, this.logger, getLogLevel());
+		this.retrying = false;
+	}
+
+	public void onPartitionsAssigned(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+		if (this.retrying) {
+			consumer.pause(consumer.assignment());
+		}
 	}
 
 }

spring-kafka/src/test/java/org/springframework/kafka/listener/FallbackBatchErrorHandlerTests.java

@@ -21,6 +21,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.BDDMockito.willAnswer;
+import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -38,6 +39,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.common.TopicPartition;
 import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
 
 import org.springframework.kafka.KafkaException;
 import org.springframework.util.backoff.FixedBackOff;
@@ -168,4 +170,39 @@ void exitOnContainerStop() {
 		assertThat(this.invoked).isEqualTo(1);
 	}
 
+	@Test
+	void rePauseOnRebalance() {
+		this.invoked = 0;
+		List<ConsumerRecord<?, ?>> recovered = new ArrayList<>();
+		FallbackBatchErrorHandler eh = new FallbackBatchErrorHandler(new FixedBackOff(0L, 1L), (cr, ex) ->  {
+			recovered.add(cr);
+		});
+		Map<TopicPartition, List<ConsumerRecord<Object, Object>>> map = new HashMap<>();
+		map.put(new TopicPartition("foo", 0),
+				Collections.singletonList(new ConsumerRecord<>("foo", 0, 0L, "foo", "bar")));
+		map.put(new TopicPartition("foo", 1),
+				Collections.singletonList(new ConsumerRecord<>("foo", 1, 0L, "foo", "bar")));
+		ConsumerRecords<?, ?> records = new ConsumerRecords<>(map);
+		Consumer<?, ?> consumer = mock(Consumer.class);
+		willAnswer(inv -> {
+			eh.onPartitionsAssigned(consumer, List.of(new TopicPartition("foo", 0), new TopicPartition("foo", 1)));
+			return records;
+		}).given(consumer).poll(any());
+		MessageListenerContainer container = mock(MessageListenerContainer.class);
+		given(container.isRunning()).willReturn(true);
+		eh.handle(new RuntimeException(), records, consumer, container, () -> {
+			this.invoked++;
+			throw new RuntimeException();
+		});
+		assertThat(this.invoked).isEqualTo(1);
+		assertThat(recovered).hasSize(2);
+		InOrder inOrder = inOrder(consumer);
+		inOrder.verify(consumer).pause(any());
+		inOrder.verify(consumer).poll(any());
+		inOrder.verify(consumer).pause(any());
+		inOrder.verify(consumer).resume(any());
+		verify(consumer, times(3)).assignment();
+		verifyNoMoreInteractions(consumer);
+	}
+
 }