GH-3276: Support async retry with @RetryableTopic

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

@@ -38,6 +38,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledFuture;
@@ -73,6 +74,7 @@
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.internals.RecordHeader;
 
+import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.context.ApplicationContext;
@@ -840,6 +842,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private volatile long lastPoll = System.currentTimeMillis();
 
+		private final ConcurrentLinkedDeque<FailedRecordTuple<K, V>> failedRecords = new ConcurrentLinkedDeque();
+
 		@SuppressWarnings(UNCHECKED)
 		ListenerConsumer(GenericMessageListener<?> listener, ListenerType listenerType,
 				ObservationRegistry observationRegistry) {
@@ -895,6 +899,15 @@ else if (listener instanceof MessageListener) {
 				this.wantsFullRecords = false;
 				this.pollThreadStateProcessor = setUpPollProcessor(false);
 				this.observationEnabled = this.containerProperties.isObservationEnabled();
+
+				if (!AopUtils.isAopProxy(listener)) {
+					BiConsumer<ConsumerRecord<K, V>, RuntimeException> callbackForAsyncFailureQueue =
+							(cRecord, ex) -> {
+								FailedRecordTuple<K, V> failedRecord = new FailedRecordTuple<>(cRecord, ex);
+								this.failedRecords.addLast(failedRecord);
+							};
+					this.listener.setCallbackForAsyncFailure(callbackForAsyncFailureQueue);
+				}
 			}
 			else {
 				throw new IllegalArgumentException("Listener must be one of 'MessageListener', "
@@ -1294,6 +1307,15 @@ public void run() {
 			boolean failedAuthRetry = false;
 			this.lastReceive = System.currentTimeMillis();
 			while (isRunning()) {
+
+				try {
+					handleAsyncFailure();
+				}
+				catch (Exception e) {
+					ListenerConsumer.this.logger.error(
+							"Failed to process async retry messages. skip this time, try it again next loop.");
+				}
+
 				try {
 					pollAndInvoke();
 					if (failedAuthRetry) {
@@ -1435,6 +1457,21 @@ protected void pollAndInvoke() {
 			}
 		}
 
+		protected void handleAsyncFailure() {
+			List<FailedRecordTuple> copyFailedRecords = new ArrayList<>();
+			while (!this.failedRecords.isEmpty()) {
+				FailedRecordTuple failedRecordTuple = this.failedRecords.pollFirst();
+				copyFailedRecords.add(failedRecordTuple);
+			}
+
+			// If any copied and failed record fails to complete due to an unexpected error,
+			// We will give up on retrying with the remaining copied and failed Records.
+			if (!copyFailedRecords.isEmpty()) {
+				copyFailedRecords.forEach(failedRecordTuple ->
+											this.invokeErrorHandlerBySingleRecord(failedRecordTuple.record(), failedRecordTuple.ex()));
+			}
+		}
+
 		private void doProcessCommits() {
 			if (!this.autoCommit && !this.isRecordAck) {
 				try {
@@ -2827,6 +2864,42 @@ private void doInvokeOnMessage(final ConsumerRecord<K, V> recordArg) {
 			}
 		}
 
+		private void invokeErrorHandlerBySingleRecord(final ConsumerRecord<K, V> cRecord, RuntimeException rte) {
+			if (this.commonErrorHandler.seeksAfterHandling() || rte instanceof CommitFailedException) {
+				try {
+					if (this.producer == null) {
+						processCommits();
+					}
+				}
+				catch (Exception ex) { // NO SONAR
+					this.logger.error(ex, "Failed to commit before handling error");
+				}
+				List<ConsumerRecord<?, ?>> records = new ArrayList<>();
+				records.add(cRecord);
+				this.commonErrorHandler.handleRemaining(rte, records, this.consumer,
+														KafkaMessageListenerContainer.this.thisOrParentContainer);
+			}
+			else {
+				boolean handled = false;
+				try {
+					handled = this.commonErrorHandler.handleOne(rte, cRecord, this.consumer,
+																KafkaMessageListenerContainer.this.thisOrParentContainer);
+				}
+				catch (Exception ex) {
+					this.logger.error(ex, "ErrorHandler threw unexpected exception");
+				}
+				Map<TopicPartition, List<ConsumerRecord<K, V>>> records = new LinkedHashMap<>();
+				if (!handled) {
+					records.computeIfAbsent(new TopicPartition(cRecord.topic(), cRecord.partition()),
+											tp -> new ArrayList<>()).add(cRecord);
+				}
+				if (!records.isEmpty()) {
+					this.remainingRecords = new ConsumerRecords<>(records);
+					this.pauseForPending = true;
+				}
+			}
+		}
+
 		private void invokeErrorHandler(final ConsumerRecord<K, V> cRecord,
 				Iterator<ConsumerRecord<K, V>> iterator, RuntimeException rte) {
 
@@ -3913,4 +3986,6 @@ private static class StopAfterFenceException extends KafkaException {
 
 	}
 
+	record FailedRecordTuple<K, V>(ConsumerRecord<K, V> record, RuntimeException ex) { };
+
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2019 the original author or authors.
+ * Copyright 2015-2024 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.
@@ -16,6 +16,8 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.function.BiConsumer;
+
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 /**
@@ -26,8 +28,13 @@
  *
  * @author Marius Bogoevici
  * @author Gary Russell
+ * @author Sanghyeok An
  */
 @FunctionalInterface
 public interface MessageListener<K, V> extends GenericMessageListener<ConsumerRecord<K, V>> {
 
+	default void setCallbackForAsyncFailure(
+			BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback) {
+		//
+	}
 }

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/KafkaBackoffAwareMessageListenerAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2023 the original author or authors.
+ * Copyright 2018-2024 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 java.time.Clock;
 import java.time.Instant;
 import java.util.Optional;
+import java.util.function.BiConsumer;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -43,6 +44,7 @@
  * @param <K> the record key type.
  * @param <V> the record value type.
  * @author Tomaz Fernandes
+ * @author Sanghyeok An
  * @since 2.7
  *
  */
@@ -143,4 +145,9 @@ public void onMessage(ConsumerRecord<K, V> data, Acknowledgment acknowledgment)
 	public void onMessage(ConsumerRecord<K, V> data, Consumer<?, ?> consumer) {
 		onMessage(data, null, consumer);
 	}
+
+	@Override
+	public void setCallbackForAsyncFailure(BiConsumer<ConsumerRecord<K, V>, RuntimeException> callbackForAsyncFailureQueue) {
+		this.delegate.setCallbackForAsyncFailure(callbackForAsyncFailureQueue);
+	}
 }

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/MessagingMessageListenerAdapter.java

@@ -28,6 +28,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 import org.apache.commons.logging.LogFactory;
@@ -91,6 +93,7 @@
  * @author Wang ZhiYang
  * @author Huijin Hong
  * @author Soby Chacko
+ * @author Sanghyeok An
  */
 public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerSeekAware, AsyncRepliesAware {
 
@@ -152,6 +155,8 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private String correlationHeaderName = KafkaHeaders.CORRELATION_ID;
 
+	private BiConsumer<ConsumerRecord<K, V>, RuntimeException> callbackForAsyncFailureQueue;
+
 	/**
 	 * Create an instance with the provided bean and method.
 	 * @param bean the bean.
@@ -665,16 +670,26 @@ protected void asyncFailure(Object request, @Nullable Acknowledgment acknowledgm
 			Throwable t, Message<?> source) {
 
 		try {
+			Throwable cause = t instanceof CompletionException ? t.getCause() : t;
 			handleException(request, acknowledgment, consumer, source,
-					new ListenerExecutionFailedException(createMessagingErrorMessage(
-							"Async Fail", source.getPayload()), t));
+							new ListenerExecutionFailedException(createMessagingErrorMessage(
+									"Async Fail", source.getPayload()), cause));
 		}
 		catch (Throwable ex) {
 			this.logger.error(t, () -> "Future, Mono, or suspend function was completed with an exception for " + source);
 			acknowledge(acknowledgment);
+			if (canAsyncRetry(request, ex)) {
+				ConsumerRecord<K, V> record = (ConsumerRecord<K, V>) request;
+				this.callbackForAsyncFailureQueue.accept(record, (RuntimeException) ex);
+			}
 		}
 	}
 
+	private boolean canAsyncRetry(Object request, Throwable exception) {
+		// The async retry with @RetryableTopic only support in case of SingleRecord Listener.
+		return request instanceof ConsumerRecord && exception instanceof RuntimeException;
+	}
+
 	protected void handleException(Object records, @Nullable Acknowledgment acknowledgment, Consumer<?, ?> consumer,
 			Message<?> message, ListenerExecutionFailedException e) {
 
@@ -869,6 +884,10 @@ private boolean rawByParameterIsType(Type parameterType, Type type) {
 		return parameterType instanceof ParameterizedType pType && pType.getRawType().equals(type);
 	}
 
+	protected void setAsyncFailureCallback(BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback) {
+		this.callbackForAsyncFailureQueue = asyncRetryCallback;
+	}
+
 	/**
 	 * Root object for reply expression evaluation.
 	 * @param request the request.

spring-kafka/src/main/java/org/springframework/kafka/listener/adapter/RecordMessagingMessageListenerAdapter.java

@@ -17,6 +17,7 @@
 package org.springframework.kafka.listener.adapter;
 
 import java.lang.reflect.Method;
+import java.util.function.BiConsumer;
 
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -85,4 +86,9 @@ public void onMessage(ConsumerRecord<K, V> record, @Nullable Acknowledgment ackn
 		invoke(record, acknowledgment, consumer, message);
 	}
 
+	@Override
+	public void setCallbackForAsyncFailure(
+			BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback) {
+		setAsyncFailureCallback(asyncRetryCallback);
+	}
 }