spring-projectsGH-1704: Interceptor Improvements

Resolves spring-projects#1704

- Extract `ThreadStateProcessor` and `ConsumerAwareThreadStateProcessor`
- Avoid all the if/else tests around calling the interceptor common methods
- Add `afterRecord` to the record interceptor so sleuth can defer cleaning up until after
  the error handler.

Author: garyrussell

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

@@ -22,6 +22,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 import org.springframework.kafka.listener.ContainerProperties.EOSMode;
+import org.springframework.kafka.support.ThreadStateProcessor;
 
 /**
  * Invoked by a listener container with remaining, unprocessed, records
@@ -39,7 +40,7 @@
  *
  */
 @FunctionalInterface
-public interface AfterRollbackProcessor<K, V> {
+public interface AfterRollbackProcessor<K, V> extends ThreadStateProcessor {
 
 	/**
 	 * Process the remaining records. Recoverable will be true if the container is
@@ -63,14 +64,6 @@ public interface AfterRollbackProcessor<K, V> {
 	void process(List<ConsumerRecord<K, V>> records, Consumer<K, V> consumer,
 			MessageListenerContainer container, Exception exception, boolean recoverable, EOSMode eosMode);
 
-	/**
-	 * Optional method to clear thread state; will be called just before a consumer
-	 * thread terminates.
-	 * @since 2.2
-	 */
-	default void clearThreadState() {
-	}
-
 	/**
 	 * Return true to invoke
 	 * {@link #process(List, Consumer, MessageListenerContainer, Exception, boolean, ContainerProperties.EOSMode)}

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

@@ -32,7 +32,7 @@
  *
  */
 @FunctionalInterface
-public interface BatchInterceptor<K, V> extends BeforeAfterPollProcessor<K, V> {
+public interface BatchInterceptor<K, V> extends ConsumerAwareThreadStateProcessor {
 
 	/**
 	 * Perform some action on the records or return a different one. If null is returned

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

@@ -23,6 +23,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 
+import org.springframework.kafka.support.ThreadStateProcessor;
 import org.springframework.kafka.support.TopicPartitionOffset;
 
 /**
@@ -33,7 +34,7 @@
  * @since 2.8
  *
  */
-public interface CommonErrorHandler extends DeliveryAttemptAware {
+public interface CommonErrorHandler extends DeliveryAttemptAware, ThreadStateProcessor {
 
 	/**
 	 * Return false if this error handler should only receive the current failed record;
@@ -125,13 +126,6 @@ default int deliveryAttempt(TopicPartitionOffset topicPartitionOffset) {
 		return 0;
 	}
 
-	/**
-	 * Optional method to clear thread state; will be called just before a consumer
-	 * thread terminates.
-	 */
-	default void clearThreadState() {
-	}
-
 	/**
 	 * Return true if the offset should be committed for a handled error (no exception
 	 * thrown).

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

@@ -75,12 +75,12 @@ public void failure(ConsumerRecords<K, V> records, Exception exception, Consumer
 	}
 
 	@Override
-	public void beforePoll(Consumer<K, V> consumer) {
-		this.delegates.forEach(del -> del.beforePoll(consumer));
+	public void setupThreadState(Consumer<?, ?> consumer) {
+		this.delegates.forEach(del -> del.setupThreadState(consumer));
 	}
 
 	@Override
-	public void clearThreadState(Consumer<K, V> consumer) {
+	public void clearThreadState(Consumer<?, ?> consumer) {
 		this.delegates.forEach(del -> del.clearThreadState(consumer));
 	}
 

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

@@ -78,12 +78,13 @@ public void failure(ConsumerRecord<K, V> record, Exception exception, Consumer<K
 	}
 
 	@Override
-	public void beforePoll(Consumer<K, V> consumer) {
-		this.delegates.forEach(del -> del.beforePoll(consumer));
+	public void setupThreadState(Consumer<?, ?> consumer) {
+		this.delegates.forEach(del -> del.setupThreadState(consumer));
 	}
 
 	@Override
-	public void clearThreadState(Consumer<K, V> consumer) {
+	public void clearThreadState(Consumer<?, ?> consumer) {
 		this.delegates.forEach(del -> del.clearThreadState(consumer));
 	}
+
 }

spring-kafka/src/main/java/org/springframework/kafka/listener/BeforeAfterPollProcessor.java renamed to spring-kafka/src/main/java/org/springframework/kafka/listener/ConsumerAwareThreadStateProcessor.java

@@ -18,43 +18,34 @@
 
 import org.apache.kafka.clients.consumer.Consumer;
 
+import org.springframework.kafka.support.ThreadStateProcessor;
+
 /**
- * An interceptor for consumer poll operation.
- *
- * @param <K> the key type.
- * @param <V> the value type.
+ * A {@link ThreadStateProcessor} involving a {@link Consumer}.
  *
- * @author Gary Russell
  * @author Karol Dowbecki
- * @author Artem Bilan
+ * @author Gary Russell
  * @since 2.8
  *
  */
-public interface BeforeAfterPollProcessor<K, V> {
+public interface ConsumerAwareThreadStateProcessor extends ThreadStateProcessor {
 
 	/**
-	 * Called before consumer is polled.
-	 * <p>
-	 * It can be used to set up thread-bound resources which will be available for the
-	 * entire duration of the consumer poll operation e.g. logging with MDC.
-	 * </p>
+	 * Call to set up thread-bound resources which will be available for the
+	 * entire duration of enclosed operation involving a {@link Consumer}.
 	 *
 	 * @param consumer the consumer.
 	 */
-	default void beforePoll(Consumer<K, V> consumer) {
+	default void setupThreadState(Consumer<?, ?> consumer) {
 	}
 
 	/**
-	 * Called after records were processed by listener and error handler.
-	 * <p>
-	 * It can be used to clear thread-bound resources which were set up in {@link #beforePoll(Consumer)}.
-	 * This is the last method called by the {@link MessageListenerContainer} before the next record
-	 * processing cycle starts.
-	 * </p>
+	 * Call to clear thread-bound resources which were set up in
+	 * {@link #beforePoll(Consumer)}.
 	 *
 	 * @param consumer the consumer.
 	 */
-	default void clearThreadState(Consumer<K, V> consumer) {
+	default void clearThreadState(Consumer<?, ?> consumer) {
 	}
 
 }

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

@@ -632,6 +632,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final BatchInterceptor<K, V> commonBatchInterceptor = getBatchInterceptor();
 
+		private final ConsumerAwareThreadStateProcessor pollThreadStateProcessor;
+
 		private final ConsumerSeekCallback seekCallback = new InitialOrIdleSeekCallback();
 
 		private final long maxPollInterval;
@@ -746,12 +748,14 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 				this.batchListener = (BatchMessageListener<K, V>) listener;
 				this.isBatchListener = true;
 				this.wantsFullRecords = this.batchListener.wantsPollResult();
+				this.pollThreadStateProcessor = setUpPollProcessor(true);
 			}
 			else if (listener instanceof MessageListener) {
 				this.listener = (MessageListener<K, V>) listener;
 				this.batchListener = null;
 				this.isBatchListener = false;
 				this.wantsFullRecords = false;
+				this.pollThreadStateProcessor = setUpPollProcessor(false);
 			}
 			else {
 				throw new IllegalArgumentException("Listener must be one of 'MessageListener', "
@@ -802,6 +806,19 @@ else if (listener instanceof MessageListener) {
 			this.pausedPartitions = new HashSet<>();
 		}
 
+		@Nullable
+		private ConsumerAwareThreadStateProcessor setUpPollProcessor(boolean batch) {
+			if (batch) {
+				if (this.commonBatchInterceptor != null) {
+					return this.commonBatchInterceptor;
+				}
+			}
+			else if (this.commonRecordInterceptor != null) {
+				return this.commonRecordInterceptor;
+			}
+			return null;
+		}
+
 		@Nullable
 		private CommonErrorHandler determineCommonErrorHandler(@Nullable GenericErrorHandler<?> errHandler) {
 			CommonErrorHandler common = getCommonErrorHandler();
@@ -1314,22 +1331,8 @@ private void invokeIfHaveRecords(@Nullable ConsumerRecords<K, V> records) {
 		}
 
 		private void clearThreadState() {
-			if (this.isBatchListener) {
-				interceptClearThreadState(this.commonBatchInterceptor);
-			}
-			else {
-				interceptClearThreadState(this.commonRecordInterceptor);
-			}
-		}
-
-		private void interceptClearThreadState(BeforeAfterPollProcessor<K, V> processor) {
-			if (processor != null) {
-				try {
-					processor.clearThreadState(this.consumer);
-				}
-				catch (Exception e) {
-					this.logger.error(e, "BeforeAfterPollProcessor.clearThreadState threw an exception");
-				}
+			if (this.pollThreadStateProcessor != null) {
+				this.pollThreadStateProcessor.clearThreadState(this.consumer);
 			}
 		}
 
@@ -1480,22 +1483,8 @@ private ConsumerRecords<K, V> pollConsumer() {
 		}
 
 		private void beforePoll() {
-			if (this.isBatchListener) {
-				interceptBeforePoll(this.commonBatchInterceptor);
-			}
-			else {
-				interceptBeforePoll(this.commonRecordInterceptor);
-			}
-		}
-
-		private void interceptBeforePoll(BeforeAfterPollProcessor<K, V> processor) {
-			if (processor != null) {
-				try {
-					processor.beforePoll(this.consumer);
-				}
-				catch (Exception e) {
-					this.logger.error(e, "BeforeAfterPollProcessor.beforePoll threw an exception");
-				}
+			if (this.pollThreadStateProcessor != null) {
+				this.pollThreadStateProcessor.setupThreadState(this.consumer);
 			}
 		}
 
@@ -2294,6 +2283,9 @@ private void invokeRecordListenerInTx(final ConsumerRecords<K, V> records) {
 						TransactionSupport.clearTransactionIdSuffix();
 					}
 				}
+				if (this.commonRecordInterceptor != null) {
+					this.commonRecordInterceptor.afterRecord(record, this.consumer);
+				}
 				if (this.nackSleep >= 0) {
 					handleNack(records, record);
 					break;
@@ -2374,6 +2366,9 @@ private void doInvokeWithRecords(final ConsumerRecords<K, V> records) {
 				}
 				this.logger.trace(() -> "Processing " + ListenerUtils.recordToString(record));
 				doInvokeRecordListener(record, iterator);
+				if (this.commonRecordInterceptor !=  null) {
+					this.commonRecordInterceptor.afterRecord(record, this.consumer);
+				}
 				if (this.nackSleep >= 0) {
 					handleNack(records, record);
 					break;

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

@@ -33,7 +33,7 @@
  *
  */
 @FunctionalInterface
-public interface RecordInterceptor<K, V> extends BeforeAfterPollProcessor<K, V> {
+public interface RecordInterceptor<K, V> extends ConsumerAwareThreadStateProcessor {
 
 	/**
 	 * Perform some action on the record or return a different one. If null is returned
@@ -81,4 +81,15 @@ default void success(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
 	default void failure(ConsumerRecord<K, V> record, Exception exception, Consumer<K, V> consumer) {
 	}
 
+	/**
+	 * Called when processing the record is complete either
+	 * {@link #success(ConsumerRecord, Consumer)} or
+	 * {@link #failure(ConsumerRecord, Exception, Consumer)}.
+	 * @param record the record.
+	 * @param consumer the consumer.
+	 * @since 2.8
+	 */
+	default void afterRecord(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
+	}
+
 }

spring-kafka/src/main/java/org/springframework/kafka/support/ThreadStateProcessor.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2021 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.
+ */
+
+package org.springframework.kafka.support;
+
+/**
+ * A general interface for managing thread-bound resources.
+ *
+ * @author Gary Russell
+ * @since 2.8
+ *
+ */
+public interface ThreadStateProcessor {
+
+	/**
+	 * Call to set up thread-bound resources which will be available until
+	 * {@link #clearThreadState()} is called.
+	 */
+	default void setupThreadState() {
+	}
+
+	/**
+	 * Call to clear thread-bound resources which were set up in {@link #beforePoll()}
+	 * or by other operations that store thread state.
+	 */
+	default void clearThreadState() {
+	}
+
+}