crunch-kafka/src/main/java/org/apache/crunch/kafka/KafkaRecordsIterable.java - crunch - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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
  * <p>
  * http://www.apache.org/licenses/LICENSE-2.0
  * <p>
  * 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.apache.crunch.kafka;

 import org.apache.crunch.Pair;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.RetriableException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Properties;
 import java.util.Set;


 class KafkaRecordsIterable<K, V> implements Iterable<Pair<K, V>> {

   /**
    * Logger
    */
   private static final Logger LOG = LoggerFactory.getLogger(KafkaRecordsIterable.class);

   /**
    * The Kafka consumer responsible for retrieving messages.
    */
   private final Consumer<K, V> consumer;

   /**
    * The starting positions of the iterable for the topic.
    */
   private final Map<TopicPartition, Pair<Long, Long>> offsets;

   /**
    * Tracks if the iterable is empty.
    */
   private final boolean isEmpty;

   /**
    * The poll time between each request to Kafka
    */
   private final long scanPollTime;

   private final int maxRetryAttempts;

   /**
    * Creates the iterable that will pull values for a collection of topics using the provided {@code consumer} between
    * the {@code startOffsets} and {@code stopOffsets}.
    * @param consumer The consumer for pulling the data from Kafka.  The consumer will be closed automatically once all
    *                 of the records have been consumed.
    * @param offsets offsets for pulling data
    * @param properties properties for tweaking the behavior of the iterable.
    * @throws IllegalArgumentException if any of the arguments are {@code null} or empty.
    */
   public KafkaRecordsIterable(Consumer<K, V> consumer, Map<TopicPartition, Pair<Long, Long>> offsets,
                               Properties properties) {
     if (consumer == null) {
       throw new IllegalArgumentException("The 'consumer' cannot be 'null'.");
     }
     this.consumer = consumer;

     if (properties == null) {
       throw new IllegalArgumentException("The 'properties' cannot be 'null'.");
     }

     String retryString = properties.getProperty(KafkaUtils.KAFKA_RETRY_ATTEMPTS_KEY,
         KafkaUtils.KAFKA_RETRY_ATTEMPTS_DEFAULT_STRING);
     maxRetryAttempts = Integer.parseInt(retryString);

     if (offsets == null || offsets.isEmpty()) {
       throw new IllegalArgumentException("The 'offsets' cannot 'null' or empty.");
     }

     //filter out any topics and partitions that do not have offset ranges that will produce data.
     Map<TopicPartition, Pair<Long, Long>> filteredOffsets = new HashMap<>();
     for (Map.Entry<TopicPartition, Pair<Long, Long>> entry : offsets.entrySet()) {
       Pair<Long, Long> value = entry.getValue();
       //if start is less than one less than stop then there is data to be had
       if(value.first() < value.second()){
         filteredOffsets.put(entry.getKey(), value);
       }else{
         LOG.debug("Removing offsets for {} because start is not less than the end offset.", entry.getKey());
       }
     }

     //check to make sure that based on the offsets there is data to retrieve, otherwise false.
     //there will be data if the start offsets are less than stop offsets
     isEmpty = filteredOffsets.isEmpty();
     if (isEmpty) {
       LOG.warn("Iterable for Kafka for is empty because offsets are empty.");
     }

     //assign this
     this.offsets = filteredOffsets;

     scanPollTime = Long.parseLong(properties.getProperty(KafkaSource.CONSUMER_POLL_TIMEOUT_KEY,
         Long.toString(KafkaSource.CONSUMER_POLL_TIMEOUT_DEFAULT)));
   }

   @Override
   public Iterator<Pair<K, V>> iterator() {
     if (isEmpty) {
       LOG.debug("Returning empty iterator since offsets align.");
       return Collections.emptyIterator();
     }
     //Assign consumer to all of the partitions
     LOG.debug("Assigning topics and partitions and seeking to start offsets.");

     consumer.assign(new LinkedList<>(offsets.keySet()));
     //hack so maybe look at removing this
     consumer.poll(0);
     for (Map.Entry<TopicPartition, Pair<Long, Long>> entry : offsets.entrySet()) {
       consumer.seek(entry.getKey(), entry.getValue().first());
     }

     return new RecordsIterator<K, V>(consumer, offsets, scanPollTime, maxRetryAttempts);
   }

   private static class RecordsIterator<K, V> implements Iterator<Pair<K, V>> {

     private final Consumer<K, V> consumer;
     private final Map<TopicPartition, Pair<Long, Long>> offsets;
     private final long pollTime;
     private final int maxNumAttempts;
     private ConsumerRecords<K, V> records;
     private Iterator<ConsumerRecord<K, V>> currentIterator;
     private final Set<TopicPartition> remainingPartitions;

     private Pair<K, V> next;

     public RecordsIterator(Consumer<K, V> consumer,
                            Map<TopicPartition, Pair<Long, Long>> offsets, long pollTime, int maxNumRetries) {
       this.consumer = consumer;
       remainingPartitions = new HashSet<>(offsets.keySet());
       this.offsets = offsets;
       this.pollTime = pollTime;
       this.maxNumAttempts = maxNumRetries;
     }

     @Override
     public boolean hasNext() {
       if (next != null)
         return true;

       //if partitions to consume then pull next value
       if (remainingPartitions.size() > 0) {
         next = getNext();
       }

       return next != null;
     }

     @Override
     public Pair<K, V> next() {
       if (next == null) {
         next = getNext();
       }

       if (next != null) {
         Pair<K, V> returnedNext = next;
         //prime for next call
         next = getNext();
         //return the current next
         return returnedNext;
       } else {
         throw new NoSuchElementException("No more elements.");
       }
     }

     @Override
     public void remove() {
       throw new UnsupportedOperationException("remove is not supported.");
     }

     /**
      * Gets the current iterator.
      *
      * @return the current iterator or {@code null} if there are no more values to consume.
      */
     private Iterator<ConsumerRecord<K, V>> getIterator() {
       if (!remainingPartitions.isEmpty()) {
         if (currentIterator != null && currentIterator.hasNext()) {
           return currentIterator;
         }
         LOG.debug("Retrieving next set of records.");
         int numTries = 0;
         boolean notSuccess = false;
         while(!notSuccess && numTries < maxNumAttempts) {
           try {
             records = consumer.poll(pollTime);
             notSuccess = true;
           }catch(RetriableException re){
             numTries++;
             if(numTries < maxNumAttempts) {
               LOG.warn("Error pulling messages from Kafka. Retrying with attempt {}", numTries, re);
             }else{
               LOG.error("Error pulling messages from Kafka. Exceeded maximum number of attempts {}", maxNumAttempts, re);
               throw re;
             }
           }
         }
         if (records == null || records.isEmpty()) {
           LOG.debug("Retrieved empty records.");
           currentIterator = null;
           return null;
         }
         currentIterator = records.iterator();
         return currentIterator;
       }

       LOG.debug("No more partitions to consume therefore not retrieving any more records.");
       return null;
     }

     /**
      * Internal method for retrieving the next value to retrieve.
      *
      * @return {@code null} if there are no more values to retrieve otherwise the next event.
      */
     private Pair<K, V> getNext() {
       while (!remainingPartitions.isEmpty()) {
         Iterator<ConsumerRecord<K, V>> iterator = getIterator();

         while (iterator != null && iterator.hasNext()) {
           ConsumerRecord<K, V> record = iterator.next();
           TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition());
           long offset = record.offset();

           if (withinRange(topicPartition, offset)) {
             LOG.debug("Retrieving value for {} with offset {}.", topicPartition, offset);
             return Pair.of(record.key(), record.value());
           }
           LOG.debug("Value for {} with offset {} is outside of range skipping.", topicPartition, offset);
         }
       }

       LOG.debug("Closing the consumer because there are no more remaining partitions.");
       consumer.close();

       LOG.debug("Consumed data from all partitions.");
       return null;

     }

     /**
      * Checks whether the value for {@code topicPartition} with an {@code offset} is within scan range.  If
      * the value is not then {@code false} is returned otherwise {@code true}.
      *
      * @param topicPartion The partition for the offset
      * @param offset the offset in the partition
      * @return {@code true} if the value is within the expected consumption range, otherwise {@code false}.
      */
     private boolean withinRange(TopicPartition topicPartion, long offset) {
       long endOffset = offsets.get(topicPartion).second();
       //end offsets are one higher than the last written value.
       boolean emit = offset < endOffset;
       if (offset >= endOffset - 1) {
         if (LOG.isDebugEnabled()) {
           LOG.debug("Completed consuming partition {} with offset {} and ending offset {}.",
               new Object[]{topicPartion, offset, endOffset});
         }
         remainingPartitions.remove(topicPartion);
         consumer.pause(topicPartion);
       }
       LOG.debug("Value for partition {} and offset {} is within range.", topicPartion, offset);
       return emit;
     }
   }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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
	* <p>
	* http://www.apache.org/licenses/LICENSE-2.0
	* <p>
	* 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.apache.crunch.kafka;

	import org.apache.crunch.Pair;
	import org.apache.kafka.clients.consumer.Consumer;
	import org.apache.kafka.clients.consumer.ConsumerRecord;
	import org.apache.kafka.clients.consumer.ConsumerRecords;
	import org.apache.kafka.common.TopicPartition;
	import org.apache.kafka.common.errors.RetriableException;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.util.Collections;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.Iterator;
	import java.util.LinkedList;
	import java.util.Map;
	import java.util.NoSuchElementException;
	import java.util.Properties;
	import java.util.Set;


	class KafkaRecordsIterable<K, V> implements Iterable<Pair<K, V>> {

	/**
	* Logger
	*/
	private static final Logger LOG = LoggerFactory.getLogger(KafkaRecordsIterable.class);

	/**
	* The Kafka consumer responsible for retrieving messages.
	*/
	private final Consumer<K, V> consumer;

	/**
	* The starting positions of the iterable for the topic.
	*/
	private final Map<TopicPartition, Pair<Long, Long>> offsets;

	/**
	* Tracks if the iterable is empty.
	*/
	private final boolean isEmpty;

	/**
	* The poll time between each request to Kafka
	*/
	private final long scanPollTime;

	private final int maxRetryAttempts;

	/**
	* Creates the iterable that will pull values for a collection of topics using the provided {@code consumer} between
	* the {@code startOffsets} and {@code stopOffsets}.
	* @param consumer The consumer for pulling the data from Kafka. The consumer will be closed automatically once all
	* of the records have been consumed.
	* @param offsets offsets for pulling data
	* @param properties properties for tweaking the behavior of the iterable.
	* @throws IllegalArgumentException if any of the arguments are {@code null} or empty.
	*/
	public KafkaRecordsIterable(Consumer<K, V> consumer, Map<TopicPartition, Pair<Long, Long>> offsets,
	Properties properties) {
	if (consumer == null) {
	throw new IllegalArgumentException("The 'consumer' cannot be 'null'.");
	}
	this.consumer = consumer;

	if (properties == null) {
	throw new IllegalArgumentException("The 'properties' cannot be 'null'.");
	}

	String retryString = properties.getProperty(KafkaUtils.KAFKA_RETRY_ATTEMPTS_KEY,
	KafkaUtils.KAFKA_RETRY_ATTEMPTS_DEFAULT_STRING);
	maxRetryAttempts = Integer.parseInt(retryString);

	if (offsets == null \|\| offsets.isEmpty()) {
	throw new IllegalArgumentException("The 'offsets' cannot 'null' or empty.");
	}

	//filter out any topics and partitions that do not have offset ranges that will produce data.
	Map<TopicPartition, Pair<Long, Long>> filteredOffsets = new HashMap<>();
	for (Map.Entry<TopicPartition, Pair<Long, Long>> entry : offsets.entrySet()) {
	Pair<Long, Long> value = entry.getValue();
	//if start is less than one less than stop then there is data to be had
	if(value.first() < value.second()){
	filteredOffsets.put(entry.getKey(), value);
	}else{
	LOG.debug("Removing offsets for {} because start is not less than the end offset.", entry.getKey());
	}
	}

	//check to make sure that based on the offsets there is data to retrieve, otherwise false.
	//there will be data if the start offsets are less than stop offsets
	isEmpty = filteredOffsets.isEmpty();
	if (isEmpty) {
	LOG.warn("Iterable for Kafka for is empty because offsets are empty.");
	}

	//assign this
	this.offsets = filteredOffsets;

	scanPollTime = Long.parseLong(properties.getProperty(KafkaSource.CONSUMER_POLL_TIMEOUT_KEY,
	Long.toString(KafkaSource.CONSUMER_POLL_TIMEOUT_DEFAULT)));
	}

	@Override
	public Iterator<Pair<K, V>> iterator() {
	if (isEmpty) {
	LOG.debug("Returning empty iterator since offsets align.");
	return Collections.emptyIterator();
	}
	//Assign consumer to all of the partitions
	LOG.debug("Assigning topics and partitions and seeking to start offsets.");

	consumer.assign(new LinkedList<>(offsets.keySet()));
	//hack so maybe look at removing this
	consumer.poll(0);
	for (Map.Entry<TopicPartition, Pair<Long, Long>> entry : offsets.entrySet()) {
	consumer.seek(entry.getKey(), entry.getValue().first());
	}

	return new RecordsIterator<K, V>(consumer, offsets, scanPollTime, maxRetryAttempts);
	}

	private static class RecordsIterator<K, V> implements Iterator<Pair<K, V>> {

	private final Consumer<K, V> consumer;
	private final Map<TopicPartition, Pair<Long, Long>> offsets;
	private final long pollTime;
	private final int maxNumAttempts;
	private ConsumerRecords<K, V> records;
	private Iterator<ConsumerRecord<K, V>> currentIterator;
	private final Set<TopicPartition> remainingPartitions;

	private Pair<K, V> next;

	public RecordsIterator(Consumer<K, V> consumer,
	Map<TopicPartition, Pair<Long, Long>> offsets, long pollTime, int maxNumRetries) {
	this.consumer = consumer;
	remainingPartitions = new HashSet<>(offsets.keySet());
	this.offsets = offsets;
	this.pollTime = pollTime;
	this.maxNumAttempts = maxNumRetries;
	}

	@Override
	public boolean hasNext() {
	if (next != null)
	return true;

	//if partitions to consume then pull next value
	if (remainingPartitions.size() > 0) {
	next = getNext();
	}

	return next != null;
	}

	@Override
	public Pair<K, V> next() {
	if (next == null) {
	next = getNext();
	}

	if (next != null) {
	Pair<K, V> returnedNext = next;
	//prime for next call
	next = getNext();
	//return the current next
	return returnedNext;
	} else {
	throw new NoSuchElementException("No more elements.");
	}
	}

	@Override
	public void remove() {
	throw new UnsupportedOperationException("remove is not supported.");
	}

	/**
	* Gets the current iterator.
	*
	* @return the current iterator or {@code null} if there are no more values to consume.
	*/
	private Iterator<ConsumerRecord<K, V>> getIterator() {
	if (!remainingPartitions.isEmpty()) {
	if (currentIterator != null && currentIterator.hasNext()) {
	return currentIterator;
	}
	LOG.debug("Retrieving next set of records.");
	int numTries = 0;
	boolean notSuccess = false;
	while(!notSuccess && numTries < maxNumAttempts) {
	try {
	records = consumer.poll(pollTime);
	notSuccess = true;
	}catch(RetriableException re){
	numTries++;
	if(numTries < maxNumAttempts) {
	LOG.warn("Error pulling messages from Kafka. Retrying with attempt {}", numTries, re);
	}else{
	LOG.error("Error pulling messages from Kafka. Exceeded maximum number of attempts {}", maxNumAttempts, re);
	throw re;
	}
	}
	}
	if (records == null \|\| records.isEmpty()) {
	LOG.debug("Retrieved empty records.");
	currentIterator = null;
	return null;
	}
	currentIterator = records.iterator();
	return currentIterator;
	}

	LOG.debug("No more partitions to consume therefore not retrieving any more records.");
	return null;
	}

	/**
	* Internal method for retrieving the next value to retrieve.
	*
	* @return {@code null} if there are no more values to retrieve otherwise the next event.
	*/
	private Pair<K, V> getNext() {
	while (!remainingPartitions.isEmpty()) {
	Iterator<ConsumerRecord<K, V>> iterator = getIterator();

	while (iterator != null && iterator.hasNext()) {
	ConsumerRecord<K, V> record = iterator.next();
	TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition());
	long offset = record.offset();

	if (withinRange(topicPartition, offset)) {
	LOG.debug("Retrieving value for {} with offset {}.", topicPartition, offset);
	return Pair.of(record.key(), record.value());
	}
	LOG.debug("Value for {} with offset {} is outside of range skipping.", topicPartition, offset);
	}
	}

	LOG.debug("Closing the consumer because there are no more remaining partitions.");
	consumer.close();

	LOG.debug("Consumed data from all partitions.");
	return null;

	}

	/**
	* Checks whether the value for {@code topicPartition} with an {@code offset} is within scan range. If
	* the value is not then {@code false} is returned otherwise {@code true}.
	*
	* @param topicPartion The partition for the offset
	* @param offset the offset in the partition
	* @return {@code true} if the value is within the expected consumption range, otherwise {@code false}.
	*/
	private boolean withinRange(TopicPartition topicPartion, long offset) {
	long endOffset = offsets.get(topicPartion).second();
	//end offsets are one higher than the last written value.
	boolean emit = offset < endOffset;
	if (offset >= endOffset - 1) {
	if (LOG.isDebugEnabled()) {
	LOG.debug("Completed consuming partition {} with offset {} and ending offset {}.",
	new Object[]{topicPartion, offset, endOffset});
	}
	remainingPartitions.remove(topicPartion);
	consumer.pause(topicPartion);
	}
	LOG.debug("Value for partition {} and offset {} is within range.", topicPartion, offset);
	return emit;
	}
	}
	}