GH-1747: Add ConsumerAwareRecordInterceptor

Resolves #1747

Also add `Consumer` to the new 2.7 methods and batch interceptor.

* Deprecate old 'intercept()`.

Author: garyrussell

spring-kafka-docs/src/main/asciidoc/kafka.adoc

@@ -1009,6 +1009,8 @@ Access to the `Consumer` object is provided.
 IMPORTANT: The `Consumer` object is not thread-safe.
 You must only invoke its methods on the thread that calls the listener.
 
+IMPORTANT: You should not execute any `Consumer<?, ?>` methods that affect the consumer's positions and or committed offsets in your listener; the container needs to manage such information.
+
 [[message-listener-container]]
 ===== Message Listener Containers
 
@@ -1024,8 +1026,12 @@ Starting with version 2.2.7, you can add a `RecordInterceptor` to the listener c
 If the interceptor returns null, the listener is not called.
 Starting with version 2.7, it has additional methods which are called after the listener exits (normally, or by throwing an exception).
 Also, starting with version 2.7, there is now a `BatchInterceptor`, providing similar functionality for <<batch-listeners>>.
+In addition, the `ConsumerAwareRecordInterceptor` (and `BatchInterceptor`) provide access to the `Consumer<?, ?>`.
+This might be used, for example, to access the consumer metrics in the interceptor.
+
+IMPORTANT: You should not execute any methods that affect the consumer's positions and or committed offsets in these interceptors; the container needs to manage such information.
 
-Starting with version 2.3, the `CompositeRecordInterceptor` and `CompositeBatchInterceptor` can be used to invoke multiple interceptors.
+The `CompositeRecordInterceptor` and `CompositeBatchInterceptor` can be used to invoke multiple interceptors.
 
 By default, when using transactions, the interceptor is invoked after the transaction has started.
 Starting with version 2.3.4, you can set the listener container's `interceptBeforeTx` property to invoke the interceptor before the transaction has started instead.

spring-kafka-docs/src/main/asciidoc/whats-new.adoc

@@ -25,6 +25,7 @@ Error handlers and after rollback processors that extend `FailedRecordProcessor`
 See See <<after-rollback>>, <<seek-to-current>>, and <<recovering-batch-eh>> for more information.
 
 The `RecordInterceptor` now has additional methods called after the listener returns (normally, or by throwing an exception).
+It also has a sub-interface `ConsumerAwareRecordInterceptor`.
 In addition, there is now a `BatchInterceptor` for batch listeners.
 See <<message-listener-container>> for more information.
 

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

@@ -16,6 +16,7 @@
 
 package org.springframework.kafka.listener;
 
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 
 import org.springframework.lang.Nullable;
@@ -37,24 +38,27 @@ public interface BatchInterceptor<K, V> {
 	 * Perform some action on the records or return a different one. If null is returned
 	 * the records will be skipped. Invoked before the listener.
 	 * @param records the records.
+	 * @param consumer the consumer.
 	 * @return the records or null.
 	 */
 	@Nullable
-	ConsumerRecords<K, V> intercept(ConsumerRecords<K, V> records);
+	ConsumerRecords<K, V> intercept(ConsumerRecords<K, V> records, Consumer<K, V> consumer);
 
 	/**
 	 * Called after the listener exits normally.
 	 * @param records the records.
+	 * @param consumer the consumer.
 	 */
-	default void success(ConsumerRecords<K, V> records) {
+	default void success(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
 	}
 
 	/**
 	 * Called after the listener throws an exception.
 	 * @param records the records.
 	 * @param exception the exception.
+	 * @param consumer the consumer.
 	 */
-	default void failure(ConsumerRecords<K, V> records, Exception exception) {
+	default void failure(ConsumerRecords<K, V> records, Exception exception, Consumer<K, V> consumer) {
 	}
 
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019 the original author or authors.
+ * Copyright 2019-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.
@@ -20,6 +20,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 
 import org.springframework.util.Assert;
@@ -52,22 +53,25 @@ public CompositeBatchInterceptor(BatchInterceptor<K, V>... delegates) {
 	}
 
 	@Override
-	public ConsumerRecords<K, V> intercept(ConsumerRecords<K, V> records) {
+	public ConsumerRecords<K, V> intercept(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
 		ConsumerRecords<K, V> recordsToIntercept = records;
 		for (BatchInterceptor<K, V> delegate : this.delegates) {
-			recordsToIntercept = delegate.intercept(recordsToIntercept);
+			recordsToIntercept = delegate.intercept(recordsToIntercept, consumer);
+			if (recordsToIntercept == null) {
+				break;
+			}
 		}
 		return recordsToIntercept;
 	}
 
 	@Override
-	public void success(ConsumerRecords<K, V> records) {
-		this.delegates.forEach(del -> del.success(records));
+	public void success(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.success(records, consumer));
 	}
 
 	@Override
-	public void failure(ConsumerRecords<K, V> records, Exception exception) {
-		this.delegates.forEach(del -> del.failure(records, exception));
+	public void failure(ConsumerRecords<K, V> records, Exception exception, Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.failure(records, exception, consumer));
 	}
 
 }

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

@@ -20,8 +20,10 @@
 import java.util.Arrays;
 import java.util.Collection;
 
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
@@ -36,7 +38,7 @@
  * @since 2.3
  *
  */
-public class CompositeRecordInterceptor<K, V> implements RecordInterceptor<K, V> {
+public class CompositeRecordInterceptor<K, V> implements ConsumerAwareRecordInterceptor<K, V> {
 
 	private final Collection<RecordInterceptor<K, V>> delegates = new ArrayList<>();
 
@@ -53,22 +55,26 @@ public CompositeRecordInterceptor(RecordInterceptor<K, V>... delegates) {
 	}
 
 	@Override
-	public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record) {
+	@Nullable
+	public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
 		ConsumerRecord<K, V> recordToIntercept = record;
 		for (RecordInterceptor<K, V> delegate : this.delegates) {
-			recordToIntercept = delegate.intercept(recordToIntercept);
+			recordToIntercept = delegate.intercept(recordToIntercept, consumer);
+			if (recordToIntercept == null) {
+				break;
+			}
 		}
 		return recordToIntercept;
 	}
 
 	@Override
-	public void success(ConsumerRecord<K, V> record) {
-		this.delegates.forEach(del -> del.success(record));
+	public void success(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.success(record, consumer));
 	}
 
 	@Override
-	public void failure(ConsumerRecord<K, V> record, Exception exception) {
-		this.delegates.forEach(del -> del.failure(record, exception));
+	public void failure(ConsumerRecord<K, V> record, Exception exception, Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.failure(record, exception, consumer));
 	}
 
 }

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

@@ -0,0 +1,48 @@
+/*
+ * 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.listener;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * A {@link RecordInterceptor} that has access to the {@link Consumer}.
+ *
+ * @param <K> the key type.
+ * @param <V> the value type.
+ *
+ * @author Gary Russell
+ * @since 2.7
+ *
+ */
+@FunctionalInterface
+public interface ConsumerAwareRecordInterceptor<K, V> extends RecordInterceptor<K, V> {
+
+	@SuppressWarnings("deprecation")
+	@Override
+	@Nullable
+	default ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record) {
+		throw new UnsupportedOperationException("Container should never call this");
+	}
+
+	@Override
+	@Nullable
+	ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record, Consumer<K, V> consumer);
+
+}

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

@@ -1789,7 +1789,7 @@ protected void doInTransactionWithoutResult(TransactionStatus status) {
 		}
 
 		private void batchAfterRollback(final ConsumerRecords<K, V> records,
-				final List<ConsumerRecord<K, V>> recordList, RuntimeException e,
+				@Nullable final List<ConsumerRecord<K, V>> recordList, RuntimeException e,
 				AfterRollbackProcessor<K, V> afterRollbackProcessorToUse) {
 
 			RuntimeException rollbackException = decorateException(e);
@@ -1888,10 +1888,10 @@ private void batchInterceptAfter(ConsumerRecords<K, V> records, @Nullable Except
 			if (this.commonBatchInterceptor != null) {
 				try {
 					if (exception == null) {
-						this.commonBatchInterceptor.success(records);
+						this.commonBatchInterceptor.success(records, this.consumer);
 					}
 					else {
-						this.commonBatchInterceptor.failure(records, exception);
+						this.commonBatchInterceptor.failure(records, exception, this.consumer);
 					}
 				}
 				catch (Exception e) {
@@ -1961,7 +1961,7 @@ private void invokeBatchOnMessageWithRecordsOrList(final ConsumerRecords<K, V> r
 
 			ConsumerRecords<K, V> records = recordsArg;
 			if (this.batchInterceptor != null) {
-				records = this.batchInterceptor.intercept(recordsArg);
+				records = this.batchInterceptor.intercept(recordsArg, this.consumer);
 			}
 			if (this.wantsFullRecords) {
 				this.batchListener.onMessage(records, // NOSONAR
@@ -2143,7 +2143,7 @@ private void doInvokeWithRecords(final ConsumerRecords<K, V> records) {
 		private ConsumerRecords<K, V> checkEarlyIntercept(ConsumerRecords<K, V> nextArg) {
 			ConsumerRecords<K, V> next = nextArg;
 			if (this.earlyBatchInterceptor != null) {
-				next = this.earlyBatchInterceptor.intercept(next);
+				next = this.earlyBatchInterceptor.intercept(next, this.consumer);
 				if (next == null) {
 					this.logger.debug(() -> "RecordInterceptor returned null, skipping: "
 						+ nextArg + " with " + nextArg.count() + " records");
@@ -2156,7 +2156,7 @@ private ConsumerRecords<K, V> checkEarlyIntercept(ConsumerRecords<K, V> nextArg)
 		private ConsumerRecord<K, V> checkEarlyIntercept(ConsumerRecord<K, V> nextArg) {
 			ConsumerRecord<K, V> next = nextArg;
 			if (this.earlyRecordInterceptor != null) {
-				next = this.earlyRecordInterceptor.intercept(next);
+				next = this.earlyRecordInterceptor.intercept(next, this.consumer);
 				if (next == null) {
 					this.logger.debug(() -> "RecordInterceptor returned null, skipping: "
 						+ ListenerUtils.recordToString(nextArg));
@@ -2258,10 +2258,10 @@ private void recordInterceptAfter(ConsumerRecord<K, V> records, @Nullable Except
 			if (this.commonRecordInterceptor != null) {
 				try {
 					if (exception == null) {
-						this.commonRecordInterceptor.success(records);
+						this.commonRecordInterceptor.success(records, this.consumer);
 					}
 					else {
-						this.commonRecordInterceptor.failure(records, exception);
+						this.commonRecordInterceptor.failure(records, exception, this.consumer);
 					}
 				}
 				catch (Exception e) {
@@ -2301,7 +2301,7 @@ private void invokeOnMessage(final ConsumerRecord<K, V> record) {
 		private void doInvokeOnMessage(final ConsumerRecord<K, V> recordArg) {
 			ConsumerRecord<K, V> record = recordArg;
 			if (this.recordInterceptor != null) {
-				record = this.recordInterceptor.intercept(record);
+				record = this.recordInterceptor.intercept(record, this.consumer);
 			}
 			if (record == null) {
 				this.logger.debug(() -> "RecordInterceptor returned null, skipping: "

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

@@ -16,6 +16,7 @@
 
 package org.springframework.kafka.listener;
 
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 import org.springframework.lang.Nullable;
@@ -39,25 +40,45 @@ public interface RecordInterceptor<K, V> {
 	 * the record will be skipped. Invoked before the listener.
 	 * @param record the record.
 	 * @return the record or null.
+	 * @deprecated in favor of {@link #intercept(ConsumerRecord, Consumer)} which will
+	 * become the required method in a future release.
 	 */
+	@Deprecated
 	@Nullable
 	ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record);
 
+	/**
+	 * Perform some action on the record or return a different one. If null is returned
+	 * the record will be skipped. Invoked before the listener.
+	 * @param record the record.
+	 * @param consumer the consumer.
+	 * @return the record or null.
+	 * @since 2.7
+	 */
+	@Nullable
+	default ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record,
+			@SuppressWarnings("unused") Consumer<K, V> consumer) {
+
+		return intercept(record);
+	}
+
 	/**
 	 * Called after the listener exits normally.
 	 * @param record the record.
+	 * @param consumer the consumer.
 	 * @since 2.7
 	 */
-	default void success(ConsumerRecord<K, V> record) {
+	default void success(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
 	}
 
 	/**
 	 * Called after the listener throws an exception.
 	 * @param record the record.
 	 * @param exception the exception.
+	 * @param consumer the consumer.
 	 * @since 2.7
 	 */
-	default void failure(ConsumerRecord<K, V> record, Exception exception) {
+	default void failure(ConsumerRecord<K, V> record, Exception exception, Consumer<K, V> consumer) {
 	}
 
 }

spring-kafka/src/test/java/org/springframework/kafka/annotation/EnableKafkaIntegrationTests.java

@@ -1084,7 +1084,7 @@ public KafkaListenerContainerFactory<?> batchFactory() {
 			factory.setRecordFilterStrategy(recordFilter());
 			// always send to the same partition so the replies are in order for the test
 			factory.setReplyTemplate(partitionZeroReplyTemplate());
-			factory.setBatchInterceptor(records -> {
+			factory.setBatchInterceptor((records, consumer) -> {
 				this.batchIntercepted = true;
 				return records;
 			});