GH-1704 Broader Batch/RecordInterceptor (#1912)

* GH-1704: Broader Batch/RecordInterceptor

* GH-1704: More pessimistic finally

* Precise logger message

* Renaming methods and improving Javadoc.

Still need to work on race condition in the tests because beforePoll() and afterRecordsProcessed() is called while the test is setting up data.

* Fixing tests

* Fixing tests

* Moving beforePoll() and adding InOrder tests

* Reverting integration tests

* Extracting BeforeAfterPollProcessor

* Renaming afterPoll to clearThreadState

* Fixing compiler warnings in tests

Author: Karol Dowbecki

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

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

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

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2016-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.listener;
+
+import org.apache.kafka.clients.consumer.Consumer;
+
+/**
+ * An interceptor for consumer poll operation.
+ *
+ * @param <K> the key type.
+ * @param <V> the value type.
+ *
+ * @author Gary Russell
+ * @author Karol Dowbecki
+ * @author Artem Bilan
+ * @since 2.8
+ *
+ */
+public interface BeforeAfterPollProcessor<K, V> {
+
+	/**
+	 * 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>
+	 *
+	 * @param consumer the consumer.
+	 */
+	default void beforePoll(Consumer<K, V> 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>
+	 *
+	 * @param consumer the consumer.
+	 */
+	default void clearThreadState(Consumer<K, V> consumer) {
+	}
+
+}

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

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

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

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

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

@@ -1236,6 +1236,9 @@ public void run() {
 				catch (Exception e) {
 					handleConsumerException(e);
 				}
+				finally {
+					clearThreadState();
+				}
 			}
 			wrapUp(exitThrowable);
 		}
@@ -1310,6 +1313,26 @@ 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");
+				}
+			}
+		}
+
 		private void checkIdlePartitions() {
 			Set<TopicPartition> partitions = this.consumer.assignment();
 			partitions.forEach(this::checkIdlePartition);
@@ -1447,6 +1470,7 @@ private ConsumerRecords<K, V> doPoll() {
 		}
 
 		private ConsumerRecords<K, V> pollConsumer() {
+			beforePoll();
 			try {
 				return this.consumer.poll(this.pollTimeout);
 			}
@@ -1455,6 +1479,26 @@ 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");
+				}
+			}
+		}
+
 		private synchronized void captureOffsets(ConsumerRecords<K, V> records) {
 			if (this.offsetsInThisBatch != null && records.count() > 0) {
 				this.offsetsInThisBatch.clear();
@@ -2082,7 +2126,7 @@ private void batchInterceptAfter(ConsumerRecords<K, V> records, @Nullable Except
 					}
 				}
 				catch (Exception e) {
-					this.logger.error(e, "BatchInterceptor threw an exception");
+					this.logger.error(e, "BatchInterceptor.success/failure threw an exception");
 				}
 			}
 		}
@@ -2463,7 +2507,7 @@ private void recordInterceptAfter(ConsumerRecord<K, V> records, @Nullable Except
 					}
 				}
 				catch (Exception e) {
-					this.logger.error(e, "RecordInterceptor threw an exception");
+					this.logger.error(e, "RecordInterceptor.success/failure threw an exception");
 				}
 			}
 		}

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

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