GH-3276: Support async retry with @RetryableTopic

spring-kafka/src/main/java/org/springframework/kafka/core/FailedRecordTuple.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2016-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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * ...
+ * @author Sanghyeok An
+ *
+ * @since 3.3
+ */
+
+package org.springframework.kafka.core;
+
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+public record FailedRecordTuple<K, V>(
+		ConsumerRecord<K, V> record,
+		RuntimeException ex) {
+
+};

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;
@@ -82,6 +84,7 @@
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.kafka.KafkaException;
 import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.kafka.core.FailedRecordTuple;
 import org.springframework.kafka.core.KafkaAdmin;
 import org.springframework.kafka.core.KafkaResourceHolder;
 import org.springframework.kafka.event.ConsumerFailedToStartEvent;
@@ -840,6 +843,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 +900,12 @@ else if (listener instanceof MessageListener) {
 				this.wantsFullRecords = false;
 				this.pollThreadStateProcessor = setUpPollProcessor(false);
 				this.observationEnabled = this.containerProperties.isObservationEnabled();
+
+				if (!AopUtils.isAopProxy(listener)) {
+					final java.util.function.Consumer<FailedRecordTuple<K, V>> callbackForAsyncFailureQueue =
+							(fRecord) -> this.failedRecords.addLast(fRecord);
+					this.listener.setCallbackForAsyncFailureQueue(callbackForAsyncFailureQueue);
+				}
 			}
 			else {
 				throw new IllegalArgumentException("Listener must be one of 'MessageListener', "
@@ -1294,6 +1305,16 @@ public void run() {
 			boolean failedAuthRetry = false;
 			this.lastReceive = System.currentTimeMillis();
 			while (isRunning()) {
+
+				try {
+					handleAsyncFailure();
+				}
+				catch (Exception e) {
+					// TODO: Need to improve error handling.
+					// TODO: Need to determine how to handle a failed message.
+					this.logger.error("Failed to process re-try messages. ");
+				}
+
 				try {
 					pollAndInvoke();
 					if (failedAuthRetry) {
@@ -1435,6 +1456,19 @@ protected void pollAndInvoke() {
 			}
 		}
 
+		protected void handleAsyncFailure() {
+			List<FailedRecordTuple> copyFailedRecords = new ArrayList<>();
+			while (!this.failedRecords.isEmpty()) {
+				FailedRecordTuple failedRecordTuple = this.failedRecords.pollFirst();
+				copyFailedRecords.add(failedRecordTuple);
+			}
+
+			if (!copyFailedRecords.isEmpty()) {
+				copyFailedRecords.forEach(failedRecordTuple ->
+											this.invokeErrorHandlerBySingleRecord(failedRecordTuple.record(), failedRecordTuple.ex()));
+			}
+		}
+
 		private void doProcessCommits() {
 			if (!this.autoCommit && !this.isRecordAck) {
 				try {
@@ -2827,6 +2861,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) {
 

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,8 +16,12 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.function.Consumer;
+
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
+import org.springframework.kafka.core.FailedRecordTuple;
+
 /**
  * Listener for handling individual incoming Kafka messages.
  *
@@ -26,8 +30,12 @@
  *
  * @author Marius Bogoevici
  * @author Gary Russell
+ * @author Sanghyeok An
  */
 @FunctionalInterface
 public interface MessageListener<K, V> extends GenericMessageListener<ConsumerRecord<K, V>> {
 
+	default void setCallbackForAsyncFailureQueue(Consumer<FailedRecordTuple<K, V>> 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.
@@ -25,6 +25,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.TopicPartition;
 
+import org.springframework.kafka.core.FailedRecordTuple;
 import org.springframework.kafka.listener.AcknowledgingConsumerAwareMessageListener;
 import org.springframework.kafka.listener.KafkaBackoffException;
 import org.springframework.kafka.listener.KafkaConsumerBackoffManager;
@@ -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 setCallbackForAsyncFailureQueue(java.util.function.Consumer<FailedRecordTuple<K, V>> callbackForAsyncFailureQueue) {
+		this.delegate.setCallbackForAsyncFailureQueue(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;
@@ -45,6 +47,7 @@
 import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.expression.spel.support.StandardEvaluationContext;
 import org.springframework.expression.spel.support.StandardTypeConverter;
+import org.springframework.kafka.core.FailedRecordTuple;
 import org.springframework.kafka.core.KafkaTemplate;
 import org.springframework.kafka.listener.ConsumerSeekAware;
 import org.springframework.kafka.listener.KafkaListenerErrorHandler;
@@ -91,6 +94,7 @@
  * @author Wang ZhiYang
  * @author Huijin Hong
  * @author Soby Chacko
+ * @author Sanghyeok An
  */
 public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerSeekAware, AsyncRepliesAware {
 
@@ -152,6 +156,10 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private String correlationHeaderName = KafkaHeaders.CORRELATION_ID;
 
+	private BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback;
+
+	private java.util.function.Consumer<FailedRecordTuple<K, V>> callbackForAsyncFailureQueue;
+
 	/**
 	 * Create an instance with the provided bean and method.
 	 * @param bean the bean.
@@ -665,16 +673,35 @@ protected void asyncFailure(Object request, @Nullable Acknowledgment acknowledgm
 			Throwable t, Message<?> source) {
 
 		try {
-			handleException(request, acknowledgment, consumer, source,
-					new ListenerExecutionFailedException(createMessagingErrorMessage(
-							"Async Fail", source.getPayload()), t));
+			if (t instanceof CompletionException) {
+				// For CompletableFuture Object
+				handleException(request, acknowledgment, consumer, source,
+								new ListenerExecutionFailedException(createMessagingErrorMessage(
+										"Async Fail", source.getPayload()), t.getCause()));
+			}
+			else {
+				// For Mono Object.
+				handleException(request, acknowledgment, consumer, source,
+								new ListenerExecutionFailedException(createMessagingErrorMessage(
+										"Async Fail", source.getPayload()), t));
+			}
 		}
 		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;
+				FailedRecordTuple<K, V> failedRecordTuple = new FailedRecordTuple<>(record, (RuntimeException) ex);
+				this.callbackForAsyncFailureQueue.accept(failedRecordTuple);
+			}
 		}
 	}
 
+	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 +896,10 @@ private boolean rawByParameterIsType(Type parameterType, Type type) {
 		return parameterType instanceof ParameterizedType pType && pType.getRawType().equals(type);
 	}
 
+	public void putInAsyncFailureQueue(java.util.function.Consumer<FailedRecordTuple<K, V>> callbackForAsyncFailureQueue) {
+		this.callbackForAsyncFailureQueue = callbackForAsyncFailureQueue;
+	}
+
 	/**
 	 * Root object for reply expression evaluation.
 	 * @param request the request.

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

@@ -21,6 +21,7 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
+import org.springframework.kafka.core.FailedRecordTuple;
 import org.springframework.kafka.listener.AcknowledgingConsumerAwareMessageListener;
 import org.springframework.kafka.listener.KafkaListenerErrorHandler;
 import org.springframework.kafka.support.Acknowledgment;
@@ -85,4 +86,9 @@ public void onMessage(ConsumerRecord<K, V> record, @Nullable Acknowledgment ackn
 		invoke(record, acknowledgment, consumer, message);
 	}
 
+	@Override
+	public void setCallbackForAsyncFailureQueue(
+			java.util.function.Consumer<FailedRecordTuple<K, V>> asyncRetryCallback) {
+		putInAsyncFailureQueue(asyncRetryCallback);
+	}
 }