GH-3542: Adds the ability to add record interceptors instead of override them

Fixes: #3542
Issue link: #3542

What?
Change `RecordInterceptor` to `List<RecordInterceptor>` in
`MessageListenerContainer`.

Why?
To allow adding multiple `RecordInterceptor` instances instead of overriding the existing one.
Currently, only a single `RecordInterceptor` is supported.
Users may want to register multiple `RecordInterceptors`.
There are some workarounds, but they are not clean or ideal solutions.

By supporting `List<RecordInterceptor`>, users can add their own
interceptors via `setRecordInterceptor(...)`.

Signed-off-by: Sanghyeok An <[email protected]>

Author: chickenchickenlove

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

@@ -16,6 +16,7 @@
 
 package org.springframework.kafka.listener;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
@@ -114,7 +115,8 @@ public abstract class AbstractMessageListenerContainer<K, V>
 
 	private int topicCheckTimeout = DEFAULT_TOPIC_CHECK_TIMEOUT;
 
-	private @Nullable RecordInterceptor<K, V> recordInterceptor;
+	// TODO: HERE ASH
+	private List<RecordInterceptor<K, V>> recordInterceptors = new ArrayList<>();
 
 	private @Nullable BatchInterceptor<K, V> batchInterceptor;
 
@@ -460,8 +462,8 @@ public void setKafkaAdmin(KafkaAdmin kafkaAdmin) {
 		this.kafkaAdmin = kafkaAdmin;
 	}
 
-	protected @Nullable RecordInterceptor<K, V> getRecordInterceptor() {
-		return this.recordInterceptor;
+	protected List<RecordInterceptor<K, V>> getRecordInterceptors() {
+		return this.recordInterceptors;
 	}
 
 	/**
@@ -472,7 +474,9 @@ public void setKafkaAdmin(KafkaAdmin kafkaAdmin) {
 	 * @see #setInterceptBeforeTx(boolean)
 	 */
 	public void setRecordInterceptor(@Nullable RecordInterceptor<K, V> recordInterceptor) {
-		this.recordInterceptor = recordInterceptor;
+		if (recordInterceptor != null) {
+			this.recordInterceptors.add(recordInterceptor);
+		}
 	}
 
 	protected @Nullable BatchInterceptor<K, V> getBatchInterceptor() {

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

@@ -60,6 +60,7 @@
  * @author Tomaz Fernandes
  * @author Wang Zhiyang
  * @author Lokesh Alamuri
+ * @author Sanghyeok An
  */
 public class ConcurrentMessageListenerContainer<K, V> extends AbstractMessageListenerContainer<K, V> {
 
@@ -282,7 +283,9 @@ private void configureChildContainer(int index, KafkaMessageListenerContainer<K,
 		container.setClientIdSuffix(this.concurrency > 1 || this.alwaysClientIdSuffix ? "-" + index : "");
 		container.setCommonErrorHandler(getCommonErrorHandler());
 		container.setAfterRollbackProcessor(getAfterRollbackProcessor());
-		container.setRecordInterceptor(getRecordInterceptor());
+		for (RecordInterceptor<K, V> recordInterceptor : getRecordInterceptors()) {
+			container.setRecordInterceptor(recordInterceptor);
+		}
 		container.setBatchInterceptor(getBatchInterceptor());
 		container.setInterceptBeforeTx(isInterceptBeforeTx());
 		container.setListenerInfo(getListenerInfo());

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

@@ -714,17 +714,17 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final @Nullable Duration syncCommitTimeout;
 
-		private final @Nullable RecordInterceptor<K, V> recordInterceptor =
+		private final List<RecordInterceptor<K, V>> recordInterceptors =
 				!isInterceptBeforeTx() || this.transactionManager == null
-						? getRecordInterceptor()
-						: null;
+						? getRecordInterceptors()
+						: new ArrayList<>();
 
-		private final @Nullable RecordInterceptor<K, V> earlyRecordInterceptor =
+		private final List<RecordInterceptor<K, V>> earlyRecordInterceptors =
 				isInterceptBeforeTx() && this.transactionManager != null
-						? getRecordInterceptor()
-						: null;
+						? getRecordInterceptors()
+						: new ArrayList<>();
 
-		private final @Nullable RecordInterceptor<K, V> commonRecordInterceptor = getRecordInterceptor();
+		private final List<RecordInterceptor<K, V>> commonRecordInterceptors = getRecordInterceptors();
 
 		private final @Nullable BatchInterceptor<K, V> batchInterceptor =
 				!isInterceptBeforeTx() || this.transactionManager == null
@@ -738,7 +738,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final @Nullable BatchInterceptor<K, V> commonBatchInterceptor = getBatchInterceptor();
 
-		private final @Nullable ThreadStateProcessor pollThreadStateProcessor;
+		private final List<ThreadStateProcessor> pollThreadStateProcessor;
 
 		private final ConsumerSeekCallback seekCallback = new InitialOrIdleSeekCallback();
 
@@ -1040,9 +1040,20 @@ private void obtainClusterId() {
 			}
 		}
 
-		@Nullable
-		private ThreadStateProcessor setUpPollProcessor(boolean batch) {
-			return batch ? this.commonBatchInterceptor : this.commonRecordInterceptor;
+		private List<ThreadStateProcessor> setUpPollProcessor(boolean batch) {
+			if (batch) {
+				if (this.commonBatchInterceptor != null) {
+					List<ThreadStateProcessor> threadStateProcessors = new ArrayList<>();
+					threadStateProcessors.add(this.commonBatchInterceptor);
+					return threadStateProcessors;
+				}
+				else {
+					return new ArrayList<>();
+				}
+			}
+			else {
+				return new ArrayList<>(this.commonRecordInterceptors);
+			}
 		}
 
 		@Nullable
@@ -1548,9 +1559,7 @@ private void invokeIfHaveRecords(@Nullable ConsumerRecords<K, V> records) {
 		}
 
 		private void clearThreadState() {
-			if (this.pollThreadStateProcessor != null) {
-				this.pollThreadStateProcessor.clearThreadState(this.consumer);
-			}
+			this.pollThreadStateProcessor.forEach(threadStateProcessor -> threadStateProcessor.clearThreadState(this.consumer));
 		}
 
 		private void checkIdlePartitions() {
@@ -1708,9 +1717,7 @@ private ConsumerRecords<K, V> pollConsumer() {
 		}
 
 		private void beforePoll() {
-			if (this.pollThreadStateProcessor != null) {
-				this.pollThreadStateProcessor.setupThreadState(this.consumer);
-			}
+			this.pollThreadStateProcessor.forEach(threadStateProcessor -> threadStateProcessor.setupThreadState(this.consumer));
 		}
 
 		private synchronized void captureOffsets(ConsumerRecords<K, V> records) {
@@ -2548,9 +2555,7 @@ private void invokeRecordListenerInTx(final ConsumerRecords<K, V> records) {
 					this.logger.error(ex, "Transaction rolled back");
 					recordAfterRollback(iterator, cRecord, ex);
 				}
-				if (this.commonRecordInterceptor != null) {
-					this.commonRecordInterceptor.afterRecord(cRecord, this.consumer);
-				}
+				this.commonRecordInterceptors.forEach(interceptor -> interceptor.afterRecord(cRecord, this.consumer));
 				if (this.nackSleepDurationMillis >= 0) {
 					handleNack(records, cRecord);
 					break;
@@ -2627,9 +2632,7 @@ private void doInvokeWithRecords(final ConsumerRecords<K, V> records) {
 				}
 				this.logger.trace(() -> "Processing " + KafkaUtils.format(cRecord));
 				doInvokeRecordListener(cRecord, iterator);
-				if (this.commonRecordInterceptor !=  null) {
-					this.commonRecordInterceptor.afterRecord(cRecord, this.consumer);
-				}
+				this.commonRecordInterceptors.forEach(interceptor -> interceptor.afterRecord(cRecord, this.consumer));
 				if (this.nackSleepDurationMillis >= 0) {
 					handleNack(records, cRecord);
 					break;
@@ -2680,14 +2683,16 @@ private ConsumerRecords<K, V> checkEarlyIntercept(ConsumerRecords<K, V> nextArg)
 		private ConsumerRecord<K, V> checkEarlyIntercept(ConsumerRecord<K, V> recordArg) {
 			internalHeaders(recordArg);
 			ConsumerRecord<K, V> cRecord = recordArg;
-			if (this.earlyRecordInterceptor != null) {
-				cRecord = this.earlyRecordInterceptor.intercept(cRecord, this.consumer);
+
+			for (RecordInterceptor<K, V> earlyRecordInterceptor : this.earlyRecordInterceptors) {
+				cRecord = earlyRecordInterceptor.intercept(cRecord, this.consumer);
 				if (cRecord == null) {
 					this.logger.debug(() -> "RecordInterceptor returned null, skipping: "
-						+ KafkaUtils.format(recordArg));
+											+ KafkaUtils.format(recordArg));
 					ackCurrent(recordArg);
-					this.earlyRecordInterceptor.success(recordArg, this.consumer);
-					this.earlyRecordInterceptor.afterRecord(recordArg, this.consumer);
+					earlyRecordInterceptor.success(recordArg, this.consumer);
+					earlyRecordInterceptor.afterRecord(recordArg, this.consumer);
+					break;
 				}
 			}
 			return cRecord;
@@ -2848,13 +2853,13 @@ private void commitOffsetsIfNeededAfterHandlingError(final ConsumerRecord<K, V>
 		}
 
 		private void recordInterceptAfter(ConsumerRecord<K, V> records, @Nullable Exception exception) {
-			if (this.commonRecordInterceptor != null) {
+			if (!this.commonRecordInterceptors.isEmpty()) {
 				try {
 					if (exception == null) {
-						this.commonRecordInterceptor.success(records, this.consumer);
+						this.commonRecordInterceptors.forEach(interceptor -> interceptor.success(records, this.consumer));
 					}
 					else {
-						this.commonRecordInterceptor.failure(records, exception, this.consumer);
+						this.commonRecordInterceptors.forEach(interceptor -> interceptor.failure(records, exception, this.consumer));
 					}
 				}
 				catch (Exception e) {
@@ -2888,8 +2893,11 @@ private void invokeOnMessage(final ConsumerRecord<K, V> cRecord) {
 
 		private void doInvokeOnMessage(final ConsumerRecord<K, V> recordArg) {
 			ConsumerRecord<K, V> cRecord = recordArg;
-			if (this.recordInterceptor != null) {
-				cRecord = this.recordInterceptor.intercept(cRecord, this.consumer);
+			for (RecordInterceptor<K, V> recordInterceptor : this.recordInterceptors) {
+				cRecord = recordInterceptor.intercept(cRecord, this.consumer);
+				if (cRecord == null) {
+					break;
+				}
 			}
 			if (cRecord == null) {
 				this.logger.debug(() -> "RecordInterceptor returned null, skipping: "