diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java index 03129750bf..f7e0a47782 100644 --- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java +++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java @@ -17,6 +17,7 @@ package org.springframework.integration.kafka.inbound; import java.lang.reflect.Type; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -26,7 +27,7 @@ import java.util.List; import java.util.Map; import java.util.Set; import java.util.TreeSet; -import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.apache.commons.logging.Log; @@ -83,6 +84,11 @@ public class KafkaMessageSource extends AbstractMessageSource private static final long DEFAULT_POLL_TIMEOUT = 50L; + private static final long MIN_ASSIGN_TIMEOUT = 2000L; + + private final Supplier minTimeoutProvider = + () -> Duration.ofMillis(Math.max(this.pollTimeout.toMillis() * 20, MIN_ASSIGN_TIMEOUT)); + private final Log logger = LogFactory.getLog(getClass()); private final ConsumerFactory consumerFactory; @@ -99,7 +105,7 @@ public class KafkaMessageSource extends AbstractMessageSource private String clientId = "message.source"; - private long pollTimeout = DEFAULT_POLL_TIMEOUT; + private Duration pollTimeout = Duration.ofMillis(DEFAULT_POLL_TIMEOUT); private RecordMessageConverter messageConverter = new MessagingMessageConverter(); @@ -113,6 +119,10 @@ public class KafkaMessageSource extends AbstractMessageSource private boolean running; + private boolean assigned; + + private Duration assignTimeout = this.minTimeoutProvider.get(); + public KafkaMessageSource(ConsumerFactory consumerFactory, String... topics) { this(consumerFactory, new KafkaAckCallbackFactory<>(), topics); } @@ -152,7 +162,7 @@ public class KafkaMessageSource extends AbstractMessageSource } protected long getPollTimeout() { - return this.pollTimeout; + return this.pollTimeout.toMillis(); } /** @@ -160,7 +170,8 @@ public class KafkaMessageSource extends AbstractMessageSource * @param pollTimeout the poll timeout. */ public void setPollTimeout(long pollTimeout) { - this.pollTimeout = pollTimeout; + this.pollTimeout = Duration.ofMillis(pollTimeout); + this.assignTimeout = this.minTimeoutProvider.get(); } protected RecordMessageConverter getMessageConverter() { @@ -275,7 +286,7 @@ public class KafkaMessageSource extends AbstractMessageSource ConsumerRecord record; TopicPartition topicPartition; synchronized (this.consumerMonitor) { - ConsumerRecords records = this.consumer.poll(this.pollTimeout); + ConsumerRecords records = this.consumer.poll(this.assigned ? this.pollTimeout : this.assignTimeout); if (records == null || records.count() == 0) { return null; } @@ -323,6 +334,7 @@ public class KafkaMessageSource extends AbstractMessageSource @Override public void onPartitionsAssigned(Collection partitions) { + KafkaMessageSource.this.assigned = true; if (KafkaMessageSource.this.logger.isInfoEnabled()) { KafkaMessageSource.this.logger.info("Partitions assigned: " + partitions); } @@ -343,8 +355,9 @@ public class KafkaMessageSource extends AbstractMessageSource private void stopConsumer() { synchronized (this.consumerMonitor) { if (this.consumer != null) { - this.consumer.close(30, TimeUnit.SECONDS); + this.consumer.close(); this.consumer = null; + this.assigned = false; } } } diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java index bf4e93ca34..465f89e759 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java @@ -18,7 +18,6 @@ package org.springframework.integration.kafka.inbound; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.BDDMockito.given; @@ -28,6 +27,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import java.lang.reflect.Type; +import java.time.Duration; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; @@ -445,7 +445,7 @@ public class MessageDrivenAdapterTests { ConsumerRecords consumerRecords = new ConsumerRecords<>(records); ConsumerRecords emptyRecords = new ConsumerRecords<>(Collections.emptyMap()); AtomicBoolean first = new AtomicBoolean(true); - given(consumer.poll(anyLong())).willAnswer(i -> { + given(consumer.poll(any(Duration.class))).willAnswer(i -> { Thread.sleep(50); return first.getAndSet(false) ? consumerRecords : emptyRecords; }); diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java index d3818d3db4..1389ef076a 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceIntegrationTests.java @@ -55,7 +55,7 @@ public class MessageSourceIntegrationTests { @Test public void testSource() throws Exception { - Map consumerProps = KafkaTestUtils.consumerProps("foo", "false", embeddedKafka); + Map consumerProps = KafkaTestUtils.consumerProps("testSource", "false", embeddedKafka); consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 2); consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); DefaultKafkaConsumerFactory consumerFactory = new DefaultKafkaConsumerFactory<>(consumerProps); diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java index b6b62f3112..c807bab1ef 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java @@ -20,7 +20,6 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyCollection; -import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.BDDMockito.given; @@ -31,6 +30,7 @@ import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; +import java.time.Duration; import java.util.Arrays; import java.util.Collections; import java.util.HashSet; @@ -38,7 +38,6 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import org.apache.commons.logging.Log; @@ -105,7 +104,7 @@ public class MessageSourceTests { ConsumerRecords cr3 = new ConsumerRecords(records3); ConsumerRecords cr4 = new ConsumerRecords(records4); ConsumerRecords cr5 = new ConsumerRecords(Collections.emptyMap()); - given(consumer.poll(anyLong())).willReturn(cr1, cr2, cr3, cr4, cr5); + given(consumer.poll(any(Duration.class))).willReturn(cr1, cr2, cr3, cr4, cr5); ConsumerFactory consumerFactory = mock(ConsumerFactory.class); willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) .getConfigurationProperties(); @@ -135,16 +134,16 @@ public class MessageSourceTests { source.destroy(); InOrder inOrder = inOrder(consumer); inOrder.verify(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(1L))); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(2L))); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(3L))); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(4L))); - inOrder.verify(consumer).poll(anyLong()); - inOrder.verify(consumer).close(30, TimeUnit.SECONDS); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).close(); inOrder.verifyNoMoreInteractions(); } @@ -189,7 +188,7 @@ public class MessageSourceTests { ConsumerRecords cr5 = new ConsumerRecords(records5); ConsumerRecords cr6 = new ConsumerRecords(records6); ConsumerRecords cr7 = new ConsumerRecords(Collections.emptyMap()); - given(consumer.poll(anyLong())).willReturn(cr1, cr2, cr3, cr4, cr5, cr6, cr7); + given(consumer.poll(any(Duration.class))).willReturn(cr1, cr2, cr3, cr4, cr5, cr6, cr7); ConsumerFactory consumerFactory = mock(ConsumerFactory.class); willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) .getConfigurationProperties(); @@ -223,18 +222,18 @@ public class MessageSourceTests { source.destroy(); InOrder inOrder = inOrder(consumer); inOrder.verify(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).paused(); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).paused(); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).paused(); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).paused(); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(3L))); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(6L))); - inOrder.verify(consumer).poll(anyLong()); - inOrder.verify(consumer).close(30, TimeUnit.SECONDS); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).close(); inOrder.verifyNoMoreInteractions(); } @@ -263,7 +262,7 @@ public class MessageSourceTests { new ConsumerRecord("foo", 0, 1L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "bar"))); ConsumerRecords cr2 = new ConsumerRecords(records2); ConsumerRecords cr3 = new ConsumerRecords(Collections.emptyMap()); - given(consumer.poll(anyLong())).willReturn(cr1, cr1, cr2, cr2, cr3); + given(consumer.poll(any(Duration.class))).willReturn(cr1, cr1, cr2, cr2, cr3); ConsumerFactory consumerFactory = mock(ConsumerFactory.class); willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) .getConfigurationProperties(); @@ -290,16 +289,16 @@ public class MessageSourceTests { source.destroy(); assertThat(received).isNull(); InOrder inOrder = inOrder(consumer); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).seek(topicPartition, 0L); // rollback - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(1L))); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).seek(topicPartition, 1L); // rollback - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(2L))); - inOrder.verify(consumer).poll(anyLong()); - inOrder.verify(consumer).close(30, TimeUnit.SECONDS); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).close(); inOrder.verifyNoMoreInteractions(); } @@ -328,7 +327,7 @@ public class MessageSourceTests { new ConsumerRecord("foo", 0, 1L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "bar"))); ConsumerRecords cr2 = new ConsumerRecords(records2); ConsumerRecords cr3 = new ConsumerRecords(Collections.emptyMap()); - given(consumer.poll(anyLong())).willReturn(cr1, cr2, cr1, cr2, cr3); + given(consumer.poll(any(Duration.class))).willReturn(cr1, cr2, cr1, cr2, cr3); ConsumerFactory consumerFactory = mock(ConsumerFactory.class); willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) .getConfigurationProperties(); @@ -364,9 +363,9 @@ public class MessageSourceTests { source.destroy(); assertThat(received1).isNull(); InOrder inOrder = inOrder(consumer, log1, log2); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).paused(); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).seek(topicPartition, 0L); // rollback inOrder.verify(log1).isWarnEnabled(); ArgumentCaptor captor = ArgumentCaptor.forClass(String.class); @@ -380,12 +379,12 @@ public class MessageSourceTests { assertThat(captor.getValue()) .contains("Cannot commit offset for ConsumerRecord") .contains("; an earlier offset was rolled back"); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(1L))); - inOrder.verify(consumer).poll(anyLong()); + inOrder.verify(consumer).poll(any(Duration.class)); inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(2L))); - inOrder.verify(consumer).poll(anyLong()); - inOrder.verify(consumer).close(30, TimeUnit.SECONDS); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).close(); inOrder.verifyNoMoreInteractions(); }