Add KafkaItemReader

Resolves BATCH-2764

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java

@@ -0,0 +1,188 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *          http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.springframework.batch.item.kafka;
+
+import java.time.Duration;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.ExecutionContext;
+import org.springframework.batch.item.ItemReader;
+import org.springframework.batch.item.ItemStreamException;
+import org.springframework.batch.item.kafka.support.AutoCommitOffsetsProvider;
+import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
+
+/**
+ * <p>
+ * An {@link ItemReader} implementation for Apache Kafka.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ *
+ */
+public class KafkaItemReader<K, V> extends AbstractItemCountingItemStreamItemReader<V> implements InitializingBean {
+
+	private static final String TOPIC_PARTITION_OFFSET = "topic.partition.offset";
+
+	private static final long DEFAULT_POLL_TIMEOUT = 50L;
+
+	private static final long MIN_ASSIGN_TIMEOUT = 2000L;
+
+	private final Supplier<Duration> assignTimeoutProvider = () -> Duration
+			.ofMillis(Math.max(this.pollTimeout.toMillis() * 20, MIN_ASSIGN_TIMEOUT));
+
+	private Duration pollTimeout = Duration.ofMillis(DEFAULT_POLL_TIMEOUT);
+
+	private List<TopicPartition> topicPartitions;
+
+	private List<String> topics;
+
+	private ConsumerFactory<K, V> consumerFactory;
+
+	private Consumer<K, V> consumer;
+
+	private OffsetsProvider offsetsProvider;
+
+	private AtomicBoolean assigned = new AtomicBoolean(false);
+
+	private Map<TopicPartition, Long> offsets;
+
+	private Iterator<ConsumerRecord<K, V>> records;
+
+	public KafkaItemReader() {
+		super();
+		setName(ClassUtils.getShortName(KafkaItemReader.class));
+	}
+
+	public void setPollTimeout(long pollTimeout) {
+		Assert.isTrue(pollTimeout >= 0, "'pollTimeout' must no be negative.");
+		this.pollTimeout = Duration.ofMillis(pollTimeout);
+	}
+
+	public void setTopicPartitions(List<TopicPartition> topicPartitions) {
+		this.topicPartitions = topicPartitions;
+	}
+
+	public void setTopics(List<String> topics) {
+		this.topics = topics;
+	}
+
+	public void setConsumerFactory(ConsumerFactory<K, V> consumerFactory) {
+		this.consumerFactory = consumerFactory;
+	}
+
+	public void setOffsetsProvider(OffsetsProvider offsetsProvider) {
+		this.offsetsProvider = offsetsProvider;
+	}
+
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		Assert.state(topicPartitions != null || topics != null, "Either 'topicPartitions' or 'topics' must be provided.");
+		Assert.state(topicPartitions == null || topics == null, "Both 'topicPartitions' and 'topics' cannot be specified together.");
+		Assert.notNull(consumerFactory, "'consumerFactory' must not be null.");
+		Assert.notNull(offsetsProvider, "'offsetsProvider' must not be null.");
+		if  (consumerFactory.isAutoCommit()) {
+			Assert.state(offsetsProvider instanceof AutoCommitOffsetsProvider, "'AutoCommitOffsetsProvider' must be used if 'consumerFactory' is set to auto commit.");
+		}
+	}
+
+	@Override
+	@SuppressWarnings("unchecked")
+	public void open(ExecutionContext executionContext) throws ItemStreamException {
+		super.open(executionContext);
+		try {
+			if (isSaveState() && executionContext.containsKey(TOPIC_PARTITION_OFFSET)) {
+				offsets = (Map<TopicPartition, Long>) executionContext.get(TOPIC_PARTITION_OFFSET);
+			}
+			else {
+				offsets = offsetsProvider.get(topicPartitions);
+			}
+
+			if (offsets != null && !offsets.isEmpty()) {
+				offsets.forEach(consumer::seek);
+			}
+		}
+		catch (Exception e) {
+			throw new ItemStreamException("Failed to initialize the reader", e);
+		}
+	}
+
+	@Override
+	protected void doOpen() throws Exception {
+		consumer = consumerFactory.createConsumer();
+		if (topics != null) {
+			topicPartitions = topics.stream()
+					.flatMap(topic -> consumer.partitionsFor(topic).stream())
+					.map(partitionInfo -> new TopicPartition(partitionInfo.topic(), partitionInfo.partition()))
+					.collect(Collectors.toList());
+		}
+		consumer.assign(topicPartitions);
+		offsetsProvider.setConsumer(consumer);
+	}
+
+	@Override
+	protected void jumpToItem(int itemIndex) throws Exception {
+	}
+
+	@Override
+	protected V doRead() throws Exception {
+		if (records == null || !records.hasNext()) {
+			records = doPoll();
+		}
+		if (records.hasNext()) {
+			ConsumerRecord<K, V> record = records.next();
+			offsets.put(new TopicPartition(record.topic(), record.partition()), record.offset());
+			return record.value();
+		}
+		return null;
+	}
+
+	protected Iterator<ConsumerRecord<K, V>> doPoll() {
+		return consumer.poll(assigned.getAndSet(true) ? pollTimeout : assignTimeoutProvider.get()).iterator();
+	}
+
+	@Override
+	public void update(ExecutionContext executionContext) throws ItemStreamException {
+		super.update(executionContext);
+		if (isSaveState()) {
+			Assert.notNull(executionContext, "ExecutionContext must not be null");
+			executionContext.put(TOPIC_PARTITION_OFFSET, offsets);
+		}
+	}
+
+	@Override
+	protected void doClose() throws Exception {
+		records = null;
+		offsets = null;
+		assigned.set(false);
+		if (consumer != null) {
+			consumer.close();
+		}
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/OffsetsProvider.java

@@ -0,0 +1,37 @@
+package org.springframework.batch.item.kafka;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+
+/**
+ * A convenient strategy for providing initial externally stored offsets for the {@link KafkaItemReader} to
+ * seeks to.
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public interface OffsetsProvider {
+
+	/**
+	 * <p>
+	 * Map of assigned topic-partitions and the offset, or read position, at which the {@link KafkaItemReader} should
+	 * start.
+	 * </p>
+	 *
+	 * @param topicPartitions list of assigned topic-partitions to get offset for
+	 * @return map of offset by topic-partition
+	 */
+	Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions);
+
+	/**
+	 * <p>
+	 * Inject a {@link Consumer} that can be used to fetch internally stored offsets and/or topic-partition assignment.
+	 * </p>
+	 *
+	 * @param consumer the {@link Consumer} to set
+	 */
+	void setConsumer(Consumer<?, ?> consumer);
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *          http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.springframework.batch.item.kafka.builder;
+
+import java.util.List;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.KafkaItemReader;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+import org.springframework.batch.item.kafka.support.AutoCommitOffsetsProvider;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.util.Assert;
+
+/**
+ * A builder implementation for the {@link KafkaItemReader}
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ * @see KafkaItemReader
+ */
+public class KafkaItemReaderBuilder<K, V> {
+
+	private ConsumerFactory<K, V> consumerFactory;
+
+	private List<TopicPartition> topicPartitions;
+
+	private List<String> topics;
+
+	private OffsetsProvider offsetsProvider = new AutoCommitOffsetsProvider();
+
+	private long pollTimeout = 50L;
+
+	private boolean saveState = true;
+
+	private String name;
+
+	private int maxItemCount = Integer.MAX_VALUE;
+
+	/**
+	 * Configure if the state of the {@link org.springframework.batch.item.ItemStreamSupport} should be persisted within
+	 * the {@link org.springframework.batch.item.ExecutionContext} for restart purposes.
+	 *
+	 * @param saveState defaults to true
+	 * @return The current instance of the builder.
+	 */
+	public KafkaItemReaderBuilder<K, V> saveState(boolean saveState) {
+		this.saveState = saveState;
+		return this;
+	}
+
+	/**
+	 * The name used to calculate the key within the {@link org.springframework.batch.item.ExecutionContext}. Required
+	 * if {@link #saveState(boolean)} is set to true.
+	 *
+	 * @param name name of the reader instance
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.ItemStreamSupport#setName(String)
+	 */
+	public KafkaItemReaderBuilder<K, V> name(String name) {
+		this.name = name;
+		return this;
+	}
+
+	/**
+	 * Configure the max number of items to be read.
+	 *
+	 * @param maxItemCount the max items to be read
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader#setMaxItemCount(int)
+	 */
+	public KafkaItemReaderBuilder<K, V> maxItemCount(int maxItemCount) {
+		this.maxItemCount = maxItemCount;
+		return this;
+	}
+
+	/**
+	 * The {@link ConsumerFactory} implementation to produce a new {@link Consumer} instance for the reader.
+	 *
+	 * @param consumerFactory
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setConsumerFactory(ConsumerFactory)
+	 */
+	public KafkaItemReaderBuilder<K, V> consumerFactory(ConsumerFactory<K, V> consumerFactory) {
+		this.consumerFactory = consumerFactory;
+		return this;
+	}
+
+	/**
+	 * A list of {@link TopicPartition}s to manually assign the consumer.
+	 *
+	 * @param topicPartitions list of partitions to assign the consumer
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setTopicPartitions(List)
+	 */
+	public KafkaItemReaderBuilder<K, V> topicPartitions(List<TopicPartition> topicPartitions) {
+		this.topicPartitions = topicPartitions;
+		return this;
+	}
+
+	/**
+	 * A list of topics to manually assign the consumer.
+	 *
+	 * @param topics list of topics to assign the consumer
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setTopics(List)
+	 */
+	public KafkaItemReaderBuilder<K, V> topics(List<String> topics) {
+		this.topics = topics;
+		return this;
+	}
+
+	/**
+	 * The {@link OffsetsProvider} implementation to provide initial offsets.
+	 *
+	 * @param offsetsProvider
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setOffsetsProvider(OffsetsProvider)
+	 */
+	public KafkaItemReaderBuilder<K,V> offsetsProvider(OffsetsProvider offsetsProvider)  {
+		this.offsetsProvider = offsetsProvider;
+		return this;
+	}
+
+	/**
+	 * Set the pollTimeout for the poll() operations.
+	 *
+	 * @param pollTimeout default to 50ms
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setPollTimeout(long)
+	 */
+	public KafkaItemReaderBuilder<K, V> pollTimeout(long pollTimeout) {
+		this.pollTimeout = pollTimeout;
+		return this;
+	}
+
+	public KafkaItemReader<K, V> build() {
+		if (this.saveState) {
+			Assert.hasText(this.name, "A name is required when saveState is set to true");
+		}
+		Assert.state(this.topicPartitions != null || this.topics != null, "Either 'topicPartitions' or 'topics' must be provided.");
+		Assert.state(this.topicPartitions == null || this.topics == null, "Both 'topicPartitions' and 'topics' cannot be specified together.");
+		Assert.isTrue(this.pollTimeout >= 0, "pollTimeout must not be negative.");
+		Assert.notNull(this.consumerFactory, "'consumerFactory' must not be null.");
+		Assert.notNull(this.offsetsProvider, "'offsetsProvider' must not be null.");
+		if  (this.consumerFactory.isAutoCommit()) {
+			Assert.state(this.offsetsProvider instanceof AutoCommitOffsetsProvider, "'AutoCommitOffsetsProvider' must be used if 'consumerFactory' is set to auto commit.");
+		}
+
+		KafkaItemReader<K, V> reader = new KafkaItemReader<>();
+		reader.setConsumerFactory(this.consumerFactory);
+		reader.setTopicPartitions(this.topicPartitions);
+		reader.setTopics(this.topics);
+		reader.setOffsetsProvider(this.offsetsProvider);
+		reader.setPollTimeout(this.pollTimeout);
+		reader.setSaveState(this.saveState);
+		reader.setName(this.name);
+		reader.setMaxItemCount(this.maxItemCount);
+		return reader;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/AutoCommitOffsetsProvider.java

@@ -0,0 +1,32 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * Noop implementation of {@link OffsetsProvider} to use with automatic offset committing. This is the only
+ * implementation allowed when the KafkaConsumer has 'enable.auto.commit' set to true. See 'auto.offset.reset' config
+ * for initial offset strategy options.
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public class AutoCommitOffsetsProvider implements OffsetsProvider {
+
+	/**
+	 * @return null, preventing a call to
+	 * {@link org.apache.kafka.clients.consumer.Consumer#seek(org.apache.kafka.common.TopicPartition, long)}
+	 */
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return null;
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/BeginningOffsetsProvider.java

@@ -0,0 +1,34 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns the earliest offsets for the given topic-partitions.
+ * Equivalent of using 'auto.offset.reset' set to 'earliest' with 'enable.auto.commit' set to false.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @see org.apache.kafka.clients.consumer.KafkaConsumer#beginningOffsets(Collection)
+ * @since 4.2
+ */
+public class BeginningOffsetsProvider implements OffsetsProvider {
+
+	private Consumer<?, ?> consumer;
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return this.consumer.beginningOffsets(topicPartitions);
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+		this.consumer = consumer;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/SimpleOffsetsProvider.java

@@ -0,0 +1,34 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns default or provided offsets for the given topic-partitions.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public class SimpleOffsetsProvider implements OffsetsProvider {
+
+	private Map<TopicPartition, Long> offsets;
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return this.offsets;
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+	}
+
+	public void setOffsets(Map<TopicPartition, Long> offsets) {
+		this.offsets = offsets;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/TimestampOffsetsProvider.java

@@ -0,0 +1,45 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns offsets for the given topic-partitions by timestamp.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @see org.apache.kafka.clients.consumer.KafkaConsumer#offsetsForTimes(java.util.Map)
+ * @since 4.2
+ */
+public class TimestampOffsetsProvider implements OffsetsProvider {
+
+	private final Long timestampToSearch;
+
+	private Consumer<?, ?> consumer;
+
+	public TimestampOffsetsProvider(Long timestampToSearch) {
+		this.timestampToSearch = timestampToSearch;
+	}
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		Map<TopicPartition, Long> timestampsToSearch = topicPartitions.stream()
+				.collect(Collectors.toMap(Function.identity(), topicPartition -> timestampToSearch));
+		return consumer.offsetsForTimes(timestampsToSearch).entrySet().stream()
+				.filter(entry -> entry.getValue() != null)
+				.collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().offset()));
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+		this.consumer = consumer;
+	}
+}