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;
@@ -105,7 +107,9 @@
 import org.springframework.kafka.listener.ContainerProperties.AckMode;
 import org.springframework.kafka.listener.ContainerProperties.AssignmentCommitOption;
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
+import org.springframework.kafka.listener.adapter.AbstractDelegatingMessageListenerAdapter;
 import org.springframework.kafka.listener.adapter.AsyncRepliesAware;
+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;
@@ -167,6 +171,7 @@
  * @author Mikael Carlstedt
  * @author Borahm Lee
  * @author Lokesh Alamuri
+ * @author Sanghyeok An
  */
 public class KafkaMessageListenerContainer<K, V> // NOSONAR line count
 		extends AbstractMessageListenerContainer<K, V> implements ConsumerPauseResumeEventPublisher {
@@ -840,6 +845,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 +902,15 @@ 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 AbstractDelegatingMessageListenerAdapter<?>) {
+					AbstractDelegatingMessageListenerAdapter<MessageListener<K, V>> genListener =
+							(AbstractDelegatingMessageListenerAdapter<MessageListener<K, V>>) this.genericListener;
+					if (genListener.getDelegate() instanceof RecordMessagingMessageListenerAdapter<K, V> adapterListener) {
+						adapterListener.setCallbackForAsyncFailure(this::callbackForAsyncFailure);
+					}
+				}
 			}
 			else {
 				throw new IllegalArgumentException("Listener must be one of 'MessageListener', "
@@ -1294,6 +1310,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 +1460,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 +2875,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) {
 
@@ -3299,6 +3385,10 @@ private Collection<ConsumerRecord<K, V>> getHighestOffsetRecords(ConsumerRecords
 					.values();
 		}
 
+		private void callbackForAsyncFailure(ConsumerRecord<K, V> cRecord, RuntimeException ex) {
+			this.failedRecords.addLast(new FailedRecordTuple<>(cRecord, ex));
+		}
+
 		@Override
 		public void seek(String topic, int partition, long offset) {
 			this.seeks.add(new TopicPartitionOffset(topic, partition, offset));
@@ -3913,4 +4003,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/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,18 @@ private boolean rawByParameterIsType(Type parameterType, Type type) {
 		return parameterType instanceof ParameterizedType pType && pType.getRawType().equals(type);
 	}
 
+	/**
+	 * Set the retry callback for failures of both {@link CompletableFuture} and {@link Mono}.
+	 * {@link MessagingMessageListenerAdapter#asyncFailure(Object, Acknowledgment, Consumer, Throwable, Message)}
+	 * will invoke {@link MessagingMessageListenerAdapter#callbackForAsyncFailure} when
+	 * {@link CompletableFuture} or {@link Mono} fails to complete.
+	 * @param asyncRetryCallback the callback for async retry.
+	 * @since 3.3
+	 */
+	public void setCallbackForAsyncFailure(BiConsumer<ConsumerRecord<K, V>, RuntimeException> asyncRetryCallback) {
+		this.callbackForAsyncFailure = asyncRetryCallback;
+	}
+
 	/**
 	 * Root object for reply expression evaluation.
 	 * @param request the request.