Github user srdo commented on a diff in the pull request:
https://github.com/apache/storm/pull/1835#discussion_r93108651
--- Diff:
external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/internal/fetcher/ManualKafkaRecordsFetcher.java
---
@@ -0,0 +1,88 @@
+package org.apache.storm.kafka.spout.internal.fetcher;
+
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.storm.kafka.spout.KafkaPartitionReader;
+import org.apache.storm.kafka.spout.KafkaSpout;
+import org.apache.storm.kafka.spout.TopicPartitionComparator;
+
+import java.util.*;
+
+public class ManualKafkaRecordsFetcher<K, V> implements
KafkaRecordsFetcher<K, V> {
+ private static final Comparator<TopicPartition>
KAFKA_TOPIC_PARTITION_COMPARATOR = new TopicPartitionComparator();
+
+ private final KafkaConsumer<K, V> consumer;
+ private final int thisTaskIndex;
+ private final int totalTaskCount;
+ private final KafkaPartitionReader partitionReader;
+ private final KafkaSpout.Timer partitionRefreshTimer;
+ private final PartitionAssignmentChangeListener
partitionAssignmentChangeListener;
+ private Set<TopicPartition> myPartitions = Collections.emptySet();
+
+ public ManualKafkaRecordsFetcher(KafkaConsumer<K, V> consumer,
+ int thisTaskIndex,
+ int totalTaskCount,
+ KafkaPartitionReader partitionReader,
+ KafkaSpout.Timer
partitionRefreshTimer,
+ PartitionAssignmentChangeListener
partitionAssignmentChangeListener) {
+ this.consumer = consumer;
+ this.thisTaskIndex = thisTaskIndex;
+ this.totalTaskCount = totalTaskCount;
+ this.partitionReader = partitionReader;
+ this.partitionRefreshTimer = partitionRefreshTimer;
+ this.partitionAssignmentChangeListener =
partitionAssignmentChangeListener;
+
+ doRefreshMyPartitions();
+ }
+
+ private void refreshMyPartitionsIfNeed() {
+ if (!partitionRefreshTimer.isExpiredResetOnTrue()) {
+ return;
+ }
+
+ doRefreshMyPartitions();
+ }
+
+ private void doRefreshMyPartitions() {
+ List<TopicPartition> topicPartitions =
partitionReader.readPartitions(consumer);
+ Collections.sort(topicPartitions,
KAFKA_TOPIC_PARTITION_COMPARATOR);
+
+ Set<TopicPartition> curPartitions = new
HashSet<>(topicPartitions.size()/totalTaskCount+1);
+ for (int i=thisTaskIndex; i<topicPartitions.size();
i+=totalTaskCount) {
+ curPartitions.add(topicPartitions.get(i));
+ }
+
+ if (!myPartitions.equals(curPartitions) && myPartitions!=null) {
+
partitionAssignmentChangeListener.onPartitionAssignmentChange(myPartitions,
curPartitions);
+ }
+
+ myPartitions = curPartitions;
+
+ consumer.assign(myPartitions);
+ }
+
+ @Override
+ public ConsumerRecords<K, V> fetchRecords(long fetchTimeoutMs) {
+ refreshMyPartitionsIfNeed();
+
+ return consumer.poll(fetchTimeoutMs);
+ }
+
+ public interface PartitionAssignmentChangeListener {
+ void onPartitionAssignmentChange(Set<TopicPartition>
oldPartitions, Set<TopicPartition> newPartitions);
+ }
+
+ public static PartitionAssignmentChangeListener listenerOf(final
ConsumerRebalanceListener consumerRebalanceListener) {
+ return new PartitionAssignmentChangeListener() {
--- End diff --
You don't have to change it, but on master we're targeting Java 8, so
lambdas are an option for this kind of thing :)
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
