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;
@@ -106,6 +108,8 @@
 import org.springframework.kafka.listener.ContainerProperties.AssignmentCommitOption;
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
 import org.springframework.kafka.listener.adapter.AsyncRepliesAware;
+import org.springframework.kafka.listener.adapter.KafkaBackoffAwareMessageListenerAdapter;
+import org.springframework.kafka.listener.adapter.RecordMessagingMessageListenerAdapter;
 import org.springframework.kafka.support.Acknowledgment;
 import org.springframework.kafka.support.KafkaHeaders;
 import org.springframework.kafka.support.KafkaUtils;
@@ -840,6 +844,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 +901,22 @@ else if (listener instanceof MessageListener) {
 				this.wantsFullRecords = false;
 				this.pollThreadStateProcessor = setUpPollProcessor(false);
 				this.observationEnabled = this.containerProperties.isObservationEnabled();
+
+				if (!AopUtils.isAopProxy(this.genericListener) &&
+					this.genericListener instanceof KafkaBackoffAwareMessageListenerAdapter<?, ?>) {
+					KafkaBackoffAwareMessageListenerAdapter<K, V> genListener =
+							(KafkaBackoffAwareMessageListenerAdapter<K, V>) this.genericListener;
+					if (genListener.getDelegate() instanceof RecordMessagingMessageListenerAdapter<K, V>) {
+
+						RecordMessagingMessageListenerAdapter<K, V> recordAdapterListener =
+								(RecordMessagingMessageListenerAdapter<K, V>) genListener.getDelegate();
+
+						BiConsumer<ConsumerRecord<K, V>, RuntimeException> callbackForAsyncFailure =
+								(cRecord, ex) -> this.failedRecords.addLast(new FailedRecordTuple<>(cRecord, ex));
+						recordAdapterListener.setCallbackForAsyncFailure(callbackForAsyncFailure);
+					}
+
+				}
 			}
 			else {
 				throw new IllegalArgumentException("Listener must be one of 'MessageListener', "
@@ -1294,6 +1316,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 +1466,29 @@ protected void pollAndInvoke() {
 			}
 		}
 
+		protected void handleAsyncFailure() {
+			List<FailedRecordTuple<K, V>> copyFailedRecords = new ArrayList<>();
+			while (!this.failedRecords.isEmpty()) {
+				FailedRecordTuple<K, V> 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.
+			for (FailedRecordTuple<K, V> copyFailedRecord : copyFailedRecords) {
+				try {
+					invokeErrorHandlerBySingleRecord(copyFailedRecord);
+				}
+				catch (Exception e) {
+					this.logger.warn(() ->
+								"Async failed record failed to complete, thus skip it. record :"
+								+ copyFailedRecord.toString()
+								+ ", Exception : "
+								+ e.getMessage());
+				}
+			}
+		}
+
 		private void doProcessCommits() {
 			if (!this.autoCommit && !this.isRecordAck) {
 				try {
@@ -2827,6 +2881,44 @@ private void doInvokeOnMessage(final ConsumerRecord<K, V> recordArg) {
 			}
 		}
 
+		private void invokeErrorHandlerBySingleRecord(final FailedRecordTuple<K, V> failedRecordTuple) {
+			final ConsumerRecord<K, V> cRecord = failedRecordTuple.record;
+			RuntimeException rte = failedRecordTuple.ex;
+			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 +4005,6 @@ private static class StopAfterFenceException extends KafkaException {
 
 	}
 
+	private 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.

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.
@@ -143,4 +143,5 @@ public void onMessage(ConsumerRecord<K, V> data, Acknowledgment acknowledgment)
 	public void onMessage(ConsumerRecord<K, V> data, Consumer<?, ?> consumer) {
 		onMessage(data, null, consumer);
 	}
+
 }

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,9 @@ public abstract class MessagingMessageListenerAdapter<K, V> implements ConsumerS
 
 	private String correlationHeaderName = KafkaHeaders.CORRELATION_ID;
 
+	@Nullable
+	private BiConsumer<ConsumerRecord<K, V>, RuntimeException> callbackForAsyncFailure;
+
 	/**
 	 * Create an instance with the provided bean and method.
 	 * @param bean the bean.
@@ -665,16 +671,27 @@ 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) &&
+				Objects.nonNull(this.callbackForAsyncFailure)) {
+				ConsumerRecord<K, V> record = (ConsumerRecord<K, V>) request;
+				this.callbackForAsyncFailure.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 +886,10 @@ private boolean rawByParameterIsType(Type parameterType, Type type) {
 		return parameterType instanceof ParameterizedType pType && pType.getRawType().equals(type);
 	}
 
+	public void setCallbackForAsyncFailure(BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback) {
+		this.callbackForAsyncFailure = asyncRetryCallback;
+	}
+
 	/**
 	 * Root object for reply expression evaluation.
 	 * @param request the request.