spring-projectsGH-1704: Broader Batch/RecordInterceptor

Author: Karol Dowbecki

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

@@ -61,4 +61,21 @@ default void success(ConsumerRecords<K, V> records, Consumer<K, V> consumer) {
 	default void failure(ConsumerRecords<K, V> records, Exception exception, Consumer<K, V> consumer) {
 	}
 
+	/**
+	 * Called before consumer is polled.
+	 * @param consumer the consumer.
+	 * @since 2.8
+	 */
+	default void beforePoll(Consumer<K, V> consumer) {
+	}
+
+	/**
+	 * Called after listener and error handler were invoked. Last action before the
+	 * next consumer polling.
+	 * @param consumer the consumer.
+	 * @since 2.8
+	 */
+	default void finishInvoke(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 finishInvoke(Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.finishInvoke(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 finishInvoke(Consumer<K, V> consumer) {
+		this.delegates.forEach(del -> del.finishInvoke(consumer));
+	}
 }

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

@@ -1291,6 +1291,7 @@ protected void pollAndInvoke() {
 			}
 
 			invokeIfHaveRecords(records);
+			finishInvoke();
 		}
 
 		private void invokeIfHaveRecords(@Nullable ConsumerRecords<K, V> records) {
@@ -1310,6 +1311,37 @@ private void invokeIfHaveRecords(@Nullable ConsumerRecords<K, V> records) {
 			}
 		}
 
+		private void finishInvoke() {
+			if (this.isBatchListener) {
+				batchInterceptFinishInvoke();
+			}
+			else {
+				recordInterceptFinishInvoke();
+			}
+		}
+
+		private void batchInterceptFinishInvoke() {
+			if (this.commonBatchInterceptor != null) {
+				try {
+					this.commonBatchInterceptor.finishInvoke(this.consumer);
+				}
+				catch (Exception e) {
+					this.logger.error(e, "BatchInterceptor threw an exception");
+				}
+			}
+		}
+
+		private void recordInterceptFinishInvoke() {
+			if (this.commonRecordInterceptor != null) {
+				try {
+					this.commonRecordInterceptor.finishInvoke(this.consumer);
+				}
+				catch (Exception e) {
+					this.logger.error(e, "RecordInterceptor threw an exception");
+				}
+			}
+		}
+
 		private void checkIdlePartitions() {
 			Set<TopicPartition> partitions = this.consumer.assignment();
 			partitions.forEach(this::checkIdlePartition);
@@ -1422,6 +1454,7 @@ private ConsumerRecords<K, V> doPoll() {
 			ConsumerRecords<K, V> records;
 			if (this.isBatchListener && this.subBatchPerPartition) {
 				if (this.batchIterator == null) {
+					batchInterceptBeforePoll();
 					this.lastBatch = pollConsumer();
 					captureOffsets(this.lastBatch);
 					if (this.lastBatch.count() == 0) {
@@ -1439,6 +1472,7 @@ private ConsumerRecords<K, V> doPoll() {
 				}
 			}
 			else {
+				beforePoll();
 				records = pollConsumer();
 				captureOffsets(records);
 				checkRebalanceCommits();
@@ -1455,6 +1489,37 @@ private ConsumerRecords<K, V> pollConsumer() {
 			}
 		}
 
+		private void beforePoll() {
+			if (this.isBatchListener) {
+				batchInterceptBeforePoll();
+			}
+			else {
+				recordInterceptBeforePoll();
+			}
+		}
+
+		private void batchInterceptBeforePoll() {
+			if (this.commonBatchInterceptor != null) {
+				try {
+					this.commonBatchInterceptor.beforePoll(this.consumer);
+				}
+				catch (Exception e) {
+					this.logger.error(e, "BatchInterceptor threw an exception");
+				}
+			}
+		}
+
+		private void recordInterceptBeforePoll() {
+			if (this.commonRecordInterceptor != null) {
+				try {
+					this.commonRecordInterceptor.beforePoll(this.consumer);
+				}
+				catch (Exception e) {
+					this.logger.error(e, "RecordInterceptor threw an exception");
+				}
+			}
+		}
+
 		private synchronized void captureOffsets(ConsumerRecords<K, V> records) {
 			if (this.offsetsInThisBatch != null && records.count() > 0) {
 				this.offsetsInThisBatch.clear();

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

@@ -81,4 +81,21 @@ 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 before consumer is polled.
+	 * @param consumer the consumer.
+	 * @since 2.8
+	 */
+	default void beforePoll(Consumer<K, V> consumer) {
+	}
+
+	/**
+	 * Called after listener and error handler were invoked. Last action before the
+	 * next consumer polling.
+	 * @param consumer the consumer.
+	 * @since 2.8
+	 */
+	default void finishInvoke(Consumer<K, V> consumer) {
+	}
+
 }

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

@@ -98,6 +98,7 @@
 import org.springframework.kafka.event.ListenerContainerIdleEvent;
 import org.springframework.kafka.event.ListenerContainerNoLongerIdleEvent;
 import org.springframework.kafka.listener.AbstractConsumerSeekAware;
+import org.springframework.kafka.listener.BatchInterceptor;
 import org.springframework.kafka.listener.CommonLoggingErrorHandler;
 import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
 import org.springframework.kafka.listener.ConsumerAwareErrorHandler;
@@ -111,6 +112,7 @@
 import org.springframework.kafka.listener.KafkaMessageListenerContainer;
 import org.springframework.kafka.listener.ListenerExecutionFailedException;
 import org.springframework.kafka.listener.MessageListenerContainer;
+import org.springframework.kafka.listener.RecordInterceptor;
 import org.springframework.kafka.listener.adapter.FilteringMessageListenerAdapter;
 import org.springframework.kafka.listener.adapter.MessagingMessageListenerAdapter;
 import org.springframework.kafka.listener.adapter.RecordFilterStrategy;
@@ -180,7 +182,8 @@
 		"annotated25", "annotated25reply1", "annotated25reply2", "annotated26", "annotated27", "annotated28",
 		"annotated29", "annotated30", "annotated30reply", "annotated31", "annotated32", "annotated33",
 		"annotated34", "annotated35", "annotated36", "annotated37", "foo", "manualStart", "seekOnIdle",
-		"annotated38", "annotated38reply", "annotated39", "annotated40", "annotated41" })
+		"annotated38", "annotated38reply", "annotated39", "annotated40", "annotated41", "annotated42",
+		"annotated43" })
 @TestPropertySource(properties = "spel.props=fetch.min.bytes=420000,max.poll.records=10")
 public class EnableKafkaIntegrationTests {
 
@@ -600,7 +603,23 @@ public void testBatchRecords() throws Exception {
 		assertThat(this.config.listen12Message.getPayload()).isInstanceOf(List.class);
 		List<?> errorPayload = (List<?>) this.config.listen12Message.getPayload();
 		assertThat(errorPayload.size()).isGreaterThanOrEqualTo(1);
+
+	}
+
+	@Test
+	public void testBatchInterceptor() throws Exception {
+		template.send("annotated43", null, "foo");
+		template.send("annotated43", null, "bar");
+		assertThat(this.listener.latch43.await(30, TimeUnit.SECONDS)).isTrue();
+		assertThat(this.listener.payload).isInstanceOf(List.class);
+		List<?> list = (List<?>) this.listener.payload;
+		assertThat(list.size()).isGreaterThan(0);
+		assertThat(list.get(0)).isInstanceOf(ConsumerRecord.class);
+		assertThat(this.config.batchBeforePoll).isTrue();
 		assertThat(this.config.batchIntercepted).isTrue();
+		assertThat(this.config.batchSuccess).isTrue();
+		assertThat(this.config.batchFailure).isFalse();
+		assertThat(this.config.batchFinishInvoke).isTrue();
 	}
 
 	@Test
@@ -865,7 +884,6 @@ public void testKeyConversion() throws Exception {
 		this.bytesKeyTemplate.send("annotated36", "foo".getBytes(), "bar");
 		assertThat(this.listener.keyLatch.await(30, TimeUnit.SECONDS)).isTrue();
 		assertThat(this.listener.convertedKey).isEqualTo("foo");
-		assertThat(this.config.intercepted).isTrue();
 		try {
 			assertThat(this.meterRegistry.get("kafka.consumer.coordinator.join.total")
 					.tag("consumerTag", "bytesString")
@@ -893,6 +911,17 @@ public void testKeyConversion() throws Exception {
 				.isInstanceOf(CommonLoggingErrorHandler.class);
 	}
 
+	@Test
+	public void testRecordInterceptor() throws Exception {
+		this.bytesKeyTemplate.send("annotated42", "fooRecordInterceptor".getBytes(), "barRecordInterceptor");
+		assertThat(this.listener.latch42.await(30, TimeUnit.SECONDS)).isTrue();
+		assertThat(this.config.recordBeforePoll).isTrue();
+		assertThat(this.config.recordIntercepted).isTrue();
+		assertThat(this.config.recordSuccess).isTrue();
+		assertThat(this.config.recordFailure).isFalse();
+		assertThat(this.config.recordFinishInvoke).isTrue();
+	}
+
 	@Test
 	public void testProjection() throws InterruptedException {
 		template.send("annotated37", 0, "{ \"username\" : \"SomeUsername\", \"user\" : { \"name\" : \"SomeName\"}}");
@@ -984,10 +1013,26 @@ public static class Config implements KafkaListenerConfigurer {
 
 		volatile Throwable globalErrorThrowable;
 
-		volatile boolean intercepted;
+		volatile boolean recordBeforePoll;
+
+		volatile boolean recordIntercepted;
+
+		volatile boolean recordSuccess;
+
+		volatile boolean recordFailure;
+
+		volatile boolean recordFinishInvoke;
+
+		volatile boolean batchBeforePoll;
 
 		volatile boolean batchIntercepted;
 
+		volatile boolean batchSuccess;
+
+		volatile boolean batchFailure;
+
+		volatile boolean batchFinishInvoke;
+
 		@Autowired
 		private EmbeddedKafkaBroker embeddedKafka;
 
@@ -1132,9 +1177,32 @@ public KafkaListenerContainerFactory<?> bytesStringListenerContainerFactory() {
 			ConcurrentKafkaListenerContainerFactory<byte[], String> factory =
 					new ConcurrentKafkaListenerContainerFactory<>();
 			factory.setConsumerFactory(bytesStringConsumerFactory());
-			factory.setRecordInterceptor(record -> {
-				this.intercepted = true;
-				return record;
+			factory.setRecordInterceptor(new RecordInterceptor<>() {
+				@Override
+				public ConsumerRecord<byte[], String> intercept(ConsumerRecord<byte[], String> record) {
+					Config.this.recordIntercepted = true;
+					return record;
+				}
+
+				@Override
+				public void success(ConsumerRecord<byte[], String> record, Consumer<byte[], String> consumer) {
+					Config.this.recordSuccess = true;
+				}
+
+				@Override
+				public void failure(ConsumerRecord<byte[], String> record, Exception exception, Consumer<byte[], String> consumer) {
+					Config.this.recordFailure = true;
+				}
+
+				@Override
+				public void beforePoll(Consumer<byte[], String> consumer) {
+					Config.this.recordBeforePoll = true;
+				}
+
+				@Override
+				public void finishInvoke(Consumer<byte[], String> consumer) {
+					Config.this.recordFinishInvoke = true;
+				}
 			});
 			factory.setCommonErrorHandler(new CommonLoggingErrorHandler());
 			return factory;
@@ -1149,9 +1217,32 @@ 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, consumer) -> {
-				this.batchIntercepted = true;
-				return records;
+			factory.setBatchInterceptor(new BatchInterceptor<>() {
+				@Override
+				public ConsumerRecords<Integer, String> intercept(ConsumerRecords<Integer, String> records, Consumer<Integer, String> consumer) {
+					Config.this.batchIntercepted = true;
+					return records;
+				}
+
+				@Override
+				public void success(ConsumerRecords<Integer, String> records, Consumer<Integer, String> consumer) {
+					Config.this.batchSuccess = true;
+				}
+
+				@Override
+				public void failure(ConsumerRecords<Integer, String> records, Exception exception, Consumer<Integer, String> consumer) {
+					Config.this.batchFailure = true;
+				}
+
+				@Override
+				public void beforePoll(Consumer<Integer, String> consumer) {
+					Config.this.batchBeforePoll = true;
+				}
+
+				@Override
+				public void finishInvoke(Consumer<Integer, String> consumer) {
+					Config.this.batchFinishInvoke = true;
+				}
 			});
 			return factory;
 		}
@@ -1734,6 +1825,10 @@ static class Listener implements ConsumerSeekAware {
 
 		final CountDownLatch latch21 = new CountDownLatch(1);
 
+		final CountDownLatch latch42 = new CountDownLatch(1);
+
+		final CountDownLatch latch43 = new CountDownLatch(1);
+
 		final CountDownLatch validationLatch = new CountDownLatch(1);
 
 		final CountDownLatch eventLatch = new CountDownLatch(1);
@@ -2090,6 +2185,17 @@ public void bytesKey(String in, @Header(KafkaHeaders.RECEIVED_MESSAGE_KEY) Strin
 			this.keyLatch.countDown();
 		}
 
+		@KafkaListener(topics = "annotated42", containerFactory = "bytesStringListenerContainerFactory")
+		public void recordInterceptorListener(String in) {
+			this.latch42.countDown();
+		}
+
+		@KafkaListener(topics = "annotated43", containerFactory = "batchFactory")
+		public void batchInterceptorListener(List<ConsumerRecord<?, ?>> records) {
+			this.payload = records;
+			this.latch43.countDown();
+		}
+
 		@KafkaListener(topics = "annotated29")
 		public void anonymousListener(String in) {
 		}