GH-280: Add topicPattern support to KafkaMessageSource
Fixes https://github.com/spring-projects/spring-integration-kafka/issues/280 * Refactored to topics/topicPattern to be a property to prevent constructor telescoping * Add support for manual partition assignment - Cleanup after rebasing from #283 - Add additional tests * Fix checkstyle * Remove redundant factory methods * Add better test for static assignment * Use ConsumerProperties Fix typo where ContainerProperties was used instead of ConsumerProperties * Simplify construction * Remove unused import * Use ObjectUtils * Use component name for clientId * Collapse static assignment for loop * Fix lint warning
This commit is contained in:
Anshul Mehra
committed by
Artem Bilan
Artem Bilan
parent
e52c535f7c
commit
c20ef85168
@@ -55,12 +55,14 @@ import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
|
||||
import org.springframework.kafka.listener.ConsumerProperties;
|
||||
import org.springframework.kafka.support.Acknowledgment;
|
||||
import org.springframework.kafka.support.KafkaHeaders;
|
||||
import org.springframework.kafka.support.TopicPartitionOffset;
|
||||
import org.springframework.kafka.support.converter.KafkaMessageHeaders;
|
||||
import org.springframework.kafka.support.converter.MessagingMessageConverter;
|
||||
import org.springframework.kafka.support.converter.RecordMessageConverter;
|
||||
import org.springframework.lang.Nullable;
|
||||
import org.springframework.messaging.Message;
|
||||
import org.springframework.util.Assert;
|
||||
import org.springframework.util.ObjectUtils;
|
||||
import org.springframework.util.StringUtils;
|
||||
|
||||
/**
|
||||
@@ -107,17 +109,13 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
|
||||
private final KafkaAckCallbackFactory<K, V> ackCallbackFactory;
|
||||
|
||||
private final String[] topics;
|
||||
|
||||
private final Object consumerMonitor = new Object();
|
||||
|
||||
private final Map<TopicPartition, Set<KafkaAckInfo<K, V>>> inflightRecords = new ConcurrentHashMap<>();
|
||||
|
||||
private final AtomicInteger remainingCount = new AtomicInteger();
|
||||
|
||||
private String groupId;
|
||||
|
||||
private String clientId = "message.source";
|
||||
private final ConsumerProperties consumerProperties;
|
||||
|
||||
private Duration pollTimeout;
|
||||
|
||||
@@ -125,10 +123,6 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
|
||||
private Class<?> payloadType;
|
||||
|
||||
private ConsumerRebalanceListener rebalanceListener;
|
||||
|
||||
private ConsumerAwareRebalanceListener consumerAwareRebalanceListener;
|
||||
|
||||
private boolean rawMessageHeader;
|
||||
|
||||
private Duration commitTimeout;
|
||||
@@ -250,28 +244,31 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
|
||||
Assert.notNull(consumerFactory, "'consumerFactory' must not be null");
|
||||
Assert.notNull(ackCallbackFactory, "'ackCallbackFactory' must not be null");
|
||||
Assert.isTrue(consumerProperties.getTopics() != null && consumerProperties.getTopics().length > 0, "At least one topic is required");
|
||||
Assert.isTrue(
|
||||
!ObjectUtils.isEmpty(consumerProperties.getTopics())
|
||||
|| !ObjectUtils.isEmpty(consumerProperties.getTopicPartitionsToAssign())
|
||||
|| consumerProperties.getTopicPattern() != null,
|
||||
"topics, topicPattern, or topicPartitions must be provided"
|
||||
);
|
||||
this.consumerProperties = consumerProperties;
|
||||
|
||||
this.consumerFactory = fixOrRejectConsumerFactory(consumerFactory, allowMultiFetch);
|
||||
this.ackCallbackFactory = ackCallbackFactory;
|
||||
this.topics = consumerProperties.getTopics();
|
||||
this.groupId = consumerProperties.getGroupId();
|
||||
if (StringUtils.hasText(consumerProperties.getClientId())) {
|
||||
this.clientId = consumerProperties.getClientId();
|
||||
}
|
||||
this.pollTimeout = Duration.ofMillis(consumerProperties.getPollTimeout());
|
||||
this.assignTimeout = this.minTimeoutProvider.get();
|
||||
this.commitTimeout = consumerProperties.getSyncCommitTimeout();
|
||||
this.ackCallbackFactory.setCommitTimeout(consumerProperties.getSyncCommitTimeout());
|
||||
if (consumerProperties.getConsumerRebalanceListener() instanceof ConsumerAwareRebalanceListener) {
|
||||
this.consumerAwareRebalanceListener = (ConsumerAwareRebalanceListener) consumerProperties.getConsumerRebalanceListener();
|
||||
}
|
||||
else {
|
||||
this.rebalanceListener = consumerProperties.getConsumerRebalanceListener();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onInit() {
|
||||
if (!StringUtils.hasText(this.consumerProperties.getClientId())) {
|
||||
this.consumerProperties.setClientId(getComponentName());
|
||||
}
|
||||
}
|
||||
|
||||
protected String getGroupId() {
|
||||
return this.groupId;
|
||||
return this.consumerProperties.getGroupId();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -282,11 +279,11 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
*/
|
||||
@Deprecated
|
||||
public void setGroupId(String groupId) {
|
||||
this.groupId = groupId;
|
||||
this.consumerProperties.setGroupId(groupId);
|
||||
}
|
||||
|
||||
protected String getClientId() {
|
||||
return this.clientId;
|
||||
return this.consumerProperties.getClientId();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -297,7 +294,7 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
*/
|
||||
@Deprecated
|
||||
public void setClientId(String clientId) {
|
||||
this.clientId = clientId;
|
||||
this.consumerProperties.setClientId(clientId);
|
||||
}
|
||||
|
||||
protected long getPollTimeout() {
|
||||
@@ -343,7 +340,7 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
}
|
||||
|
||||
protected ConsumerRebalanceListener getRebalanceListener() {
|
||||
return this.rebalanceListener;
|
||||
return this.consumerProperties.getConsumerRebalanceListener();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -354,10 +351,7 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
*/
|
||||
@Deprecated
|
||||
public void setRebalanceListener(ConsumerRebalanceListener rebalanceListener) {
|
||||
this.rebalanceListener = rebalanceListener;
|
||||
if (rebalanceListener instanceof ConsumerAwareRebalanceListener) {
|
||||
this.consumerAwareRebalanceListener = (ConsumerAwareRebalanceListener) rebalanceListener;
|
||||
}
|
||||
this.consumerProperties.setConsumerRebalanceListener(rebalanceListener);
|
||||
}
|
||||
|
||||
@Override
|
||||
@@ -537,41 +531,103 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
|
||||
protected void createConsumer() {
|
||||
synchronized (this.consumerMonitor) {
|
||||
this.consumer = this.consumerFactory.createConsumer(this.groupId, this.clientId, null);
|
||||
boolean isConsumerAware = this.consumerAwareRebalanceListener != null;
|
||||
this.consumer.subscribe(Arrays.asList(this.topics), new ConsumerRebalanceListener() {
|
||||
this.consumer = this.consumerFactory.createConsumer(this.consumerProperties.getGroupId(),
|
||||
this.consumerProperties.getClientId(), null);
|
||||
ConsumerRebalanceListener providedRebalanceListener = this.consumerProperties
|
||||
.getConsumerRebalanceListener();
|
||||
boolean isConsumerAware = providedRebalanceListener instanceof ConsumerAwareRebalanceListener;
|
||||
ConsumerRebalanceListener rebalanceCallback = new ConsumerRebalanceListener() {
|
||||
|
||||
@Override
|
||||
public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
|
||||
KafkaMessageSource.this.assignedPartitions.clear();
|
||||
if (KafkaMessageSource.this.logger.isInfoEnabled()) {
|
||||
KafkaMessageSource.this.logger.info("Partitions revoked: " + partitions);
|
||||
KafkaMessageSource.this.logger
|
||||
.info("Partitions revoked: " + partitions);
|
||||
}
|
||||
if (isConsumerAware) {
|
||||
KafkaMessageSource.this.consumerAwareRebalanceListener.onPartitionsRevokedAfterCommit(
|
||||
KafkaMessageSource.this.consumer, partitions);
|
||||
}
|
||||
else if (KafkaMessageSource.this.rebalanceListener != null) {
|
||||
KafkaMessageSource.this.rebalanceListener.onPartitionsRevoked(partitions);
|
||||
if (providedRebalanceListener != null) {
|
||||
if (isConsumerAware) {
|
||||
((ConsumerAwareRebalanceListener) providedRebalanceListener)
|
||||
.onPartitionsRevokedAfterCommit(KafkaMessageSource.this.consumer, partitions);
|
||||
}
|
||||
else {
|
||||
providedRebalanceListener.onPartitionsRevoked(partitions);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
|
||||
KafkaMessageSource.this.assignedPartitions = new ArrayList<>(partitions);
|
||||
if (KafkaMessageSource.this.logger.isInfoEnabled()) {
|
||||
KafkaMessageSource.this.logger.info("Partitions assigned: " + partitions);
|
||||
KafkaMessageSource.this.logger
|
||||
.info("Partitions assigned: " + partitions);
|
||||
}
|
||||
if (isConsumerAware) {
|
||||
KafkaMessageSource.this.consumerAwareRebalanceListener.onPartitionsAssigned(
|
||||
KafkaMessageSource.this.consumer, partitions);
|
||||
}
|
||||
else if (KafkaMessageSource.this.rebalanceListener != null) {
|
||||
KafkaMessageSource.this.rebalanceListener.onPartitionsAssigned(partitions);
|
||||
if (providedRebalanceListener != null) {
|
||||
if (isConsumerAware) {
|
||||
((ConsumerAwareRebalanceListener) providedRebalanceListener)
|
||||
.onPartitionsAssigned(KafkaMessageSource.this.consumer, partitions);
|
||||
}
|
||||
else {
|
||||
providedRebalanceListener.onPartitionsAssigned(partitions);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
});
|
||||
};
|
||||
|
||||
if (this.consumerProperties.getTopicPattern() != null) {
|
||||
this.consumer.subscribe(this.consumerProperties.getTopicPattern(), rebalanceCallback);
|
||||
}
|
||||
else if (this.consumerProperties.getTopicPartitionsToAssign() != null) {
|
||||
List<TopicPartition> topicPartitionsToAssign = Arrays
|
||||
.stream(this.consumerProperties.getTopicPartitionsToAssign())
|
||||
.map(TopicPartitionOffset::getTopicPartition)
|
||||
.collect(Collectors.toList());
|
||||
this.consumer.assign(topicPartitionsToAssign);
|
||||
this.assignedPartitions = new ArrayList<>(topicPartitionsToAssign);
|
||||
|
||||
TopicPartitionOffset[] partitions = this.consumerProperties.getTopicPartitionsToAssign();
|
||||
|
||||
for (TopicPartitionOffset partition : partitions) {
|
||||
if (TopicPartitionOffset.SeekPosition.BEGINNING.equals(partition.getPosition())) {
|
||||
this.consumer.seekToBeginning(Collections.singleton(partition.getTopicPartition()));
|
||||
}
|
||||
else if (TopicPartitionOffset.SeekPosition.END.equals(partition.getPosition())) {
|
||||
this.consumer.seekToEnd(Collections.singleton(partition.getTopicPartition()));
|
||||
}
|
||||
else {
|
||||
TopicPartition topicPartition = partition.getTopicPartition();
|
||||
Long offset = partition.getOffset();
|
||||
if (offset != null) {
|
||||
long newOffset = offset;
|
||||
|
||||
if (offset < 0) {
|
||||
if (!partition.isRelativeToCurrent()) {
|
||||
this.consumer.seekToEnd(Collections.singleton(topicPartition));
|
||||
continue;
|
||||
}
|
||||
newOffset = Math.max(0, this.consumer.position(topicPartition) + offset);
|
||||
}
|
||||
else if (partition.isRelativeToCurrent()) {
|
||||
newOffset = this.consumer.position(topicPartition) + offset;
|
||||
}
|
||||
|
||||
try {
|
||||
this.consumer.seek(topicPartition, newOffset);
|
||||
}
|
||||
catch (Exception e) {
|
||||
this.logger.error("Failed to set initial offset for " + topicPartition
|
||||
+ " at " + newOffset + ". Position is " + this.consumer.position(topicPartition), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
this.consumer.subscribe(Arrays.asList(this.consumerProperties.getTopics()), rebalanceCallback);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -799,7 +855,7 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
|
||||
|
||||
@Override
|
||||
public String getGroupId() {
|
||||
return KafkaMessageSource.this.groupId;
|
||||
return KafkaMessageSource.this.getGroupId();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
@@ -52,21 +52,21 @@ public class KafkaInboundChannelAdapterParserTests {
|
||||
|
||||
@Test
|
||||
public void testProps() {
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "topics")).isEqualTo(new String[] { "topic1" });
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.topics")).isEqualTo(new String[] { "topic1" });
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "consumerFactory"))
|
||||
.isSameAs(this.context.getBean("consumerFactory"));
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "ackCallbackFactory"))
|
||||
.isSameAs(this.context.getBean("ackFactory"));
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "clientId")).isEqualTo("client");
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "groupId")).isEqualTo("group");
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.clientId")).isEqualTo("client");
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.groupId")).isEqualTo("group");
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "messageConverter"))
|
||||
.isSameAs(this.context.getBean("converter"));
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "payloadType")).isEqualTo(String.class);
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "rawMessageHeader", Boolean.class)).isTrue();
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "rebalanceListener"))
|
||||
assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.consumerRebalanceListener"))
|
||||
.isSameAs(this.context.getBean("rebal"));
|
||||
|
||||
assertThat(TestUtils.getPropertyValue(this.source2, "topics")).isEqualTo(new String[] { "topic1", "topic2" });
|
||||
assertThat(TestUtils.getPropertyValue(this.source2, "consumerProperties.topics")).isEqualTo(new String[] { "topic1", "topic2" });
|
||||
DefaultKafkaConsumerFactory<?, ?> cf = TestUtils.getPropertyValue(this.source2, "consumerFactory",
|
||||
DefaultKafkaConsumerFactory.class);
|
||||
assertThat(cf).isSameAs(this.context.getBean("multiFetchConsumerFactory"));
|
||||
|
||||
@@ -17,9 +17,11 @@
|
||||
package org.springframework.integration.kafka.inbound;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
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;
|
||||
@@ -29,6 +31,7 @@ import static org.mockito.BDDMockito.willReturn;
|
||||
import static org.mockito.Mockito.inOrder;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.spy;
|
||||
import static org.mockito.Mockito.times;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.time.temporal.ChronoUnit;
|
||||
@@ -40,7 +43,13 @@ 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.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
import java.util.regex.Pattern;
|
||||
import java.util.stream.Collectors;
|
||||
import java.util.stream.Stream;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.kafka.clients.consumer.Consumer;
|
||||
@@ -48,7 +57,9 @@ import org.apache.kafka.clients.consumer.ConsumerConfig;
|
||||
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.MockConsumer;
|
||||
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
|
||||
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
|
||||
import org.apache.kafka.common.TopicPartition;
|
||||
import org.apache.kafka.common.record.TimestampType;
|
||||
import org.junit.Test;
|
||||
@@ -62,8 +73,10 @@ import org.springframework.integration.acks.AcknowledgmentCallback;
|
||||
import org.springframework.integration.test.util.TestUtils;
|
||||
import org.springframework.kafka.core.ConsumerFactory;
|
||||
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
|
||||
import org.springframework.kafka.listener.ConsumerAwareRebalanceListener;
|
||||
import org.springframework.kafka.listener.ConsumerProperties;
|
||||
import org.springframework.kafka.support.KafkaHeaders;
|
||||
import org.springframework.kafka.support.TopicPartitionOffset;
|
||||
import org.springframework.kafka.test.utils.KafkaTestUtils;
|
||||
import org.springframework.messaging.Message;
|
||||
|
||||
@@ -76,6 +89,102 @@ import org.springframework.messaging.Message;
|
||||
*/
|
||||
public class MessageSourceTests {
|
||||
|
||||
@Test
|
||||
public void testIllegalArgs() {
|
||||
ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
|
||||
assertThatThrownBy(() -> new KafkaMessageSource(consumerFactory, new ConsumerProperties((Pattern) null)))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessage("topics, topicPattern, or topicPartitions must be provided");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testConsumerAwareRebalanceListener() {
|
||||
Consumer consumer = mock(Consumer.class);
|
||||
TopicPartition topicPartition = new TopicPartition("foo", 0);
|
||||
List<TopicPartition> assigned = Collections.singletonList(topicPartition);
|
||||
AtomicReference<ConsumerRebalanceListener> listener = new AtomicReference<>();
|
||||
willAnswer(i -> {
|
||||
listener.set(i.getArgument(1));
|
||||
return null;
|
||||
}).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class));
|
||||
|
||||
ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
|
||||
willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory)
|
||||
.getConfigurationProperties();
|
||||
given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer);
|
||||
ConsumerProperties consumerProperties = new ConsumerProperties("foo");
|
||||
AtomicBoolean partitionsAssignedCalled = new AtomicBoolean();
|
||||
AtomicReference<Consumer> partitionsAssignedConsumer = new AtomicReference<>();
|
||||
AtomicBoolean partitionsRevokedCalled = new AtomicBoolean();
|
||||
AtomicReference<Consumer> partitionsRevokedConsumer = new AtomicReference<>();
|
||||
consumerProperties.setConsumerRebalanceListener(new ConsumerAwareRebalanceListener() {
|
||||
@Override
|
||||
public void onPartitionsRevokedAfterCommit(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
|
||||
partitionsRevokedCalled.getAndSet(true);
|
||||
partitionsRevokedConsumer.set(consumer);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onPartitionsAssigned(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
|
||||
partitionsAssignedCalled.getAndSet(true);
|
||||
partitionsAssignedConsumer.set(consumer);
|
||||
}
|
||||
});
|
||||
KafkaMessageSource source = new KafkaMessageSource(consumerFactory, consumerProperties);
|
||||
source.setRawMessageHeader(true);
|
||||
|
||||
Message<?> received = source.receive();
|
||||
|
||||
listener.get().onPartitionsAssigned(assigned);
|
||||
assertThat(partitionsAssignedCalled.get()).isTrue();
|
||||
assertThat(partitionsAssignedConsumer.get()).isEqualTo(consumer);
|
||||
|
||||
listener.get().onPartitionsRevoked(assigned);
|
||||
assertThat(partitionsRevokedCalled.get()).isTrue();
|
||||
assertThat(partitionsRevokedConsumer.get()).isEqualTo(consumer);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testRebalanceListener() {
|
||||
Consumer consumer = mock(Consumer.class);
|
||||
TopicPartition topicPartition = new TopicPartition("foo", 0);
|
||||
List<TopicPartition> assigned = Collections.singletonList(topicPartition);
|
||||
AtomicReference<ConsumerRebalanceListener> listener = new AtomicReference<>();
|
||||
willAnswer(i -> {
|
||||
listener.set(i.getArgument(1));
|
||||
return null;
|
||||
}).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class));
|
||||
|
||||
ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
|
||||
willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory)
|
||||
.getConfigurationProperties();
|
||||
given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer);
|
||||
ConsumerProperties consumerProperties = new ConsumerProperties("foo");
|
||||
AtomicBoolean partitionsAssignedCalled = new AtomicBoolean();
|
||||
AtomicBoolean partitionsRevokedCalled = new AtomicBoolean();
|
||||
consumerProperties.setConsumerRebalanceListener(new ConsumerRebalanceListener() {
|
||||
@Override
|
||||
public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
|
||||
partitionsRevokedCalled.getAndSet(true);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
|
||||
partitionsAssignedCalled.getAndSet(true);
|
||||
}
|
||||
});
|
||||
KafkaMessageSource source = new KafkaMessageSource(consumerFactory, consumerProperties);
|
||||
source.setRawMessageHeader(true);
|
||||
|
||||
Message<?> received = source.receive();
|
||||
|
||||
listener.get().onPartitionsAssigned(assigned);
|
||||
assertThat(partitionsAssignedCalled.get()).isTrue();
|
||||
|
||||
listener.get().onPartitionsRevoked(assigned);
|
||||
assertThat(partitionsRevokedCalled.get()).isTrue();
|
||||
}
|
||||
|
||||
@SuppressWarnings({ "rawtypes", "unchecked" })
|
||||
@Test
|
||||
public void testAck() {
|
||||
@@ -553,4 +662,178 @@ public class MessageSourceTests {
|
||||
inOrder.verifyNoMoreInteractions();
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
@Test
|
||||
public void testTopicPatternBasedMessageSource() {
|
||||
MockConsumer<String, String> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
|
||||
TopicPartition topicPartition1 = new TopicPartition("abc_foo", 0);
|
||||
TopicPartition topicPartition2 = new TopicPartition("abc_foo", 1);
|
||||
TopicPartition topicPartition3 = new TopicPartition("def_foo", 0);
|
||||
TopicPartition topicPartition4 = new TopicPartition("def_foo", 1);
|
||||
List<TopicPartition> topicPartitions = Arrays
|
||||
.asList(topicPartition1, topicPartition2, topicPartition3, topicPartition4);
|
||||
|
||||
Map<TopicPartition, Long> beginningOffsets = topicPartitions.stream().collect(Collectors
|
||||
.toMap(Function.identity(), tp -> 0L));
|
||||
consumer.updateBeginningOffsets(beginningOffsets);
|
||||
|
||||
ConsumerFactory<String, String> consumerFactory = mock(ConsumerFactory.class);
|
||||
willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory)
|
||||
.getConfigurationProperties();
|
||||
given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer);
|
||||
KafkaMessageSource<String, String> source = new KafkaMessageSource<>(consumerFactory, new ConsumerProperties(Pattern.compile("[a-zA-Z0-9_]*?foo")));
|
||||
source.setRawMessageHeader(true);
|
||||
source.start();
|
||||
// force consumer creation
|
||||
source.receive();
|
||||
|
||||
consumer.rebalance(topicPartitions);
|
||||
|
||||
ConsumerRecord<String, String> record1 = new ConsumerRecord<>("abc_foo", 0, 0, null, "a");
|
||||
ConsumerRecord<String, String> record2 = new ConsumerRecord<>("abc_foo", 0, 1, null, "b");
|
||||
ConsumerRecord<String, String> record3 = new ConsumerRecord<>("abc_foo", 1, 0, null, "c");
|
||||
ConsumerRecord<String, String> record4 = new ConsumerRecord<>("def_foo", 1, 0, null, "d");
|
||||
ConsumerRecord<String, String> record5 = new ConsumerRecord<>("def_foo", 0, 0, null, "e");
|
||||
Arrays.asList(record1, record2, record3, record4, record5)
|
||||
.forEach(consumer::addRecord);
|
||||
|
||||
Message<?> received = source.receive();
|
||||
assertThat(received).isNotNull();
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class);
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record1);
|
||||
received = source.receive();
|
||||
assertThat(received).isNotNull();
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class);
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record2);
|
||||
received = source.receive();
|
||||
assertThat(received).isNotNull();
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class);
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record3);
|
||||
received = source.receive();
|
||||
assertThat(received).isNotNull();
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class);
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record4);
|
||||
received = source.receive();
|
||||
assertThat(received).isNotNull();
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class);
|
||||
assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record5);
|
||||
|
||||
source.destroy();
|
||||
}
|
||||
|
||||
@SuppressWarnings({ "rawtypes", "unchecked" })
|
||||
@Test
|
||||
public void testStaticPartitionAssignment() {
|
||||
MockConsumer<String, String> consumer = spy(new MockConsumer<>(OffsetResetStrategy.EARLIEST));
|
||||
|
||||
TopicPartition beginning = new TopicPartition("foo", 0);
|
||||
TopicPartition end = new TopicPartition("foo", 1);
|
||||
TopicPartition timestamp = new TopicPartition("foo", 2);
|
||||
TopicPartition negativeOffset = new TopicPartition("foo", 3);
|
||||
TopicPartition negativeRelativeToCurrent = new TopicPartition("foo", 4);
|
||||
TopicPartition positiveRelativeToCurrent = new TopicPartition("foo", 5);
|
||||
|
||||
List<TopicPartition> topicPartitions = Arrays.asList(beginning, end, timestamp,
|
||||
negativeOffset, negativeRelativeToCurrent, positiveRelativeToCurrent);
|
||||
|
||||
Map<TopicPartition, Long> beginningOffsets = topicPartitions.stream().collect(Collectors
|
||||
.toMap(Function.identity(), tp -> 0L));
|
||||
consumer.updateBeginningOffsets(beginningOffsets);
|
||||
Map<TopicPartition, Long> endOffsets = topicPartitions.stream().collect(Collectors
|
||||
.toMap(Function.identity(), tp -> 3L));
|
||||
consumer.updateEndOffsets(endOffsets);
|
||||
consumer.assign(topicPartitions);
|
||||
consumer.seek(timestamp, 1);
|
||||
consumer.seek(negativeRelativeToCurrent, 3);
|
||||
consumer.seek(positiveRelativeToCurrent, 1);
|
||||
|
||||
ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
|
||||
willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory)
|
||||
.getConfigurationProperties();
|
||||
given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer);
|
||||
|
||||
TopicPartitionOffset beginningTpo = new TopicPartitionOffset(beginning, null, TopicPartitionOffset.SeekPosition.BEGINNING);
|
||||
TopicPartitionOffset endTpo = new TopicPartitionOffset(end, null, TopicPartitionOffset.SeekPosition.END);
|
||||
TopicPartitionOffset timestampTpo = new TopicPartitionOffset(timestamp, null, TopicPartitionOffset.SeekPosition.TIMESTAMP);
|
||||
TopicPartitionOffset negativeOffsetTpo = new TopicPartitionOffset(negativeOffset, -1L, null);
|
||||
TopicPartitionOffset negativeRelativeToCurrentTpo = new TopicPartitionOffset(negativeRelativeToCurrent.topic(),
|
||||
negativeRelativeToCurrent.partition(), -1L, true);
|
||||
TopicPartitionOffset positiveRelativeToCurrentTpo = new TopicPartitionOffset(positiveRelativeToCurrent.topic(),
|
||||
positiveRelativeToCurrent.partition(), 1L, true);
|
||||
KafkaMessageSource source = new KafkaMessageSource(consumerFactory,
|
||||
new ConsumerProperties(beginningTpo, endTpo, timestampTpo,
|
||||
negativeOffsetTpo, negativeRelativeToCurrentTpo, positiveRelativeToCurrentTpo));
|
||||
source.setRawMessageHeader(true);
|
||||
|
||||
source.createConsumer();
|
||||
|
||||
ConsumerRecord<String, String> p0r0 = new ConsumerRecord<>("foo", 0, 0, null, "p0r0");
|
||||
ConsumerRecord<String, String> p0r1 = new ConsumerRecord<>("foo", 0, 1, null, "p0r1");
|
||||
ConsumerRecord<String, String> p0r2 = new ConsumerRecord<>("foo", 0, 2, null, "p0r2");
|
||||
ConsumerRecord<String, String> p0r3 = new ConsumerRecord<>("foo", 0, 3, null, "p0r3");
|
||||
ConsumerRecord<String, String> p1r0 = new ConsumerRecord<>("foo", 1, 0, null, "p1r0");
|
||||
ConsumerRecord<String, String> p1r1 = new ConsumerRecord<>("foo", 1, 1, null, "p1r1");
|
||||
ConsumerRecord<String, String> p1r2 = new ConsumerRecord<>("foo", 1, 2, null, "p1r2");
|
||||
ConsumerRecord<String, String> p1r3 = new ConsumerRecord<>("foo", 1, 3, null, "p1r3");
|
||||
ConsumerRecord<String, String> p2r0 = new ConsumerRecord<>("foo", 2, 0, null, "p2r0");
|
||||
ConsumerRecord<String, String> p2r1 = new ConsumerRecord<>("foo", 2, 1, null, "p2r1");
|
||||
ConsumerRecord<String, String> p2r2 = new ConsumerRecord<>("foo", 2, 2, null, "p2r2");
|
||||
ConsumerRecord<String, String> p2r3 = new ConsumerRecord<>("foo", 2, 3, null, "p2r3");
|
||||
ConsumerRecord<String, String> p3r0 = new ConsumerRecord<>("foo", 3, 0, null, "p3r0");
|
||||
ConsumerRecord<String, String> p3r1 = new ConsumerRecord<>("foo", 3, 1, null, "p3r1");
|
||||
ConsumerRecord<String, String> p3r2 = new ConsumerRecord<>("foo", 3, 2, null, "p3r2");
|
||||
ConsumerRecord<String, String> p3r3 = new ConsumerRecord<>("foo", 3, 3, null, "p3r3");
|
||||
ConsumerRecord<String, String> p4r0 = new ConsumerRecord<>("foo", 4, 0, null, "p4r0");
|
||||
ConsumerRecord<String, String> p4r1 = new ConsumerRecord<>("foo", 4, 1, null, "p4r1");
|
||||
ConsumerRecord<String, String> p4r2 = new ConsumerRecord<>("foo", 4, 2, null, "p4r2");
|
||||
ConsumerRecord<String, String> p4r3 = new ConsumerRecord<>("foo", 4, 3, null, "p4r3");
|
||||
ConsumerRecord<String, String> p5r0 = new ConsumerRecord<>("foo", 5, 0, null, "p5r0");
|
||||
ConsumerRecord<String, String> p5r1 = new ConsumerRecord<>("foo", 5, 1, null, "p5r1");
|
||||
ConsumerRecord<String, String> p5r2 = new ConsumerRecord<>("foo", 5, 2, null, "p5r2");
|
||||
ConsumerRecord<String, String> p5r3 = new ConsumerRecord<>("foo", 5, 3, null, "p5r3");
|
||||
Arrays.asList(p0r0, p0r1, p0r2, p0r3, p1r0, p1r1, p1r2, p1r3, p2r0, p2r1, p2r2, p2r3, p3r0, p3r1, p3r2, p3r3,
|
||||
p4r0, p4r1, p4r2, p4r3, p5r0, p5r1, p5r2, p5r3)
|
||||
.forEach(consumer::addRecord);
|
||||
|
||||
|
||||
Message<Object> message;
|
||||
Set<String> expected = Stream.of(
|
||||
p0r0, p0r1, p0r2, p0r3, // Seek to beginning
|
||||
p1r3, // Seek to end
|
||||
p2r1, p2r2, p2r3, // Null offset and SeekPosition.TIMESTAMP results in no change in position
|
||||
p3r3, // Negative offset ends up in seek to end
|
||||
p4r2, p4r3, // Negative offset with relative to current(3)
|
||||
p5r2, p5r3 // Positive offset with relative to current(1)
|
||||
).map(ConsumerRecord::value).collect(Collectors.toSet());
|
||||
Set<Object> received = new HashSet<>();
|
||||
while ((message = source.receive()) != null) {
|
||||
received.add(message.getHeaders().get(KafkaHeaders.RAW_DATA, ConsumerRecord.class).value());
|
||||
}
|
||||
|
||||
assertThat(received).isEqualTo(expected);
|
||||
|
||||
source.pause();
|
||||
source.receive();
|
||||
source.resume();
|
||||
source.receive();
|
||||
source.destroy();
|
||||
|
||||
InOrder inOrder = inOrder(consumer);
|
||||
inOrder.verify(consumer).assign(anyCollection());
|
||||
inOrder.verify(consumer).seekToBeginning(anyCollection());
|
||||
inOrder.verify(consumer, times(2)).seekToEnd(anyCollection());
|
||||
inOrder.verify(consumer).position(negativeRelativeToCurrent);
|
||||
inOrder.verify(consumer).seek(negativeRelativeToCurrent, 2L);
|
||||
inOrder.verify(consumer).position(positiveRelativeToCurrent);
|
||||
inOrder.verify(consumer).seek(positiveRelativeToCurrent, 2L);
|
||||
inOrder.verify(consumer).poll(any(Duration.class));
|
||||
inOrder.verify(consumer).poll(any(Duration.class));
|
||||
inOrder.verify(consumer).pause(anyCollection());
|
||||
inOrder.verify(consumer).poll(any(Duration.class));
|
||||
inOrder.verify(consumer).resume(anyCollection());
|
||||
inOrder.verify(consumer).poll(any(Duration.class));
|
||||
inOrder.verify(consumer).close();
|
||||
inOrder.verify(consumer).close(anyLong(), any(TimeUnit.class));
|
||||
inOrder.verifyNoMoreInteractions();
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user