Resolves #2588

GH-2588: ARBP support BatchListenerFailedException

* support retry and recover in `DefaultAfterRollbackProcessor.processBatch`

* add nextBackOff at `ListenerUtils`

* fix assertThat `KafkaHeaders.DLT_ORIGINAL_OFFSET` at `TransactionalContainerTests.testMaxFailures`

Author: Zhiyang.Wang1

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

@@ -32,8 +32,10 @@
 
 import org.springframework.kafka.core.KafkaOperations;
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
+import org.springframework.kafka.support.KafkaUtils;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
 import org.springframework.util.backoff.BackOff;
 import org.springframework.util.backoff.BackOffExecution;
 
@@ -63,7 +65,9 @@ public class DefaultAfterRollbackProcessor<K, V> extends FailedBatchProcessor
 
 	private final BackOff backOff;
 
-	private KafkaOperations<?, ?> kafkaTemplate;
+	private final KafkaOperations<?, ?> kafkaTemplate;
+
+	private final BiConsumer<List<ConsumerRecord<?, ?>>, Exception> recoverer;
 
 	/**
 	 * Construct an instance with the default recoverer which simply logs the record after
@@ -140,20 +144,29 @@ public DefaultAfterRollbackProcessor(@Nullable
 	public DefaultAfterRollbackProcessor(@Nullable BiConsumer<ConsumerRecord<?, ?>, Exception> recoverer,
 			BackOff backOff, @Nullable BackOffHandler backOffHandler, @Nullable KafkaOperations<?, ?> kafkaOperations,
 			boolean commitRecovered) {
-
-		super(recoverer, backOff, backOffHandler, null);
+		super(recoverer, backOff, backOffHandler, new CommonErrorHandler() { });
 		this.kafkaTemplate = kafkaOperations;
 		super.setCommitRecovered(commitRecovered);
 		checkConfig();
 		this.backOff = backOff;
+		this.recoverer = (crs, ex) -> {
+			if (!CollectionUtils.isEmpty(crs)) {
+				if (recoverer == null) {
+					this.logger.error(ex, () -> "Records discarded: " + recordsToString(crs));
+				}
+				else {
+					crs.spliterator().forEachRemaining(rec -> recoverer.accept(rec, ex));
+				}
+			}
+		};
 	}
 
 	private void checkConfig() {
 		Assert.isTrue(!isCommitRecovered() || this.kafkaTemplate != null,
 				"A KafkaOperations is required when 'commitRecovered' is true");
 	}
 
-	@SuppressWarnings({ "unchecked", "rawtypes", "deprecation" })
+	@SuppressWarnings({ "unchecked", "rawtypes"})
 	@Override
 	public void process(List<ConsumerRecord<K, V>> records, Consumer<K, V> consumer,
 			@Nullable MessageListenerContainer container, Exception exception, boolean recoverable, EOSMode eosMode) {
@@ -179,42 +192,52 @@ && 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) {
 
-		int index = handlerBatchListenerFailedException(exception, records);
-		if (index > -1) {
-			List<ConsumerRecord<?, ?>> toCommit = new ArrayList<>();
-			List<ConsumerRecord<?, ?>> remaining = new ArrayList<>();
-			for (ConsumerRecord<?, ?> datum : records) {
-				if (index-- > 0) {
-					toCommit.add(datum);
-				}
-				else {
-					remaining.add(datum);
-				}
+		long nextBackOff = ListenerUtils.nextBackOff(this.backOff, this.backOffs);
+		if (nextBackOff != BackOffExecution.STOP) {
+			SeekUtils.doSeeks((List) recordList, consumer, exception, recoverable,
+					getFailureTracker()::recovered, container, this.logger);
+			try {
+				ListenerUtils.stoppableSleep(container, nextBackOff);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+			}
+			return;
+		}
+
+		int indexArg = handlerBatchListenerFailedException(exception, records);
+		int index = indexArg > -1 ? indexArg + 1 : recordList.size();
+		List<ConsumerRecord<?, ?>> toRecover = new ArrayList<>();
+		List<ConsumerRecord<?, ?>> remaining = new ArrayList<>();
+		for (ConsumerRecord<?, ?> datum : records) {
+			if (index-- > 0) {
+				toRecover.add(datum);
+			}
+			else {
+				remaining.add(datum);
 			}
+		}
+
+		try {
+			this.recoverer.accept(toRecover, exception);
+			SeekUtils.doSeeks(remaining, consumer, exception, recoverable,
+					getFailureTracker()::recovered, container, this.logger);
 			Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
-			toCommit.forEach(rec -> offsets.compute(new TopicPartition(rec.topic(), rec.partition()),
+			toRecover.forEach(rec -> offsets.compute(new TopicPartition(rec.topic(), rec.partition()),
 					(key, val) -> ListenerUtils.createOffsetAndMetadata(container, rec.offset() + 1)));
 			if (offsets.size() > 0 && this.kafkaTemplate != null && this.kafkaTemplate.isTransactional()) {
 				this.kafkaTemplate.sendOffsetsToTransaction(offsets, consumer.groupMetadata());
 			}
-
-			// skip first record
-			if (SeekUtils.doSeeks(remaining, consumer, exception, recoverable,
-					getFailureTracker()::recovered, container, this.logger)
-						&& isCommitRecovered() && this.kafkaTemplate.isTransactional()) {
-				ConsumerRecord<?, ?> skipped = remaining.get(0);
-				this.kafkaTemplate.sendOffsetsToTransaction(
-						Collections.singletonMap(new TopicPartition(skipped.topic(), skipped.partition()),
-								createOffsetAndMetadata(container, skipped.offset() + 1)
-						), consumer.groupMetadata());
-			}
+			clearThreadState();
 		}
-		else {
-			process(recordList, consumer, container, exception, false, eosMode);
+		catch (Exception ex) {
+			process(recordList, consumer, container, exception, recoverable, eosMode);
+			logger.error(ex, () -> "Recoverer threw an exception; re-seeking batch");
 		}
 
 	}
@@ -238,4 +261,14 @@ private static OffsetAndMetadata createOffsetAndMetadata(@Nullable MessageListen
 		}
 		return ListenerUtils.createOffsetAndMetadata(container, offset);
 	}
+
+	private static String recordsToString(List<ConsumerRecord<?, ?>> records) {
+		StringBuffer sb = new StringBuffer();
+		records.spliterator().forEachRemaining(rec -> sb
+				.append(KafkaUtils.format(rec))
+				.append(','));
+		sb.deleteCharAt(sb.length() - 1);
+		return sb.toString();
+	}
+
 }

Diff for: spring-kafka/src/main/java/org/springframework/kafka/listener/FallbackBatchErrorHandler.java

@@ -181,7 +181,7 @@ public void handleBatch(Exception thrownException, ConsumerRecords<?, ?> data, C
 					.stream()
 					.collect(
 							Collectors.toMap(tp -> tp,
-									tp -> data.records(tp).get(0).offset(), (u, v) -> (long) v, LinkedHashMap::new))
+									tp -> data.records(tp).get(0).offset(), (u, v) -> v, LinkedHashMap::new))
 					.forEach(consumer::seek);
 
 			throw new KafkaException("Seek to current after exception", getLogLevel(), thrownException);

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

@@ -2183,7 +2183,7 @@ protected void doInTransactionWithoutResult(TransactionStatus status) {
 						afterRollbackProcessorToUse.processBatch(records,
 								Objects.requireNonNullElseGet(recordList, () -> createRecordList(records)),
 								ListenerConsumer.this.consumer,
-								KafkaMessageListenerContainer.this.thisOrParentContainer, e, true,
+								KafkaMessageListenerContainer.this.thisOrParentContainer, e, false,
 								ListenerConsumer.this.eosMode);
 					}
 
@@ -2201,7 +2201,7 @@ private void batchAfterRollback(final ConsumerRecords<K, V> records,
 			try {
 				afterRollbackProcessorToUse.processBatch(records,
 						Objects.requireNonNullElseGet(recordList, () -> createRecordList(records)), this.consumer,
-						KafkaMessageListenerContainer.this.thisOrParentContainer, rollbackException, true,
+						KafkaMessageListenerContainer.this.thisOrParentContainer, rollbackException, false,
 						this.eosMode);
 			}
 			catch (KafkaException ke) {

Diff for: 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.

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

@@ -130,14 +130,18 @@ void testBatchBackOff() {
 		Consumer<String, String> consumer = mock(Consumer.class);
 		given(consumer.groupMetadata()).willReturn(new ConsumerGroupMetadata("foo"));
 		MessageListenerContainer container = mock(MessageListenerContainer.class);
+		ContainerProperties containerProperties = mock(ContainerProperties.class);
+		given(container.getContainerProperties()).willReturn(containerProperties);
 		given(container.isRunning()).willReturn(true);
-		processor.processBatch(consumerRecords, records, consumer, container, illegalState, true, EOSMode.V2);
-		processor.processBatch(consumerRecords, records, consumer, container, illegalState, true, 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.processBatch(consumerRecords, records, consumer, container, illegalState, true, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, illegalState, false, EOSMode.V2);
 		verify(backOff, times(3)).start();
+		processor.clearThreadState();
+		processor.processBatch(consumerRecords, records, consumer, container, illegalState, false, EOSMode.V2);
+		verify(backOff, times(4)).start();
 
 	}
 
@@ -173,10 +177,10 @@ void testBatchListenerFailedException() {
 		given(container.getContainerProperties()).willReturn(containerProperties);
 		given(container.isRunning()).willReturn(true);
 		BatchListenerFailedException batchFailed = new BatchListenerFailedException("", 2);
-		processor.processBatch(consumerRecords, records, consumer, container, batchFailed, true, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, batchFailed, false, EOSMode.V2);
 		verify(consumer).seek(new TopicPartition("foo", 1), 1);
 		BatchListenerFailedException batchFailed2 = new BatchListenerFailedException("", record1);
-		processor.processBatch(consumerRecords, records, consumer, container, batchFailed2, true, EOSMode.V2);
+		processor.processBatch(consumerRecords, records, consumer, container, batchFailed2, false, EOSMode.V2);
 		verify(consumer).seek(new TopicPartition("foo", 0), 0);
 		verify(consumer, times(2)).seek(new TopicPartition("foo", 1), 1);
 	}

Diff for: spring-kafka/src/test/java/org/springframework/kafka/listener/TransactionalContainerTests.java

@@ -37,6 +37,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -762,7 +763,7 @@ public void accept(ConsumerRecord<?, ?> record, Consumer<?, ?> consumer, Excepti
 		assertThat(headers.get(KafkaHeaders.DLT_EXCEPTION_STACKTRACE)).isNotNull();
 		assertThat(headers.get(KafkaHeaders.DLT_EXCEPTION_STACKTRACE, byte[].class))
 				.contains("fail for max failures".getBytes());
-		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_OFFSET, byte[].class)[3]).isEqualTo((byte) 0);
+		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_OFFSET, byte[].class)[7]).isEqualTo((byte) 0);
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_PARTITION, byte[].class)[3]).isEqualTo((byte) 0);
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_TIMESTAMP, byte[].class)).isNotNull();
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_TIMESTAMP_TYPE, byte[].class)).isNotNull();
@@ -862,11 +863,18 @@ public void accept(ConsumerRecord<?, ?> record, Consumer<?, ?> consumer, Excepti
 		container.stop();
 		Consumer<Integer, String> consumer = cf.createConsumer();
 		embeddedKafka.consumeFromAnEmbeddedTopic(consumer, topic8DLT);
-		ConsumerRecord<Integer, String> dltRecord = KafkaTestUtils.getSingleRecord(consumer, topic8DLT);
-		assertThat(dltRecord.value()).isEqualTo("bar");
+		ConsumerRecords<Integer, String> dltRecords = KafkaTestUtils.getRecords(consumer, Duration.ofSeconds(60));
+		List<ConsumerRecord<Integer, String>> recordList = new ArrayList<>();
+		for (ConsumerRecord<Integer, String> record : dltRecords) {
+			recordList.add(record);
+		}
+		ConsumerRecord<Integer, String> dltRecord0 = recordList.get(0);
+		assertThat(dltRecord0.value()).isEqualTo("foo");
+		ConsumerRecord<Integer, String> dltRecord1 = recordList.get(1);
+		assertThat(dltRecord1.value()).isEqualTo("bar");
 		DefaultKafkaHeaderMapper mapper = new DefaultKafkaHeaderMapper();
 		Map<String, Object> map = new HashMap<>();
-		mapper.toHeaders(dltRecord.headers(), map);
+		mapper.toHeaders(dltRecord1.headers(), map);
 		MessageHeaders headers = new MessageHeaders(map);
 		assertThat(new String(headers.get(KafkaHeaders.DLT_EXCEPTION_FQCN, byte[].class)))
 				.contains("ListenerExecutionFailedException");
@@ -877,7 +885,7 @@ public void accept(ConsumerRecord<?, ?> record, Consumer<?, ?> consumer, Excepti
 		assertThat(headers.get(KafkaHeaders.DLT_EXCEPTION_STACKTRACE)).isNotNull();
 		assertThat(headers.get(KafkaHeaders.DLT_EXCEPTION_STACKTRACE, byte[].class))
 				.contains("fail for max failures".getBytes());
-		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_OFFSET, byte[].class)[3]).isEqualTo((byte) 0);
+		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_OFFSET, byte[].class)[7]).isEqualTo((byte) 1);
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_PARTITION, byte[].class)[3]).isEqualTo((byte) 0);
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_TIMESTAMP, byte[].class)).isNotNull();
 		assertThat(headers.get(KafkaHeaders.DLT_ORIGINAL_TIMESTAMP_TYPE, byte[].class)).isNotNull();
@@ -892,10 +900,10 @@ public void accept(ConsumerRecord<?, ?> record, Consumer<?, ?> consumer, Excepti
 		assertThat(captor.getValue()).isInstanceOf(ListenerExecutionFailedException.class)
 				.extracting(ex -> ((ListenerExecutionFailedException) ex).getGroupId())
 				.isEqualTo(group);
-		verify(afterRollbackProcessor).clearThreadState();
-		verify(dlTemplate).send(any(ProducerRecord.class));
+		verify(afterRollbackProcessor, times(2)).clearThreadState();
+		verify(dlTemplate, times(2)).send(any(ProducerRecord.class));
 		verify(dlTemplate).sendOffsetsToTransaction(
-				eq(Collections.singletonMap(new TopicPartition(topic8, 0), new OffsetAndMetadata(1L))),
+				eq(Collections.singletonMap(new TopicPartition(topic8, 0), new OffsetAndMetadata(2L))),
 				any(ConsumerGroupMetadata.class));
 		logger.info("Stop testBatchListenerMaxFailures");
 	}