Allow fetching multiple records

Resolves https://github.com/spring-projects/spring-integration-kafka/issues/267

- add an option to allow fetching multiple records per poll

* DSL for new constructor arg
This commit is contained in:
Gary Russell
2019-04-17 17:54:14 -04:00
committed by Artem Bilan
parent cd6d17e4b0
commit c25b3b45f4
8 changed files with 263 additions and 21 deletions

View File

@@ -48,6 +48,10 @@ public class KafkaInboundChannelAdapterParser extends AbstractPollingInboundChan
if (StringUtils.hasText(attribute)) {
builder.addConstructorArgReference(attribute);
}
attribute = element.getAttribute("allow-multi-fetch");
if (StringUtils.hasText(attribute)) {
builder.addConstructorArgValue(attribute);
}
builder.addConstructorArgValue(element.getAttribute("topics"));
IntegrationNamespaceUtils.setValueIfAttributeDefined(builder, element, "client-id");
IntegrationNamespaceUtils.setValueIfAttributeDefined(builder, element, "group-id");

View File

@@ -82,7 +82,24 @@ public final class Kafka {
public static <K, V> KafkaInboundChannelAdapterSpec<K, V> inboundChannelAdapter(
ConsumerFactory<K, V> consumerFactory, String... topics) {
return new KafkaInboundChannelAdapterSpec<>(consumerFactory, topics);
return inboundChannelAdapter(consumerFactory, false, topics);
}
/**
* Create an initial {@link KafkaInboundChannelAdapterSpec} with the consumer factory and
* topics.
* @param consumerFactory the consumer factory.
* @param allowMultiFetch true to fetch multiple records on each poll.
* @param topics the topic(s).
* @param <K> the Kafka message key type.
* @param <V> the Kafka message value type.
* @return the spec.
* @since 3.2
*/
public static <K, V> KafkaInboundChannelAdapterSpec<K, V> inboundChannelAdapter(
ConsumerFactory<K, V> consumerFactory, boolean allowMultiFetch, String... topics) {
return new KafkaInboundChannelAdapterSpec<>(consumerFactory, allowMultiFetch, topics);
}
/**
@@ -100,7 +117,28 @@ public final class Kafka {
ConsumerFactory<K, V> consumerFactory,
KafkaAckCallbackFactory<K, V> ackCallbackFactory, String... topics) {
return new KafkaInboundChannelAdapterSpec<>(consumerFactory, ackCallbackFactory, topics);
return inboundChannelAdapter(consumerFactory, ackCallbackFactory, false, topics);
}
/**
* Create an initial {@link KafkaInboundChannelAdapterSpec} with the consumer factory and
* topics with a custom ack callback factory.
* @param consumerFactory the consumer factory.
* @param ackCallbackFactory the callback factory.
* @param allowMultiFetch true to fetch multiple records on each poll.
* @param topics the topic(s).
* @param <K> the Kafka message key type.
* @param <V> the Kafka message value type.
* @return the spec.
* @since 3.0.1
*/
public static <K, V> KafkaInboundChannelAdapterSpec<K, V> inboundChannelAdapter(
ConsumerFactory<K, V> consumerFactory,
KafkaAckCallbackFactory<K, V> ackCallbackFactory,
boolean allowMultiFetch,
String... topics) {
return new KafkaInboundChannelAdapterSpec<>(consumerFactory, ackCallbackFactory, allowMultiFetch, topics);
}
/**

View File

@@ -38,14 +38,14 @@ import org.springframework.kafka.support.converter.RecordMessageConverter;
public class KafkaInboundChannelAdapterSpec<K, V>
extends MessageSourceSpec<KafkaInboundChannelAdapterSpec<K, V>, KafkaMessageSource<K, V>> {
KafkaInboundChannelAdapterSpec(ConsumerFactory<K, V> consumerFactory, String... topics) {
this.target = new KafkaMessageSource<>(consumerFactory, topics);
KafkaInboundChannelAdapterSpec(ConsumerFactory<K, V> consumerFactory, boolean allowMultiFetch, String... topics) {
this.target = new KafkaMessageSource<>(consumerFactory, allowMultiFetch, topics);
}
KafkaInboundChannelAdapterSpec(ConsumerFactory<K, V> consumerFactory,
KafkaAckCallbackFactory<K, V> ackCallbackFactory, String... topics) {
KafkaAckCallbackFactory<K, V> ackCallbackFactory, boolean allowMultiFetch, String... topics) {
this.target = new KafkaMessageSource<>(consumerFactory, ackCallbackFactory, topics);
this.target = new KafkaMessageSource<>(consumerFactory, ackCallbackFactory, allowMultiFetch, topics);
}
public KafkaInboundChannelAdapterSpec<K, V> groupId(String groupId) {

View File

@@ -22,11 +22,13 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.stream.Collectors;
@@ -92,6 +94,12 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
private static final long MIN_ASSIGN_TIMEOUT = 2000L;
/**
* The number of records remaining from the previous poll.
* @since 3.2
*/
public static final String REMAINING_RECORDS = KafkaHeaders.PREFIX + "remainingRecords";
private final Supplier<Duration> minTimeoutProvider =
() -> Duration.ofMillis(Math.max(this.pollTimeout.toMillis() * 20, MIN_ASSIGN_TIMEOUT));
@@ -105,6 +113,8 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
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";
@@ -135,17 +145,79 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
private volatile boolean paused;
private volatile Iterator<ConsumerRecord<K, V>> recordsIterator;
/**
* Construct an instance with the supplied parameters. Fetching multiple
* records per poll will be disabled.
*
* @param consumerFactory the consumer factory.
* @param topics the topics.
* @see #KafkaMessageSource(ConsumerFactory, KafkaAckCallbackFactory, boolean, String...)
*/
public KafkaMessageSource(ConsumerFactory<K, V> consumerFactory, String... topics) {
this(consumerFactory, new KafkaAckCallbackFactory<>(), topics);
this(consumerFactory, new KafkaAckCallbackFactory<>(), false, topics);
}
/**
* Construct an instance with the supplied parameters. Set 'allowMultiFetch' to true
* to allow up to {@code max.poll.records} to be fetched on each poll. When false
* (default) {@code max.poll.records} is coerced to 1 if the consumer factory is a
* {@link DefaultKafkaConsumerFactory} or otherwise rejected with an
* {@link IllegalArgumentException}. IMPORTANT: When true, you must call
* {@link #receive()} at a sufficient rate to consume the number of records received
* within {@code max.poll.interval.ms}. When false, you must call {@link #receive()}
* within {@code max.poll.interval.ms}. {@link #pause()} will not take effect until
* the records from the previous poll are consumed.
*
* @param consumerFactory the consumer factory.
* @param allowMultiFetch true to allow {@code max.poll.records > 1}.
* @param topics the topics.
* @since 3.2
*/
public KafkaMessageSource(ConsumerFactory<K, V> consumerFactory, boolean allowMultiFetch, String... topics) {
this(consumerFactory, new KafkaAckCallbackFactory<>(), allowMultiFetch, topics);
}
/**
* Construct an instance with the supplied parameters. Fetching multiple
* records per poll will be disabled.
*
* @param consumerFactory the consumer factory.
* @param ackCallbackFactory the ack callback factory.
* @param topics the topics.
* @see #KafkaMessageSource(ConsumerFactory, KafkaAckCallbackFactory, boolean, String...)
*/
public KafkaMessageSource(ConsumerFactory<K, V> consumerFactory,
KafkaAckCallbackFactory<K, V> ackCallbackFactory, String... topics) {
this(consumerFactory, ackCallbackFactory, false, topics);
}
/**
* Construct an instance with the supplied parameters. Set 'allowMultiFetch' to true
* to allow up to {@code max.poll.records} to be fetched on each poll. When false
* (default) {@code max.poll.records} is coerced to 1 if the consumer factory is a
* {@link DefaultKafkaConsumerFactory} or otherwise rejected with an
* {@link IllegalArgumentException}. IMPORTANT: When true, you must call
* {@link #receive()} at a sufficient rate to consume the number of records received
* within {@code max.poll.interval.ms}. When false, you must call {@link #receive()}
* within {@code max.poll.interval.ms}. {@link #pause()} will not take effect until
* the records from the previous poll are consumed.
*
* @param consumerFactory the consumer factory.
* @param ackCallbackFactory the ack callback factory.
* @param allowMultiFetch true to allow {@code max.poll.records > 1}.
* @param topics the topics.
* @since 3.2
*/
public KafkaMessageSource(ConsumerFactory<K, V> consumerFactory,
KafkaAckCallbackFactory<K, V> ackCallbackFactory, boolean allowMultiFetch, String... topics) {
Assert.notNull(consumerFactory, "'consumerFactory' must not be null");
Assert.notNull(ackCallbackFactory, "'ackCallbackFactory' must not be null");
Assert.isTrue(topics != null && topics.length > 0, "At least one topic is required");
this.consumerFactory = fixOrRejectConsumerFactory(consumerFactory);
this.consumerFactory = fixOrRejectConsumerFactory(consumerFactory, allowMultiFetch);
this.ackCallbackFactory = ackCallbackFactory;
this.topics = topics;
}
@@ -261,11 +333,12 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
this.ackCallbackFactory.setCommitTimeout(commitTimeout);
}
private ConsumerFactory<K, V> fixOrRejectConsumerFactory(ConsumerFactory<K, V> suppliedConsumerFactory) {
private ConsumerFactory<K, V> fixOrRejectConsumerFactory(ConsumerFactory<K, V> suppliedConsumerFactory,
boolean allowMultiFetch) {
Object maxPoll = suppliedConsumerFactory.getConfigurationProperties()
.get(ConsumerConfig.MAX_POLL_RECORDS_CONFIG);
if (maxPoll == null || (maxPoll instanceof Number && ((Number) maxPoll).intValue() != 1)
|| (maxPoll instanceof String && Integer.parseInt((String) maxPoll) != 1)) {
if (!allowMultiFetch && (maxPoll == null || maxPollGtrOne(maxPoll))) {
if (!suppliedConsumerFactory.getClass().getName().equals(DefaultKafkaConsumerFactory.class.getName())) {
throw new IllegalArgumentException("Custom consumer factory is not configured with '"
+ ConsumerConfig.MAX_POLL_RECORDS_CONFIG + " = 1'");
@@ -291,6 +364,18 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
}
}
private boolean maxPollGtrOne(Object maxPoll) {
return maxPollNumberGtrOne(maxPoll) || maxPollStringGtr1(maxPoll);
}
private boolean maxPollNumberGtrOne(Object maxPoll) {
return maxPoll instanceof Number && ((Number) maxPoll).intValue() != 1;
}
private boolean maxPollStringGtr1(Object maxPoll) {
return maxPoll instanceof String && Integer.parseInt((String) maxPoll) != 1;
}
@Override
public synchronized boolean isRunning() {
return this.running;
@@ -339,19 +424,27 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
this.consumer.resume(this.assignedPartitions);
this.paused = false;
}
if (this.paused) {
if (this.paused && this.recordsIterator == null) {
this.logger.debug("Consumer is paused; no records will be returned");
}
ConsumerRecord<K, V> record;
TopicPartition topicPartition;
synchronized (this.consumerMonitor) {
ConsumerRecords<K, V> records = this.consumer.poll(this.assignedPartitions.isEmpty() ? this.assignTimeout : this.pollTimeout);
if (records == null || records.count() == 0) {
return null;
}
record = records.iterator().next();
topicPartition = new TopicPartition(record.topic(), record.partition());
if (this.recordsIterator != null) {
record = nextRecord();
}
else {
synchronized (this.consumerMonitor) {
ConsumerRecords<K, V> records = this.consumer
.poll(this.assignedPartitions.isEmpty() ? this.assignTimeout : this.pollTimeout);
if (records == null || records.count() == 0) {
return null;
}
this.remainingCount.set(records.count());
this.recordsIterator = records.iterator();
record = nextRecord();
}
}
topicPartition = new TopicPartition(record.topic(), record.partition());
KafkaAckInfo<K, V> ackInfo = new KafkaAckInfoImpl(record, topicPartition);
AcknowledgmentCallback ackCallback = this.ackCallbackFactory.createCallback(ackInfo);
this.inflightRecords.computeIfAbsent(topicPartition, tp -> Collections.synchronizedSet(new TreeSet<>()))
@@ -362,6 +455,7 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
if (message.getHeaders() instanceof KafkaMessageHeaders) {
Map<String, Object> rawHeaders = ((KafkaMessageHeaders) message.getHeaders()).getRawHeaders();
rawHeaders.put(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK, ackCallback);
rawHeaders.put(REMAINING_RECORDS, this.remainingCount.get());
if (this.rawMessageHeader) {
rawHeaders.put(KafkaHeaders.RAW_DATA, record);
}
@@ -369,7 +463,8 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
}
else {
AbstractIntegrationMessageBuilder<?> builder = getMessageBuilderFactory().fromMessage(message)
.setHeader(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK, ackCallback);
.setHeader(IntegrationMessageHeaderAccessor.ACKNOWLEDGMENT_CALLBACK, ackCallback)
.setHeader(REMAINING_RECORDS, this.remainingCount.get());
if (this.rawMessageHeader) {
builder.setHeader(KafkaHeaders.RAW_DATA, record);
}
@@ -377,6 +472,16 @@ public class KafkaMessageSource<K, V> extends AbstractMessageSource<Object> impl
}
}
private ConsumerRecord<K, V> nextRecord() {
ConsumerRecord<K, V> record;
record = this.recordsIterator.next();
if (!this.recordsIterator.hasNext()) {
this.recordsIterator = null;
}
this.remainingCount.decrementAndGet();
return record;
}
protected void createConsumer() {
synchronized (this.consumerMonitor) {
this.consumer = this.consumerFactory.createConsumer(this.groupId, this.clientId, null);

View File

@@ -245,6 +245,21 @@
<xsd:union memberTypes="xsd:int xsd:boolean"/>
</xsd:simpleType>
</xsd:attribute>
<xsd:attribute name="allow-multi-fetch">
<xsd:annotation>
<xsd:documentation>
Allow fetching multiple records per poll.
IMPORTANT: When true, you must poll the adapter at a sufficient rate to consume
the number of records received within 'max.poll.interval.ms'.
When false, you must poll the adapter within 'max.poll.interval.ms'.
Pausing the adapter will not take effect until the records from the previous poll
are consumed. Default 'false'.
</xsd:documentation>
</xsd:annotation>
<xsd:simpleType>
<xsd:union memberTypes="xsd:boolean xsd:string"/>
</xsd:simpleType>
</xsd:attribute>
</xsd:complexType>
</xsd:element>

View File

@@ -25,7 +25,8 @@
<int-kafka:inbound-channel-adapter
id="adapter2"
consumer-factory="consumerFactory"
consumer-factory="multiFetchConsumerFactory"
allow-multi-fetch="true"
topics="topic1, topic2"
group-id="group"
auto-startup="false"
@@ -41,6 +42,14 @@
</constructor-arg>
</bean>
<bean id="multiFetchConsumerFactory" class="org.springframework.kafka.core.DefaultKafkaConsumerFactory">
<constructor-arg>
<map>
<entry key="max.poll.records" value="10"/>
</map>
</constructor-arg>
</bean>
<bean id="ackFactory" class="org.mockito.Mockito" factory-method="mock">
<constructor-arg
value="org.springframework.integration.kafka.inbound.KafkaMessageSource$KafkaAckCallbackFactory"/>

View File

@@ -18,6 +18,7 @@ package org.springframework.integration.kafka.config.xml;
import static org.assertj.core.api.Assertions.assertThat;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
@@ -25,6 +26,7 @@ import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.ApplicationContext;
import org.springframework.integration.kafka.inbound.KafkaMessageSource;
import org.springframework.integration.test.util.TestUtils;
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
import org.springframework.test.annotation.DirtiesContext;
import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
@@ -65,6 +67,10 @@ public class KafkaInboundChannelAdapterParserTests {
.isSameAs(this.context.getBean("rebal"));
assertThat(TestUtils.getPropertyValue(this.source2, "topics")).isEqualTo(new String[] { "topic1", "topic2" });
DefaultKafkaConsumerFactory<?, ?> cf = TestUtils.getPropertyValue(this.source2, "consumerFactory",
DefaultKafkaConsumerFactory.class);
assertThat(cf).isSameAs(this.context.getBean("multiFetchConsumerFactory"));
assertThat(cf.getConfigurationProperties().get(ConsumerConfig.MAX_POLL_RECORDS_CONFIG)).isEqualTo("10");
}
}

View File

@@ -483,4 +483,69 @@ public class MessageSourceTests {
inOrder.verify(consumer).poll(Duration.of(2, ChronoUnit.SECONDS));
}
@SuppressWarnings({ "rawtypes", "unchecked" })
@Test
public void testAllowMulti() {
Consumer consumer = mock(Consumer.class);
TopicPartition topicPartition = new TopicPartition("foo", 0);
List<TopicPartition> assigned = Collections.singletonList(topicPartition);
willAnswer(i -> {
((ConsumerRebalanceListener) i.getArgument(1))
.onPartitionsAssigned(assigned);
return null;
}).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class));
ArgumentCaptor<Collection<TopicPartition>> partitions = ArgumentCaptor.forClass(Collection.class);
willDoNothing().given(consumer).pause(partitions.capture());
willDoNothing().given(consumer).resume(partitions.capture());
Map<TopicPartition, List<ConsumerRecord>> records = new LinkedHashMap<>();
records.put(topicPartition, Arrays.asList(
new ConsumerRecord("foo", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "foo"),
new ConsumerRecord("foo", 0, 1L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "bar"),
new ConsumerRecord("foo", 0, 2L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "baz"),
new ConsumerRecord("foo", 0, 3L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0, 0, 0, null, "qux")));
ConsumerRecords cr1 = new ConsumerRecords(records);
ConsumerRecords cr2 = new ConsumerRecords(Collections.emptyMap());
given(consumer.poll(any(Duration.class))).willReturn(cr1, cr2);
ConsumerFactory consumerFactory = mock(ConsumerFactory.class);
willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 4)).given(consumerFactory)
.getConfigurationProperties();
given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer);
KafkaMessageSource source = new KafkaMessageSource(consumerFactory, true, "foo");
source.setRawMessageHeader(true);
Message<?> received = source.receive();
assertThat(received).isNotNull();
assertThat(received.getHeaders().get(KafkaMessageSource.REMAINING_RECORDS, Integer.class)).isEqualTo(3);
StaticMessageHeaderAccessor.getAcknowledgmentCallback(received)
.acknowledge(AcknowledgmentCallback.Status.ACCEPT);
received = source.receive();
assertThat(received).isNotNull();
assertThat(received.getHeaders().get(KafkaMessageSource.REMAINING_RECORDS, Integer.class)).isEqualTo(2);
StaticMessageHeaderAccessor.getAcknowledgmentCallback(received)
.acknowledge(AcknowledgmentCallback.Status.ACCEPT);
received = source.receive();
assertThat(received).isNotNull();
assertThat(received.getHeaders().get(KafkaMessageSource.REMAINING_RECORDS, Integer.class)).isEqualTo(1);
StaticMessageHeaderAccessor.getAcknowledgmentCallback(received)
.acknowledge(AcknowledgmentCallback.Status.ACCEPT);
received = source.receive();
assertThat(received).isNotNull();
assertThat(received.getHeaders().get(KafkaMessageSource.REMAINING_RECORDS, Integer.class)).isEqualTo(0);
StaticMessageHeaderAccessor.getAcknowledgmentCallback(received)
.acknowledge(AcknowledgmentCallback.Status.ACCEPT);
received = source.receive();
assertThat(received).isNull();
source.destroy();
InOrder inOrder = inOrder(consumer);
inOrder.verify(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class));
inOrder.verify(consumer).poll(any(Duration.class));
inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(1L)));
inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(2L)));
inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(3L)));
inOrder.verify(consumer).commitSync(Collections.singletonMap(topicPartition, new OffsetAndMetadata(4L)));
inOrder.verify(consumer).poll(any(Duration.class));
inOrder.verify(consumer).close();
inOrder.verifyNoMoreInteractions();
}
}