kafka/kafka-common/src/main/java/org/apache/apex/malhar/kafka/AbstractKafkaPartitioner.java - apex-malhar - 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
  *
  *   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.apache.apex.malhar.kafka;

 import java.io.IOException;
 import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;

 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import org.apache.apex.malhar.lib.util.KryoCloneUtils;
 import org.apache.hadoop.classification.InterfaceStability;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;

 import com.google.common.base.Joiner;

 import com.datatorrent.api.DefaultPartition;
 import com.datatorrent.api.Partitioner;
 import com.datatorrent.api.StatsListener;

 /**
  * Abstract partitioner used to manage the partitions of kafka input operator.
  * It use a number of kafka consumers(one for each cluster) to get the latest partition metadata for topics that
  * the consumer subscribes and expose those to subclass which implements the assign method
  *
  * The partitioner is always stateless.
  *
  * @since 3.3.0
  */
 @InterfaceStability.Evolving
 public abstract class AbstractKafkaPartitioner implements Partitioner<AbstractKafkaInputOperator>, StatsListener
 {

   private static final Logger logger = LoggerFactory.getLogger(AbstractKafkaPartitioner.class);

   private static final String META_CONSUMER_GROUP_NAME = AbstractKafkaInputOperator.class.getName() + "META_GROUP";

   protected final String[] clusters;

   protected final String[] topics;

   protected final AbstractKafkaInputOperator prototypeOperator;

   private ArrayList<KafkaConsumer<byte[], byte[]>> metadataRefreshClients;

   // prevent null
   private final List<Set<AbstractKafkaPartitioner.PartitionMeta>> currentPartitions = new LinkedList<>();

   public AbstractKafkaPartitioner(String[] clusters, String[] topics, AbstractKafkaInputOperator prototypeOperator)
   {
     this.clusters = clusters;
     this.topics = topics;
     this.prototypeOperator = prototypeOperator;
   }

   abstract List<Set<PartitionMeta>> assign(Map<String, Map<String, List<PartitionInfo>>> metadata);

   @Override
   public Collection<Partition<AbstractKafkaInputOperator>> definePartitions(
       Collection<Partition<AbstractKafkaInputOperator>> collection, PartitioningContext partitioningContext)
   {
     initMetadataClients();

     Map<String, Map<String, List<PartitionInfo>>> metadata = new HashMap<>();

     try {
       for (int i = 0; i < clusters.length; i++) {
         metadata.put(clusters[i], new HashMap<String, List<PartitionInfo>>());
         for (String topic : topics) {
           //try several time if partitionsFor(topic) returns null or throws exception.
           //partitionsFor(topic) will return null if the topic is invalid or hasn't completed
           int tryTime = 10;
           while (tryTime-- > 0) {
             try {
               List<PartitionInfo> ptis = metadataRefreshClients.get(i).partitionsFor(topic);
               if (ptis != null) {
                 if (logger.isDebugEnabled()) {
                   logger.debug("Partition metadata for topic {} : {}", topic, Joiner.on(';').join(ptis));
                 }
                 metadata.get(clusters[i]).put(topic, ptis);
                 break;
               }

               logger.warn("Partition metadata for topic {} is null. retrying...", topic);

             } catch (Exception e) {
               logger.warn("Got Exception when trying get partition info for topic {}.", topic, e);
             }

             try {
               Thread.sleep(100);
             } catch (Exception e1) {
               //ignore
             }
           } //end while

           if (tryTime == 0) {
             throw new RuntimeException(
                 "Get partition info for topic completely failed. Please check the log file. topic name: " + topic);
           }
         }
       }
     } finally {
       closeClients();
     }

     List<Set<AbstractKafkaPartitioner.PartitionMeta>> parts = null;
     try {
       parts = assign(metadata);
     } catch (Exception e) {
       logger.error("assign() exception.", e);
       e.printStackTrace();
     }

     if (currentPartitions == parts || currentPartitions.equals(parts)) {
       logger.debug("No partition change found");
       return collection;
     } else {
       logger.info("Partition change detected: ");
       currentPartitions.clear();
       currentPartitions.addAll(parts);
       int i = 0;
       List<Partition<AbstractKafkaInputOperator>> result = new LinkedList<>();
       for (Iterator<Partition<AbstractKafkaInputOperator>> iter = collection.iterator(); iter.hasNext(); ) {
         Partition<AbstractKafkaInputOperator> nextPartition = iter.next();
         if (parts.remove(nextPartition.getPartitionedInstance().assignment())) {
           if (logger.isInfoEnabled()) {
             logger.info("[Existing] Partition {} with assignment {} ", i,
                 Joiner.on(';').join(nextPartition.getPartitionedInstance().assignment()));
           }
           result.add(nextPartition);
           i++;
         }
       }

       for (Set<AbstractKafkaPartitioner.PartitionMeta> partitionAssignment : parts) {
         if (logger.isInfoEnabled()) {
           logger.info("[New] Partition {} with assignment {} ", i,
               Joiner.on(';').join(partitionAssignment));
         }
         result.add(createPartition(partitionAssignment));
         i++;
       }

       return result;
     }
   }

   protected void closeClients()
   {
     for (KafkaConsumer<byte[], byte[]> consume : metadataRefreshClients) {
       consume.close();
     }
     metadataRefreshClients = null;
   }

   @Override
   public void partitioned(Map<Integer, Partition<AbstractKafkaInputOperator>> map)
   {

   }

   @Override
   public Response processStats(BatchedOperatorStats batchedOperatorStats)
   {
     Response response = new Response();
     response.repartitionRequired = true;
     return response;
   }

   protected Partitioner.Partition<AbstractKafkaInputOperator> createPartition(
       Set<AbstractKafkaPartitioner.PartitionMeta> partitionAssignment)
   {
     Partitioner.Partition<AbstractKafkaInputOperator> p =
         new DefaultPartition<AbstractKafkaInputOperator>(KryoCloneUtils.cloneObject(prototypeOperator));
     p.getPartitionedInstance().assign(partitionAssignment);
     return p;
   }

   /**
    *
    */
   private void initMetadataClients()
   {
     if (metadataRefreshClients != null && metadataRefreshClients.size() == clusters.length) {
       // The metadata client is active
       return;
     }

     if (clusters == null || clusters.length == 0) {
       throw new IllegalStateException("clusters can not be null");
     }

     metadataRefreshClients = new ArrayList<>(clusters.length);
     int index = 0;
     for (String c : clusters) {
       Properties prop = prototypeOperator.getConsumerProps();
       prop.put("group.id", META_CONSUMER_GROUP_NAME);
       prop.put("bootstrap.servers", c);
       prop.put("key.deserializer", ByteArrayDeserializer.class.getName());
       prop.put("value.deserializer", ByteArrayDeserializer.class.getName());
       prop.put("enable.auto.commit", "false");
       if (logger.isInfoEnabled()) {
         logger.info("Consumer Properties :  {} ", getPropertyAsString(prop));
       }
       metadataRefreshClients.add(index++, new KafkaConsumer<byte[], byte[]>(prop));
     }
   }

   /**
    * Converts the property list (key and element pairs) to String format
    * This format is used to print to a Stream for debugging.
    * @param prop
    * @return String
    */
   private String getPropertyAsString(Properties prop)
   {
     StringWriter writer = new StringWriter();
     try {
       prop.store(writer, "");
     } catch (IOException e) {
       logger.error("Cannot retrieve consumer properties for Logging : {}", e.getMessage());
     }
     return writer.getBuffer().toString();
   }

   /**
    * The key object used in the assignment map for each operator
    */
   public static class PartitionMeta
   {

     public PartitionMeta()
     {
     }

     public PartitionMeta(String cluster, String topic, int partitionId)
     {
       this.cluster = cluster;
       this.topic = topic;
       this.partitionId = partitionId;
       this.topicPartition = new TopicPartition(topic, partitionId);
     }

     private String cluster;

     private transient TopicPartition topicPartition;

     private String topic;

     private int partitionId;

     public String getCluster()
     {
       return cluster;
     }

     public int getPartitionId()
     {
       return partitionId;
     }

     public String getTopic()
     {
       return topic;
     }

     public TopicPartition getTopicPartition()
     {
       if (topicPartition == null) {
         topicPartition = new TopicPartition(topic, partitionId);
       }
       return topicPartition;
     }

     @Override
     public boolean equals(Object o)
     {
       if (this == o) {
         return true;
       }
       if (o == null || getClass() != o.getClass()) {
         return false;
       }
       PartitionMeta that = (PartitionMeta)o;
       return Objects.equals(partitionId, that.partitionId) &&
         Objects.equals(cluster, that.cluster) &&
         Objects.equals(topic, that.topic);
     }

     @Override
     public int hashCode()
     {
       return Objects.hash(cluster, topic, partitionId);
     }

     @Override
     public String toString()
     {
       return "PartitionMeta{" +
         "cluster='" + cluster + '\'' +
         ", topicPartition=" + getTopicPartition() +
         '}';
     }
   }
 }
	/**
	* 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
	*
	* 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.apache.apex.malhar.kafka;

	import java.io.IOException;
	import java.io.StringWriter;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.LinkedList;
	import java.util.List;
	import java.util.Map;
	import java.util.Objects;
	import java.util.Properties;
	import java.util.Set;

	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import org.apache.apex.malhar.lib.util.KryoCloneUtils;
	import org.apache.hadoop.classification.InterfaceStability;
	import org.apache.kafka.clients.consumer.KafkaConsumer;
	import org.apache.kafka.common.PartitionInfo;
	import org.apache.kafka.common.TopicPartition;
	import org.apache.kafka.common.serialization.ByteArrayDeserializer;

	import com.google.common.base.Joiner;

	import com.datatorrent.api.DefaultPartition;
	import com.datatorrent.api.Partitioner;
	import com.datatorrent.api.StatsListener;

	/**
	* Abstract partitioner used to manage the partitions of kafka input operator.
	* It use a number of kafka consumers(one for each cluster) to get the latest partition metadata for topics that
	* the consumer subscribes and expose those to subclass which implements the assign method
	*
	* The partitioner is always stateless.
	*
	* @since 3.3.0
	*/
	@InterfaceStability.Evolving
	public abstract class AbstractKafkaPartitioner implements Partitioner<AbstractKafkaInputOperator>, StatsListener
	{

	private static final Logger logger = LoggerFactory.getLogger(AbstractKafkaPartitioner.class);

	private static final String META_CONSUMER_GROUP_NAME = AbstractKafkaInputOperator.class.getName() + "META_GROUP";

	protected final String[] clusters;

	protected final String[] topics;

	protected final AbstractKafkaInputOperator prototypeOperator;

	private ArrayList<KafkaConsumer<byte[], byte[]>> metadataRefreshClients;

	// prevent null
	private final List<Set<AbstractKafkaPartitioner.PartitionMeta>> currentPartitions = new LinkedList<>();

	public AbstractKafkaPartitioner(String[] clusters, String[] topics, AbstractKafkaInputOperator prototypeOperator)
	{
	this.clusters = clusters;
	this.topics = topics;
	this.prototypeOperator = prototypeOperator;
	}

	abstract List<Set<PartitionMeta>> assign(Map<String, Map<String, List<PartitionInfo>>> metadata);

	@Override
	public Collection<Partition<AbstractKafkaInputOperator>> definePartitions(
	Collection<Partition<AbstractKafkaInputOperator>> collection, PartitioningContext partitioningContext)
	{
	initMetadataClients();

	Map<String, Map<String, List<PartitionInfo>>> metadata = new HashMap<>();

	try {
	for (int i = 0; i < clusters.length; i++) {
	metadata.put(clusters[i], new HashMap<String, List<PartitionInfo>>());
	for (String topic : topics) {
	//try several time if partitionsFor(topic) returns null or throws exception.
	//partitionsFor(topic) will return null if the topic is invalid or hasn't completed
	int tryTime = 10;
	while (tryTime-- > 0) {
	try {
	List<PartitionInfo> ptis = metadataRefreshClients.get(i).partitionsFor(topic);
	if (ptis != null) {
	if (logger.isDebugEnabled()) {
	logger.debug("Partition metadata for topic {} : {}", topic, Joiner.on(';').join(ptis));
	}
	metadata.get(clusters[i]).put(topic, ptis);
	break;
	}

	logger.warn("Partition metadata for topic {} is null. retrying...", topic);

	} catch (Exception e) {
	logger.warn("Got Exception when trying get partition info for topic {}.", topic, e);
	}

	try {
	Thread.sleep(100);
	} catch (Exception e1) {
	//ignore
	}
	} //end while

	if (tryTime == 0) {
	throw new RuntimeException(
	"Get partition info for topic completely failed. Please check the log file. topic name: " + topic);
	}
	}
	}
	} finally {
	closeClients();
	}

	List<Set<AbstractKafkaPartitioner.PartitionMeta>> parts = null;
	try {
	parts = assign(metadata);
	} catch (Exception e) {
	logger.error("assign() exception.", e);
	e.printStackTrace();
	}

	if (currentPartitions == parts \|\| currentPartitions.equals(parts)) {
	logger.debug("No partition change found");
	return collection;
	} else {
	logger.info("Partition change detected: ");
	currentPartitions.clear();
	currentPartitions.addAll(parts);
	int i = 0;
	List<Partition<AbstractKafkaInputOperator>> result = new LinkedList<>();
	for (Iterator<Partition<AbstractKafkaInputOperator>> iter = collection.iterator(); iter.hasNext(); ) {
	Partition<AbstractKafkaInputOperator> nextPartition = iter.next();
	if (parts.remove(nextPartition.getPartitionedInstance().assignment())) {
	if (logger.isInfoEnabled()) {
	logger.info("[Existing] Partition {} with assignment {} ", i,
	Joiner.on(';').join(nextPartition.getPartitionedInstance().assignment()));
	}
	result.add(nextPartition);
	i++;
	}
	}

	for (Set<AbstractKafkaPartitioner.PartitionMeta> partitionAssignment : parts) {
	if (logger.isInfoEnabled()) {
	logger.info("[New] Partition {} with assignment {} ", i,
	Joiner.on(';').join(partitionAssignment));
	}
	result.add(createPartition(partitionAssignment));
	i++;
	}

	return result;
	}
	}

	protected void closeClients()
	{
	for (KafkaConsumer<byte[], byte[]> consume : metadataRefreshClients) {
	consume.close();
	}
	metadataRefreshClients = null;
	}

	@Override
	public void partitioned(Map<Integer, Partition<AbstractKafkaInputOperator>> map)
	{

	}

	@Override
	public Response processStats(BatchedOperatorStats batchedOperatorStats)
	{
	Response response = new Response();
	response.repartitionRequired = true;
	return response;
	}

	protected Partitioner.Partition<AbstractKafkaInputOperator> createPartition(
	Set<AbstractKafkaPartitioner.PartitionMeta> partitionAssignment)
	{
	Partitioner.Partition<AbstractKafkaInputOperator> p =
	new DefaultPartition<AbstractKafkaInputOperator>(KryoCloneUtils.cloneObject(prototypeOperator));
	p.getPartitionedInstance().assign(partitionAssignment);
	return p;
	}

	/**
	*
	*/
	private void initMetadataClients()
	{
	if (metadataRefreshClients != null && metadataRefreshClients.size() == clusters.length) {
	// The metadata client is active
	return;
	}

	if (clusters == null \|\| clusters.length == 0) {
	throw new IllegalStateException("clusters can not be null");
	}

	metadataRefreshClients = new ArrayList<>(clusters.length);
	int index = 0;
	for (String c : clusters) {
	Properties prop = prototypeOperator.getConsumerProps();
	prop.put("group.id", META_CONSUMER_GROUP_NAME);
	prop.put("bootstrap.servers", c);
	prop.put("key.deserializer", ByteArrayDeserializer.class.getName());
	prop.put("value.deserializer", ByteArrayDeserializer.class.getName());
	prop.put("enable.auto.commit", "false");
	if (logger.isInfoEnabled()) {
	logger.info("Consumer Properties : {} ", getPropertyAsString(prop));
	}
	metadataRefreshClients.add(index++, new KafkaConsumer<byte[], byte[]>(prop));
	}
	}

	/**
	* Converts the property list (key and element pairs) to String format
	* This format is used to print to a Stream for debugging.
	* @param prop
	* @return String
	*/
	private String getPropertyAsString(Properties prop)
	{
	StringWriter writer = new StringWriter();
	try {
	prop.store(writer, "");
	} catch (IOException e) {
	logger.error("Cannot retrieve consumer properties for Logging : {}", e.getMessage());
	}
	return writer.getBuffer().toString();
	}

	/**
	* The key object used in the assignment map for each operator
	*/
	public static class PartitionMeta
	{

	public PartitionMeta()
	{
	}

	public PartitionMeta(String cluster, String topic, int partitionId)
	{
	this.cluster = cluster;
	this.topic = topic;
	this.partitionId = partitionId;
	this.topicPartition = new TopicPartition(topic, partitionId);
	}

	private String cluster;

	private transient TopicPartition topicPartition;

	private String topic;

	private int partitionId;

	public String getCluster()
	{
	return cluster;
	}

	public int getPartitionId()
	{
	return partitionId;
	}

	public String getTopic()
	{
	return topic;
	}

	public TopicPartition getTopicPartition()
	{
	if (topicPartition == null) {
	topicPartition = new TopicPartition(topic, partitionId);
	}
	return topicPartition;
	}

	@Override
	public boolean equals(Object o)
	{
	if (this == o) {
	return true;
	}
	if (o == null \|\| getClass() != o.getClass()) {
	return false;
	}
	PartitionMeta that = (PartitionMeta)o;
	return Objects.equals(partitionId, that.partitionId) &&
	Objects.equals(cluster, that.cluster) &&
	Objects.equals(topic, that.topic);
	}

	@Override
	public int hashCode()
	{
	return Objects.hash(cluster, topic, partitionId);
	}

	@Override
	public String toString()
	{
	return "PartitionMeta{" +
	"cluster='" + cluster + '\'' +
	", topicPartition=" + getTopicPartition() +
	'}';
	}
	}
	}