From 9992a6ecb2b2ebc227ceb7381842524b612397e2 Mon Sep 17 00:00:00 2001 From: Gary Russell Date: Tue, 3 Dec 2019 15:18:38 -0500 Subject: [PATCH] GH-1321: Add intercept before Tx option Resolves https://github.com/spring-projects/spring-kafka/issues/1321 Provide an option to invoke `RecordInterceptor` before the transaction starts. * Fix doc typo. --- .../AbstractMessageListenerContainer.java | 17 ++++ .../ConcurrentMessageListenerContainer.java | 1 + .../KafkaMessageListenerContainer.java | 31 ++++++- ...rentMessageListenerContainerMockTests.java | 82 +++++++++++++++++++ src/reference/asciidoc/kafka.adoc | 5 ++ 5 files changed, 133 insertions(+), 3 deletions(-) diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java index ba99e910..b4abb8bb 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java @@ -95,6 +95,8 @@ public abstract class AbstractMessageListenerContainer private RecordInterceptor recordInterceptor; + private boolean interceptBeforeTx; + private volatile boolean running = false; private volatile boolean paused; @@ -303,11 +305,26 @@ public abstract class AbstractMessageListenerContainer * Does not apply to batch listeners. * @param recordInterceptor the interceptor. * @since 2.2.7 + * @see #setInterceptBeforeTx(boolean) */ public void setRecordInterceptor(RecordInterceptor recordInterceptor) { this.recordInterceptor = recordInterceptor; } + protected boolean isInterceptBeforeTx() { + return this.interceptBeforeTx; + } + + /** + * When true, invoke the interceptor before the transaction starts. + * @param interceptBeforeTx true to intercept before the transaction. + * @since 2.3.4 + * @see #setRecordInterceptor(RecordInterceptor) + */ + public void setInterceptBeforeTx(boolean interceptBeforeTx) { + this.interceptBeforeTx = interceptBeforeTx; + } + @Override public void setupMessageListener(Object messageListener) { this.containerProperties.setMessageListener(messageListener); diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java index 812cc5f9..bfc4e739 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java @@ -163,6 +163,7 @@ public class ConcurrentMessageListenerContainer extends AbstractMessageLis container.setGenericErrorHandler(getGenericErrorHandler()); container.setAfterRollbackProcessor(getAfterRollbackProcessor()); container.setRecordInterceptor(getRecordInterceptor()); + container.setInterceptBeforeTx(isInterceptBeforeTx()); container.setEmergencyStop(() -> { stop(() -> { // NOSONAR diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java index fb372156..8b59f931 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java @@ -547,7 +547,13 @@ public class KafkaMessageListenerContainer // NOSONAR line count private final Duration syncCommitTimeout; - private final RecordInterceptor recordInterceptor = getRecordInterceptor(); + private final RecordInterceptor recordInterceptor = !isInterceptBeforeTx() + ? getRecordInterceptor() + : null; + + private final RecordInterceptor earlyRecordInterceptor = isInterceptBeforeTx() + ? getRecordInterceptor() + : null; private final ConsumerSeekCallback seekCallback = new InitialOrIdleSeekCallback(); @@ -1460,7 +1466,10 @@ public class KafkaMessageListenerContainer // NOSONAR line count private void invokeRecordListenerInTx(final ConsumerRecords records) { Iterator> iterator = records.iterator(); while (iterator.hasNext()) { - final ConsumerRecord record = iterator.next(); + final ConsumerRecord record = checkEarlyIntercept(iterator.next()); + if (record == null) { + continue; + } this.logger.trace(() -> "Processing " + record); try { TransactionSupport @@ -1532,7 +1541,10 @@ public class KafkaMessageListenerContainer // NOSONAR line count private void doInvokeWithRecords(final ConsumerRecords records) { Iterator> iterator = records.iterator(); while (iterator.hasNext()) { - final ConsumerRecord record = iterator.next(); + final ConsumerRecord record = checkEarlyIntercept(iterator.next()); + if (record == null) { + continue; + } this.logger.trace(() -> "Processing " + record); doInvokeRecordListener(record, null, iterator); if (this.nackSleep >= 0) { @@ -1542,6 +1554,19 @@ public class KafkaMessageListenerContainer // NOSONAR line count } } + private ConsumerRecord checkEarlyIntercept(ConsumerRecord nextArg) { + ConsumerRecord next = nextArg; + if (this.earlyRecordInterceptor != null) { + next = this.earlyRecordInterceptor.intercept(next); + if (next == null) { + if (this.logger.isDebugEnabled()) { + this.logger.debug("RecordInterceptor returned null, skipping: " + next); + } + } + } + return next; + } + private void handleNack(final ConsumerRecords records, final ConsumerRecord record) { if (!this.autoCommit && !this.isRecordAck) { processCommits(); diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerMockTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerMockTests.java index 641d2dc4..33d75389 100644 --- a/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerMockTests.java +++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainerMockTests.java @@ -26,6 +26,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import java.time.Duration; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -45,16 +46,21 @@ import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.OffsetAndTimestamp; +import org.apache.kafka.clients.producer.Producer; import org.apache.kafka.common.TopicPartition; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.KafkaResourceHolder; +import org.springframework.kafka.core.ProducerFactory; import org.springframework.kafka.event.ConsumerFailedToStartEvent; import org.springframework.kafka.event.ConsumerStartedEvent; import org.springframework.kafka.event.ConsumerStartingEvent; import org.springframework.kafka.test.utils.KafkaTestUtils; +import org.springframework.kafka.transaction.KafkaAwareTransactionManager; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; +import org.springframework.transaction.support.TransactionSynchronizationManager; /** * @author Gary Russell @@ -353,6 +359,82 @@ public class ConcurrentMessageListenerContainerMockTests { container.stop(); } + @Test + @DisplayName("Intercept after tx start") + void testInterceptAfterTx() throws InterruptedException { + testIntercept(false); + } + + @Test + @DisplayName("Intercept before tx start") + void testInterceptBeforeTx() throws InterruptedException { + testIntercept(true); + } + + @SuppressWarnings({ "rawtypes", "unchecked" }) + void testIntercept(boolean beforeTx) throws InterruptedException { + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + final Consumer consumer = mock(Consumer.class); + TopicPartition tp0 = new TopicPartition("foo", 0); + ConsumerRecord record = new ConsumerRecord("foo", 0, 0L, "bar", "baz"); + ConsumerRecords records = new ConsumerRecords(Collections.singletonMap(tp0, Collections.singletonList(record))); + ConsumerRecords empty = new ConsumerRecords<>(Collections.emptyMap()); + AtomicBoolean first = new AtomicBoolean(true); + willAnswer(invocation -> { + Thread.sleep(10); + return first.getAndSet(false) ? records : empty; + }).given(consumer).poll(any()); + List assignments = Arrays.asList(tp0); + willAnswer(invocation -> { + ((ConsumerRebalanceListener) invocation.getArgument(1)) + .onPartitionsAssigned(assignments); + return null; + }).given(consumer).subscribe(any(Collection.class), any()); + given(consumer.position(any())).willReturn(0L); + given(consumerFactory.createConsumer("grp", "", "-0", KafkaTestUtils.defaultPropertyOverrides())) + .willReturn(consumer); + ContainerProperties containerProperties = new ContainerProperties("foo"); + containerProperties.setGroupId("grp"); + containerProperties.setMessageListener((MessageListener) rec -> { }); + containerProperties.setMissingTopicsFatal(false); + KafkaAwareTransactionManager tm = mock(KafkaAwareTransactionManager.class); + ProducerFactory pf = mock(ProducerFactory.class); + given(tm.getProducerFactory()).willReturn(pf); + Producer producer = mock(Producer.class); + given(pf.createProducer()).willReturn(producer); + containerProperties.setTransactionManager(tm); + List order = new ArrayList<>(); + CountDownLatch latch = new CountDownLatch(3); + willAnswer(inv -> { + order.add("tx"); + TransactionSynchronizationManager.bindResource(pf, + new KafkaResourceHolder<>(producer, Duration.ofSeconds(5L))); + latch.countDown(); + return null; + }).given(tm).getTransaction(any()); + ConcurrentMessageListenerContainer container = new ConcurrentMessageListenerContainer(consumerFactory, + containerProperties); + container.setRecordInterceptor(rec -> { + order.add("interceptor"); + latch.countDown(); + return rec; + }); + container.setInterceptBeforeTx(beforeTx); + container.start(); + try { + assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue(); + if (beforeTx) { + assertThat(order).containsExactly("tx", "interceptor", "tx"); // first one is on assignment + } + else { + assertThat(order).containsExactly("tx", "tx", "interceptor"); + } + } + finally { + container.stop(); + } + } + public static class TestMessageListener1 implements MessageListener, ConsumerSeekAware { private static ThreadLocal callbacks = new ThreadLocal<>(); diff --git a/src/reference/asciidoc/kafka.adoc b/src/reference/asciidoc/kafka.adoc index f843df57..e0e7aecd 100644 --- a/src/reference/asciidoc/kafka.adoc +++ b/src/reference/asciidoc/kafka.adoc @@ -657,6 +657,11 @@ The interceptor is not invoked when the listener is a <