GH-2588: support batch recoverable DefaultAfterRollbackProcessor

* add method processBatch at `AfterRollbackProcessor`

* add opt-in property `batchRecoverAfterRollback` at `ContainerProperties`

* change format to `BatchListenerFailedException.getMessage`

* add batch recoverable after rollback unit test

Author: Zhiyang.Wang1

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2021 the original author or authors.
+ * Copyright 2018-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
 
@@ -63,6 +64,12 @@ public interface AfterRollbackProcessor<K, V> {
 	void process(List<ConsumerRecord<K, V>> records, Consumer<K, V> consumer,
 			MessageListenerContainer container, Exception exception, boolean recoverable, EOSMode eosMode);
 
+	default void processBatch(ConsumerRecords<K, V> records, List<ConsumerRecord<K, V>> recordList,
+			Consumer<K, V> consumer, MessageListenerContainer container, Exception exception,
+			boolean recoverable, ContainerProperties.EOSMode eosMode) {
+		process(recordList, consumer, container, exception, recoverable, eosMode);
+	}
+
 	/**
 	 * Optional method to clear thread state; will be called just before a consumer
 	 * thread terminates.

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

@@ -98,9 +98,9 @@ public int getIndex() {
 
 	@Override
 	public String getMessage() {
-		return super.getMessage() + (this.record != null
+		return super.getMessage() + " " + (this.record != null
 				? (this.record.topic() + "-" + this.record.partition() + "@" + this.record.offset())
-				: (" @-" + this.index));
+				: ("@-" + this.index));
 	}
 
 }

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

@@ -258,6 +258,8 @@ public enum EOSMode {
 
 	private PlatformTransactionManager transactionManager;
 
+	private boolean batchRecoverAfterRollback = false;
+
 	private int monitorInterval = DEFAULT_MONITOR_INTERVAL;
 
 	private TaskScheduler scheduler;
@@ -543,6 +545,14 @@ public void setTransactionManager(@Nullable PlatformTransactionManager transacti
 		this.transactionManager = transactionManager;
 	}
 
+	public boolean isBatchRecoverAfterRollback() {
+		return this.batchRecoverAfterRollback;
+	}
+
+	public void setBatchRecoverAfterRollback(boolean batchRecoverAfterRollback) {
+		this.batchRecoverAfterRollback = batchRecoverAfterRollback;
+	}
+
 	public int getMonitorInterval() {
 		return this.monitorInterval;
 	}

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

@@ -17,16 +17,19 @@
 package org.springframework.kafka.listener;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiConsumer;
 
 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.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 
+import org.springframework.kafka.KafkaException;
 import org.springframework.kafka.core.KafkaOperations;
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
 import org.springframework.lang.Nullable;
@@ -60,7 +63,9 @@ public class DefaultAfterRollbackProcessor<K, V> extends FailedRecordProcessor
 
 	private final BackOff backOff;
 
-	private KafkaOperations<?, ?> kafkaTemplate;
+	private final KafkaOperations<?, ?> kafkaTemplate;
+
+	private final BiConsumer<ConsumerRecords<?, ?>, Exception> recoverer;
 
 	/**
 	 * Construct an instance with the default recoverer which simply logs the record after
@@ -143,6 +148,11 @@ public DefaultAfterRollbackProcessor(@Nullable BiConsumer<ConsumerRecord<?, ?>,
 		super.setCommitRecovered(commitRecovered);
 		checkConfig();
 		this.backOff = backOff;
+		this.recoverer = (crs, ex) -> {
+			if (recoverer != null && !crs.isEmpty()) {
+				crs.spliterator().forEachRemaining(rec -> recoverer.accept(rec, ex));
+			}
+		};
 	}
 
 	private void checkConfig() {
@@ -176,6 +186,53 @@ && isCommitRecovered() && this.kafkaTemplate.isTransactional()) {
 
 	}
 
+	@SuppressWarnings({ "unchecked", "rawtypes"})
+	@Override
+	public void processBatch(ConsumerRecords<K, V> records, List<ConsumerRecord<K, V>> recordList, Consumer<K, V> consumer,
+			@Nullable MessageListenerContainer container, Exception exception, boolean recoverable, EOSMode eosMode) {
+
+		if (recoverable && isCommitRecovered()) {
+			long nextBackOff = ListenerUtils.nextBackOff(this.backOff, this.backOffs);
+			if (nextBackOff != BackOffExecution.STOP) {
+				SeekUtils.doSeeksToBegin((List) recordList, consumer, this.logger);
+				try {
+					ListenerUtils.stoppableSleep(container, nextBackOff);
+				}
+				catch (InterruptedException e) {
+					Thread.currentThread().interrupt();
+				}
+				return;
+			}
+
+			try {
+				this.recoverer.accept(records, exception);
+				Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+				records.forEach(rec -> offsets.put(new TopicPartition(rec.topic(), rec.partition()),
+						ListenerUtils.createOffsetAndMetadata(container, rec.offset() + 1)));
+				if (offsets.size() > 0 && this.kafkaTemplate != null && this.kafkaTemplate.isTransactional()) {
+					this.kafkaTemplate.sendOffsetsToTransaction(offsets, consumer.groupMetadata());
+				}
+				clearThreadState();
+			}
+			catch (Exception ex) {
+				SeekUtils.doSeeksToBegin((List) recordList, consumer, this.logger);
+				logger.error(ex, () -> "Recoverer threw an exception; re-seeking batch");
+				throw ex;
+			}
+			return;
+		}
+
+		try {
+			process(recordList, consumer, container, exception, false, eosMode);
+		}
+		catch (KafkaException ke) {
+			ke.selfLog("AfterRollbackProcessor threw an exception", this.logger);
+		}
+		catch (Exception ex) {
+			this.logger.error(ex, "AfterRollbackProcessor threw an exception");
+		}
+	}
+
 	@Override
 	public boolean isProcessInTransaction() {
 		return isCommitRecovered();

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

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
@@ -660,6 +661,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final boolean wantsFullRecords;
 
+		private final boolean wantsBatchRecoverAfterRollback;
+
 		private final boolean autoCommit;
 
 		private final boolean isManualAck;
@@ -882,6 +885,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 			this.clientId = determineClientId();
 			this.transactionTemplate = determineTransactionTemplate();
+			this.wantsBatchRecoverAfterRollback = this.containerProperties.isBatchRecoverAfterRollback();
 			this.genericListener = listener;
 			this.consumerSeekAwareListener = checkConsumerSeekAware(listener);
 			this.commitCurrentOnAssignment = determineCommitCurrent(consumerProperties,
@@ -2190,37 +2194,20 @@ private void batchRollback(final ConsumerRecords<K, V> records,
 
 					@Override
 					protected void doInTransactionWithoutResult(TransactionStatus status) {
-						batchAfterRollback(records, recordList, e, afterRollbackProcessorToUse);
+						afterRollbackProcessorToUse.processBatch(records,
+								Objects.requireNonNullElseGet(recordList, () -> createRecordList(records)),
+								ListenerConsumer.this.consumer,
+								KafkaMessageListenerContainer.this.thisOrParentContainer, e,
+								ListenerConsumer.this.wantsBatchRecoverAfterRollback, ListenerConsumer.this.eosMode);
 					}
 
 				});
 			}
 			else {
-				batchAfterRollback(records, recordList, e, afterRollbackProcessorToUse);
-			}
-		}
-
-		private void batchAfterRollback(final ConsumerRecords<K, V> records,
-				@Nullable final List<ConsumerRecord<K, V>> recordList, RuntimeException rollbackException,
-				AfterRollbackProcessor<K, V> afterRollbackProcessorToUse) {
-
-			try {
-				if (recordList == null) {
-					afterRollbackProcessorToUse.process(createRecordList(records), this.consumer,
-							KafkaMessageListenerContainer.this.thisOrParentContainer,  rollbackException, false,
-							this.eosMode);
-				}
-				else {
-					afterRollbackProcessorToUse.process(recordList, this.consumer,
-							KafkaMessageListenerContainer.this.thisOrParentContainer,  rollbackException, false,
-							this.eosMode);
-				}
-			}
-			catch (KafkaException ke) {
-				ke.selfLog("AfterRollbackProcessor threw an exception", this.logger);
-			}
-			catch (Exception ex) {
-				this.logger.error(ex, "AfterRollbackProcessor threw an exception");
+				afterRollbackProcessorToUse.processBatch(records,
+						Objects.requireNonNullElseGet(recordList, () -> createRecordList(records)), this.consumer,
+						KafkaMessageListenerContainer.this.thisOrParentContainer, e,
+						this.wantsBatchRecoverAfterRollback, this.eosMode);
 			}
 		}
 

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

@@ -126,12 +126,7 @@ public static void unrecoverableBackOff(BackOff backOff, Map<Thread, BackOffExec
 			Map<Thread, Long> lastIntervals, MessageListenerContainer container) throws InterruptedException {
 
 		Thread currentThread = Thread.currentThread();
-		BackOffExecution backOffExecution = executions.get(currentThread);
-		if (backOffExecution == null) {
-			backOffExecution = backOff.start();
-			executions.put(currentThread, backOffExecution);
-		}
-		Long interval = backOffExecution.nextBackOff();
+		Long interval = nextBackOff(backOff, executions);
 		if (interval == BackOffExecution.STOP) {
 			interval = lastIntervals.get(currentThread);
 			if (interval == null) {
@@ -144,6 +139,17 @@ public static void unrecoverableBackOff(BackOff backOff, Map<Thread, BackOffExec
 		}
 	}
 
+	public static long nextBackOff(BackOff backOff, Map<Thread, BackOffExecution> executions) {
+
+		Thread currentThread = Thread.currentThread();
+		BackOffExecution backOffExecution = executions.get(currentThread);
+		if (backOffExecution == null) {
+			backOffExecution = backOff.start();
+			executions.put(currentThread, backOffExecution);
+		}
+		return backOffExecution.nextBackOff();
+	}
+
 	/**
 	 * Sleep for the desired timeout, as long as the container continues to run.
 	 * @param container the container.

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

@@ -133,6 +133,17 @@ public static boolean doSeeks(List<ConsumerRecord<?, ?>> records, Consumer<?, ?>
 		return skipped.get();
 	}
 
+	public static void doSeeksToBegin(List<ConsumerRecord<?, ?>> records, Consumer<?, ?> consumer,
+			LogAccessor logger) {
+
+		Map<TopicPartition, Long> partitions = new LinkedHashMap<>();
+		records.forEach(record -> {
+			partitions.computeIfAbsent(new TopicPartition(record.topic(), record.partition()),
+					offset -> record.offset());
+		});
+		seekPartitions(consumer, partitions, logger);
+	}
+
 	/**
 	 * Perform seek operations on each partition.
 	 * @param consumer the consumer.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2022 the original author or authors.
+ * Copyright 2019-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,14 +27,18 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.common.TopicPartition;
 import org.junit.jupiter.api.Test;
 import org.mockito.InOrder;
@@ -100,8 +104,8 @@ void testClassifier() {
 	}
 
 	@Test
-	void testBatchBackOff() {
-		AtomicReference<ConsumerRecord<?, ?>> recovered = new AtomicReference<>();
+	void testBackOffNoBatchRecover() {
+
 		@SuppressWarnings("unchecked")
 		KafkaOperations<String, String> template = mock(KafkaOperations.class);
 		given(template.isTransactional()).willReturn(true);
@@ -118,37 +122,25 @@ void testBatchBackOff() {
 		ConsumerRecord<String, String> record1 = new ConsumerRecord<>("foo", 0, 0L, "foo", "bar");
 		ConsumerRecord<String, String> record2 = new ConsumerRecord<>("foo", 1, 1L, "foo", "bar");
 		List<ConsumerRecord<String, String>> records = Arrays.asList(record1, record2);
+		Map<TopicPartition, List<ConsumerRecord<String, String>>> map = new HashMap<>();
+		records.forEach(rec -> map.computeIfAbsent(new TopicPartition(rec.topic(), rec.partition()),
+				tp -> new ArrayList<>()).add(rec));
+		ConsumerRecords<String, String> consumerRecords = new ConsumerRecords<>(map);
 		IllegalStateException illegalState = new IllegalStateException();
 		@SuppressWarnings("unchecked")
 		Consumer<String, String> consumer = mock(Consumer.class);
 		given(consumer.groupMetadata()).willReturn(new ConsumerGroupMetadata("foo"));
 		MessageListenerContainer container = mock(MessageListenerContainer.class);
 		given(container.isRunning()).willReturn(true);
-		processor.process(records, consumer, container, illegalState, false, EOSMode.V2);
-		processor.process(records, consumer, container, illegalState, false, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, illegalState, false, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, illegalState, false, EOSMode.V2);
 		verify(backOff, times(2)).start();
 		verify(execution.get(), times(2)).nextBackOff();
 		processor.clearThreadState();
-		processor.process(records, consumer, container, illegalState, false, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, illegalState, false, EOSMode.V2);
 		verify(backOff, times(3)).start();
 	}
 
-	void testEarlyExitBackOff() {
-		DefaultAfterRollbackProcessor<String, String> processor = new DefaultAfterRollbackProcessor<>(
-				new FixedBackOff(1, 10_000));
-		@SuppressWarnings("unchecked")
-		Consumer<String, String> consumer = mock(Consumer.class);
-		ConsumerRecord<String, String> record1 = new ConsumerRecord<>("foo", 0, 0L, "foo", "bar");
-		ConsumerRecord<String, String> record2 = new ConsumerRecord<>("foo", 1, 1L, "foo", "bar");
-		List<ConsumerRecord<String, String>> records = Arrays.asList(record1, record2);
-		IllegalStateException illegalState = new IllegalStateException();
-		MessageListenerContainer container = mock(MessageListenerContainer.class);
-		given(container.isRunning()).willReturn(false);
-		long t1 = System.currentTimeMillis();
-		processor.process(records, consumer, container, illegalState, true, EOSMode.V2);
-		assertThat(System.currentTimeMillis() < t1 + 5_000);
-	}
-
 	@Test
 	void testNoEarlyExitBackOff() {
 		DefaultAfterRollbackProcessor<String, String> processor = new DefaultAfterRollbackProcessor<>(