Remove RetryListener from KafkaInbounds

* KafkaMessageDrivenChannelAdapter's ATTRIBUTES_HOLDER should be isolated.
In order to achieve custom retry in batch mode, we may to use a RetryTemplate in listener itself. But if the RetryTemplate is shared with another KafkaMessageDrivenChannelAdapter, batch mode's ATTRIBUTES_HOLDER might be over-written by another KafkaMessageDrivenChannelAdapter's IntegrationRecordMessageListener.
The situation is like shown below.
- There is only one RetryTemple bean in the application.
- There are two KafkaMessageDrivenChannelAdapters(A,B) in the application.
- A KafkaMessageDrivenChannelAdapter is batch mode and utilizing the retryTemplate in the listener.
- B KafkaMessageDrivenChannelAdapter is record mode and using the retryTemplate itself.
- (B KafkaMessageDrivenChannelAdapter's recordListener is registered in the retryTemplate.)
- When A retry is attempted in the listener, it will trigger B KafkaMessageDrivenChannelAdapter's recordListener.
- B KafkaMessageDrivenChannelAdapter's recordListener will overwrite A KafkaMessageDrivenChannelAdapter's ATTRIBUTES_HOLDER.

* should not mutate an externally provided RetryTemplate

* should not expose KafkaInboundEndpoint's methods outside the class and fix a checkstyle error.

* removing unused code and polishing

* restore retry around batch and bring back KafkaInboundEndpoint to endpoints.

* remove retry logic in batch mode and move ATTRIBUTES_HOLDER into KafkaInboundEndpoint

* remove generic type parameters of KafkaInboundEndpoint

* fix style error

---------

Co-authored-by: kurt <[email protected]>

Author: kurthong

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaInboundEndpoint.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2023 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.
@@ -17,7 +17,9 @@
 package org.springframework.integration.kafka.inbound;
 
 import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 
+import org.springframework.core.AttributeAccessor;
 import org.springframework.kafka.KafkaException;
 import org.springframework.kafka.support.Acknowledgment;
 import org.springframework.retry.RecoveryCallback;
@@ -50,6 +52,8 @@ public interface KafkaInboundEndpoint {
 	 */
 	String CONTEXT_RECORD = "record";
 
+	ThreadLocal<AttributeAccessor> ATTRIBUTES_HOLDER = new ThreadLocal<>();
+
 	/**
 	 * Execute the runnable with the retry template and recovery callback.
 	 * @param template the template.
@@ -59,21 +63,27 @@ public interface KafkaInboundEndpoint {
 	 * @param consumer the consumer.
 	 * @param runnable the runnable.
 	 */
-	default void doWithRetry(RetryTemplate template, RecoveryCallback<?> callback, Object data,
+	default void doWithRetry(RetryTemplate template, RecoveryCallback<?> callback, ConsumerRecord<?, ?> record,
 			Acknowledgment acknowledgment, Consumer<?, ?> consumer, Runnable runnable) {
 
 		try {
 			template.execute(context -> {
-				context.setAttribute(CONTEXT_RECORD, data);
-				context.setAttribute(CONTEXT_ACKNOWLEDGMENT, acknowledgment);
-				context.setAttribute(CONTEXT_CONSUMER, consumer);
+				if (context.getRetryCount() == 0) {
+					context.setAttribute(CONTEXT_RECORD, record);
+					context.setAttribute(CONTEXT_ACKNOWLEDGMENT, acknowledgment);
+					context.setAttribute(CONTEXT_CONSUMER, consumer);
+					ATTRIBUTES_HOLDER.set(context);
+				}
 				runnable.run();
 				return null;
 			}, callback);
 		}
 		catch (Exception ex) {
 			throw new KafkaException("Failed to execute runnable", ex);
 		}
+		finally {
+			ATTRIBUTES_HOLDER.remove();
+		}
 	}
 
 }

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaInboundGateway.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2022 the original author or authors.
+ * Copyright 2018-2023 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.
@@ -55,9 +55,7 @@
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.retry.RecoveryCallback;
-import org.springframework.retry.RetryCallback;
 import org.springframework.retry.RetryContext;
-import org.springframework.retry.RetryListener;
 import org.springframework.retry.support.RetryTemplate;
 import org.springframework.util.Assert;
 
@@ -78,8 +76,6 @@
 public class KafkaInboundGateway<K, V, R> extends MessagingGatewaySupport
 		implements KafkaInboundEndpoint, Pausable, OrderlyShutdownCapable {
 
-	private static final ThreadLocal<AttributeAccessor> ATTRIBUTES_HOLDER = new ThreadLocal<>();
-
 	private final IntegrationRecordMessageListener listener = new IntegrationRecordMessageListener();
 
 	private final AbstractMessageListenerContainer<K, V> messageListenerContainer;
@@ -195,7 +191,6 @@ public void setBindSourceRecord(boolean bindSourceRecord) {
 	protected void onInit() {
 		super.onInit();
 		if (this.retryTemplate != null) {
-			this.retryTemplate.registerListener(this.listener);
 			MessageChannel errorChannel = getErrorChannel();
 			if (this.recoveryCallback != null && errorChannel != null) {
 				this.recoveryCallback = new ErrorMessageSendingRecoverer(errorChannel, getErrorMessageStrategy());
@@ -261,7 +256,7 @@ public int afterShutdown() {
 	private void setAttributesIfNecessary(Object record, @Nullable Message<?> message, boolean conversionError) {
 		boolean needHolder = ATTRIBUTES_HOLDER.get() == null
 				&& (getErrorChannel() != null && (this.retryTemplate == null || conversionError));
-		boolean needAttributes = needHolder | this.retryTemplate != null;
+		boolean needAttributes = needHolder || this.retryTemplate != null;
 		if (needHolder) {
 			ATTRIBUTES_HOLDER.set(ErrorMessageUtils.getAttributeAccessor(null, null));
 		}
@@ -285,8 +280,7 @@ protected AttributeAccessor getErrorMessageAttributes(Message<?> message) {
 		}
 	}
 
-	private class IntegrationRecordMessageListener extends RecordMessagingMessageListenerAdapter<K, V>
-			implements RetryListener {
+	private class IntegrationRecordMessageListener extends RecordMessagingMessageListenerAdapter<K, V> {
 
 		IntegrationRecordMessageListener() {
 			super(null, null); // NOSONAR - out of use
@@ -332,9 +326,8 @@ private void sendAndReceive(ConsumerRecord<K, V> record, Message<?> message, Ack
 			RetryTemplate template = KafkaInboundGateway.this.retryTemplate;
 			if (template != null) {
 				doWithRetry(template, KafkaInboundGateway.this.recoveryCallback, record, acknowledgment, consumer,
-						() -> {
-							doSendAndReceive(enhanceHeadersAndSaveAttributes(message, record));
-						});
+						() -> doSendAndReceive(enhanceHeadersAndSaveAttributes(message, record))
+				);
 			}
 			else {
 				doSendAndReceive(enhanceHeadersAndSaveAttributes(message, record));
@@ -427,27 +420,6 @@ private Message<?> enhanceReply(Message<?> message, Message<?> reply) {
 			return reply;
 		}
 
-		@Override
-		public <T, E extends Throwable> boolean open(RetryContext context, RetryCallback<T, E> callback) {
-			if (KafkaInboundGateway.this.retryTemplate != null) {
-				ATTRIBUTES_HOLDER.set(context);
-			}
-			return true;
-		}
-
-		@Override
-		public <T, E extends Throwable> void close(RetryContext context, RetryCallback<T, E> callback,
-				Throwable throwable) {
-
-			ATTRIBUTES_HOLDER.remove();
-		}
-
-		@Override
-		public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
-				Throwable throwable) {
-			// Empty
-		}
-
 	}
 
 }

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageDrivenChannelAdapter.java

@@ -63,9 +63,7 @@
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.retry.RecoveryCallback;
-import org.springframework.retry.RetryCallback;
 import org.springframework.retry.RetryContext;
-import org.springframework.retry.RetryListener;
 import org.springframework.retry.support.RetryTemplate;
 import org.springframework.util.Assert;
 
@@ -85,8 +83,6 @@
 public class KafkaMessageDrivenChannelAdapter<K, V> extends MessageProducerSupport
 		implements KafkaInboundEndpoint, OrderlyShutdownCapable, Pausable {
 
-	private static final ThreadLocal<AttributeAccessor> ATTRIBUTES_HOLDER = new ThreadLocal<>();
-
 	private final AbstractMessageListenerContainer<K, V> messageListenerContainer;
 
 	private final IntegrationRecordMessageListener recordListener = new IntegrationRecordMessageListener();
@@ -298,7 +294,6 @@ protected void onInit() {
 				if (this.recoveryCallback != null && errorChannel != null) {
 					this.recoveryCallback = new ErrorMessageSendingRecoverer(errorChannel, getErrorMessageStrategy());
 				}
-				this.retryTemplate.registerListener(this.recordListener);
 			}
 			if (!doFilterInRetry && this.recordFilterStrategy != null) {
 				listener = new FilteringMessageListenerAdapter<>(listener, this.recordFilterStrategy,
@@ -366,7 +361,7 @@ public int afterShutdown() {
 	private void setAttributesIfNecessary(Object record, @Nullable Message<?> message, boolean conversionError) {
 		boolean needHolder = ATTRIBUTES_HOLDER.get() == null
 				&& (getErrorChannel() != null && (this.retryTemplate == null || conversionError));
-		boolean needAttributes = needHolder | this.retryTemplate != null;
+		boolean needAttributes = needHolder || this.retryTemplate != null;
 		if (needHolder) {
 			ATTRIBUTES_HOLDER.set(ErrorMessageUtils.getAttributeAccessor(null, null));
 		}
@@ -424,8 +419,7 @@ public enum ListenerMode {
 		batch
 	}
 
-	private class IntegrationRecordMessageListener extends RecordMessagingMessageListenerAdapter<K, V>
-			implements RetryListener {
+	private class IntegrationRecordMessageListener extends RecordMessagingMessageListenerAdapter<K, V> {
 
 		IntegrationRecordMessageListener() {
 			super(null, null); // NOSONAR - out of use
@@ -511,34 +505,12 @@ else if (KafkaMessageDrivenChannelAdapter.this.containerDeliveryAttemptPresent)
 			return messageToReturn;
 		}
 
-		@Override
-		public <T, E extends Throwable> boolean open(RetryContext context, RetryCallback<T, E> callback) {
-			if (KafkaMessageDrivenChannelAdapter.this.retryTemplate != null) {
-				ATTRIBUTES_HOLDER.set(context);
-			}
-			return true;
-		}
-
-		@Override
-		public <T, E extends Throwable> void close(RetryContext context, RetryCallback<T, E> callback,
-				Throwable throwable) {
-
-			ATTRIBUTES_HOLDER.remove();
-		}
-
-		@Override
-		public <T, E extends Throwable> void onError(RetryContext context, RetryCallback<T, E> callback,
-				Throwable throwable) {
-			// Empty
-		}
-
 	}
 
-	private class IntegrationBatchMessageListener extends BatchMessagingMessageListenerAdapter<K, V>
-			implements RetryListener {
+	private class IntegrationBatchMessageListener extends BatchMessagingMessageListenerAdapter<K, V> {
 
 		IntegrationBatchMessageListener() {
-			super(null, null); // NOSONAR - out if use
+			super(null, null); // NOSONAR - out of use
 		}
 
 		@Override
@@ -557,33 +529,10 @@ public void onMessage(List<ConsumerRecord<K, V>> records, Acknowledgment acknowl
 				message = toMessage(records, acknowledgment, consumer);
 			}
 			if (message != null) {
-				RetryTemplate template = KafkaMessageDrivenChannelAdapter.this.retryTemplate;
-				if (template != null) {
-					doWIthRetry(records, acknowledgment, consumer, message, template);
-				}
-				else {
-					sendMessageIfAny(message, records);
-				}
+				sendMessageIfAny(message, records);
 			}
 		}
 
-		private void doWIthRetry(List<ConsumerRecord<K, V>> records, Acknowledgment acknowledgment,
-				Consumer<?, ?> consumer, Message<?> message, RetryTemplate template) {
-
-			doWithRetry(template, KafkaMessageDrivenChannelAdapter.this.recoveryCallback, records, acknowledgment,
-					consumer, () -> {
-						if (KafkaMessageDrivenChannelAdapter.this.filterInRetry) {
-							List<ConsumerRecord<K, V>> filtered =
-									KafkaMessageDrivenChannelAdapter.this.recordFilterStrategy.filterBatch(records);
-							Message<?> toSend = message;
-							if (filtered.size() != records.size()) {
-								toSend = toMessage(filtered, acknowledgment, consumer);
-							}
-							sendMessageIfAny(toSend, filtered);
-						}
-					});
-		}
-
 		@Nullable
 		private Message<?> toMessage(List<ConsumerRecord<K, V>> records, Acknowledgment acknowledgment,
 				Consumer<?, ?> consumer) {
@@ -607,21 +556,6 @@ private Message<?> toMessage(List<ConsumerRecord<K, V>> records, Acknowledgment
 			return message;
 		}
 
-		@Override
-		public <T, E extends Throwable> boolean open(RetryContext context, RetryCallback<T, E> callback) {
-			if (KafkaMessageDrivenChannelAdapter.this.retryTemplate != null) {
-				ATTRIBUTES_HOLDER.set(context);
-			}
-			return true;
-		}
-
-		@Override
-		public <T, E extends Throwable> void close(RetryContext context, RetryCallback<T, E> callback,
-				Throwable throwable) {
-
-			ATTRIBUTES_HOLDER.remove();
-		}
-
 	}
 
 }