streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalTopicManager.java - kafka - 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.kafka.streams.processor.internals;

 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.I0Itec.zkclient.ZkClient;
 import org.I0Itec.zkclient.exception.ZkNoNodeException;
 import org.I0Itec.zkclient.exception.ZkNodeExistsException;
 import org.I0Itec.zkclient.serialize.ZkSerializer;
 import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.zookeeper.ZooDefs;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;

 public class InternalTopicManager {

     private static final Logger log = LoggerFactory.getLogger(InternalTopicManager.class);

     // TODO: the following ZK dependency should be removed after KIP-4
     private static final String ZK_TOPIC_PATH = "/brokers/topics";
     private static final String ZK_BROKER_PATH = "/brokers/ids";
     private static final String ZK_DELETE_TOPIC_PATH = "/admin/delete_topics";
     private static final String ZK_ENTITY_CONFIG_PATH = "/config/topics";
     // TODO: the following LogConfig dependency should be removed after KIP-4
     private static final String CLEANUP_POLICY_PROP = "cleanup.policy";
     private static final String COMPACT = "compact";

     private final ZkClient zkClient;
     private final int replicationFactor;

     private class ZKStringSerializer implements ZkSerializer {

         /**
          * @throws AssertionError if the byte String encoding type is not supported
          */
         @Override
         public byte[] serialize(Object data) {
             try {
                 return ((String) data).getBytes("UTF-8");
             } catch (UnsupportedEncodingException e) {
                 throw new AssertionError(e);
             }
         }

         /**
          * @throws AssertionError if the byte String encoding type is not supported
          */
         @Override
         public Object deserialize(byte[] bytes) {
             try {
                 if (bytes == null)
                     return null;
                 else
                     return new String(bytes, "UTF-8");
             } catch (UnsupportedEncodingException e) {
                 throw new AssertionError(e);
             }
         }
     }

     public InternalTopicManager() {
         this.zkClient = null;
         this.replicationFactor = 0;
     }

     public InternalTopicManager(String zkConnect, int replicationFactor) {
         this.zkClient = new ZkClient(zkConnect, 30 * 1000, 30 * 1000, new ZKStringSerializer());
         this.replicationFactor = replicationFactor;
     }

     public void makeReady(String topic, int numPartitions, boolean compactTopic) {
         boolean topicNotReady = true;

         while (topicNotReady) {
             Map<Integer, List<Integer>> topicMetadata = getTopicMetadata(topic);

             if (topicMetadata == null) {
                 try {
                     createTopic(topic, numPartitions, replicationFactor, compactTopic);
                 } catch (ZkNodeExistsException e) {
                     // ignore and continue
                 }
             } else {
                 if (topicMetadata.size() > numPartitions) {
                     // else if topic exists with more #.partitions than needed, delete in order to re-create it
                     try {
                         deleteTopic(topic);
                     } catch (ZkNodeExistsException e) {
                         // ignore and continue
                     }
                 } else if (topicMetadata.size() < numPartitions) {
                     // else if topic exists with less #.partitions than needed, add partitions
                     try {
                         addPartitions(topic, numPartitions - topicMetadata.size(), replicationFactor, topicMetadata);
                     } catch (ZkNoNodeException e) {
                         // ignore and continue
                     }
                 } else {
                     topicNotReady = false;
                 }
             }
         }
     }

     private List<Integer> getBrokers() {
         List<Integer> brokers = new ArrayList<>();
         for (String broker: zkClient.getChildren(ZK_BROKER_PATH)) {
             brokers.add(Integer.parseInt(broker));
         }
         Collections.sort(brokers);

         log.debug("Read brokers {} from ZK in partition assignor.", brokers);

         return brokers;
     }

     @SuppressWarnings("unchecked")
     private Map<Integer, List<Integer>> getTopicMetadata(String topic) {
         String data = zkClient.readData(ZK_TOPIC_PATH + "/" + topic, true);

         if (data == null) return null;

         try {
             ObjectMapper mapper = new ObjectMapper();

             Map<String, Object> dataMap = mapper.readValue(data, new TypeReference<Map<String, Object>>() {

             });

             Map<Integer, List<Integer>> partitions = (Map<Integer, List<Integer>>) dataMap.get("partitions");

             log.debug("Read partitions {} for topic {} from ZK in partition assignor.", partitions, topic);

             return partitions;
         } catch (IOException e) {
             throw new StreamsException("Error while reading topic metadata from ZK for internal topic " + topic, e);
         }
     }

     private void createTopic(String topic, int numPartitions, int replicationFactor, boolean compactTopic) throws ZkNodeExistsException {
         log.debug("Creating topic {} with {} partitions from ZK in partition assignor.", topic, numPartitions);
         Properties prop = new Properties();
         ObjectMapper mapper = new ObjectMapper();
         List<Integer> brokers = getBrokers();
         int numBrokers = brokers.size();
         if (numBrokers < replicationFactor) {
             log.warn("Not enough brokers found. The replication factor is reduced from " + replicationFactor + " to " +  numBrokers);
             replicationFactor = numBrokers;
         }

         Map<Integer, List<Integer>> assignment = new HashMap<>();

         for (int i = 0; i < numPartitions; i++) {
             ArrayList<Integer> brokerList = new ArrayList<>();
             for (int r = 0; r < replicationFactor; r++) {
                 int shift = r * numBrokers / replicationFactor;
                 brokerList.add(brokers.get((i + shift) % numBrokers));
             }
             assignment.put(i, brokerList);
         }
         // write out config first just like in AdminUtils.scala createOrUpdateTopicPartitionAssignmentPathInZK()
         if (compactTopic) {
             prop.put(CLEANUP_POLICY_PROP, COMPACT);
             try {
                 Map<String, Object> dataMap = new HashMap<>();
                 dataMap.put("version", 1);
                 dataMap.put("config", prop);
                 String data = mapper.writeValueAsString(dataMap);
                 zkClient.createPersistent(ZK_ENTITY_CONFIG_PATH + "/" + topic, data, ZooDefs.Ids.OPEN_ACL_UNSAFE);
             } catch (JsonProcessingException e) {
                 throw new StreamsException("Error while creating topic config in ZK for internal topic " + topic, e);
             }
         }

         // try to write to ZK with open ACL
         try {
             Map<String, Object> dataMap = new HashMap<>();
             dataMap.put("version", 1);
             dataMap.put("partitions", assignment);
             String data = mapper.writeValueAsString(dataMap);

             zkClient.createPersistent(ZK_TOPIC_PATH + "/" + topic, data, ZooDefs.Ids.OPEN_ACL_UNSAFE);
         } catch (JsonProcessingException e) {
             throw new StreamsException("Error while creating topic metadata in ZK for internal topic " + topic, e);
         }
     }

     private void deleteTopic(String topic) throws ZkNodeExistsException {
         log.debug("Deleting topic {} from ZK in partition assignor.", topic);

         zkClient.createPersistent(ZK_DELETE_TOPIC_PATH + "/" + topic, "", ZooDefs.Ids.OPEN_ACL_UNSAFE);
     }

     private void addPartitions(String topic, int numPartitions, int replicationFactor, Map<Integer, List<Integer>> existingAssignment) {
         log.debug("Adding {} partitions topic {} from ZK with existing partitions assigned as {} in partition assignor.", topic, numPartitions, existingAssignment);

         List<Integer> brokers = getBrokers();
         int numBrokers = brokers.size();
         if (numBrokers < replicationFactor) {
             log.warn("Not enough brokers found. The replication factor is reduced from " + replicationFactor + " to " +  numBrokers);
             replicationFactor = numBrokers;
         }

         int startIndex = existingAssignment.size();

         Map<Integer, List<Integer>> newAssignment = new HashMap<>(existingAssignment);

         for (int i = 0; i < numPartitions; i++) {
             ArrayList<Integer> brokerList = new ArrayList<>();
             for (int r = 0; r < replicationFactor; r++) {
                 int shift = r * numBrokers / replicationFactor;
                 brokerList.add(brokers.get((i + shift) % numBrokers));
             }
             newAssignment.put(i + startIndex, brokerList);
         }

         // try to write to ZK with open ACL
         try {
             Map<String, Object> dataMap = new HashMap<>();
             dataMap.put("version", 1);
             dataMap.put("partitions", newAssignment);

             ObjectMapper mapper = new ObjectMapper();
             String data = mapper.writeValueAsString(dataMap);

             zkClient.writeData(ZK_TOPIC_PATH + "/" + topic, data);
         } catch (JsonProcessingException e) {
             throw new StreamsException("Error while updating topic metadata in ZK for internal topic " + topic, e);
         }
     }

 }
	/**
	* 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.kafka.streams.processor.internals;

	import com.fasterxml.jackson.core.JsonProcessingException;
	import com.fasterxml.jackson.core.type.TypeReference;
	import com.fasterxml.jackson.databind.ObjectMapper;
	import org.I0Itec.zkclient.ZkClient;
	import org.I0Itec.zkclient.exception.ZkNoNodeException;
	import org.I0Itec.zkclient.exception.ZkNodeExistsException;
	import org.I0Itec.zkclient.serialize.ZkSerializer;
	import org.apache.kafka.streams.errors.StreamsException;
	import org.apache.zookeeper.ZooDefs;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.IOException;
	import java.io.UnsupportedEncodingException;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Properties;

	public class InternalTopicManager {

	private static final Logger log = LoggerFactory.getLogger(InternalTopicManager.class);

	// TODO: the following ZK dependency should be removed after KIP-4
	private static final String ZK_TOPIC_PATH = "/brokers/topics";
	private static final String ZK_BROKER_PATH = "/brokers/ids";
	private static final String ZK_DELETE_TOPIC_PATH = "/admin/delete_topics";
	private static final String ZK_ENTITY_CONFIG_PATH = "/config/topics";
	// TODO: the following LogConfig dependency should be removed after KIP-4
	private static final String CLEANUP_POLICY_PROP = "cleanup.policy";
	private static final String COMPACT = "compact";

	private final ZkClient zkClient;
	private final int replicationFactor;

	private class ZKStringSerializer implements ZkSerializer {

	/**
	* @throws AssertionError if the byte String encoding type is not supported
	*/
	@Override
	public byte[] serialize(Object data) {
	try {
	return ((String) data).getBytes("UTF-8");
	} catch (UnsupportedEncodingException e) {
	throw new AssertionError(e);
	}
	}

	/**
	* @throws AssertionError if the byte String encoding type is not supported
	*/
	@Override
	public Object deserialize(byte[] bytes) {
	try {
	if (bytes == null)
	return null;
	else
	return new String(bytes, "UTF-8");
	} catch (UnsupportedEncodingException e) {
	throw new AssertionError(e);
	}
	}
	}

	public InternalTopicManager() {
	this.zkClient = null;
	this.replicationFactor = 0;
	}

	public InternalTopicManager(String zkConnect, int replicationFactor) {
	this.zkClient = new ZkClient(zkConnect, 30 * 1000, 30 * 1000, new ZKStringSerializer());
	this.replicationFactor = replicationFactor;
	}

	public void makeReady(String topic, int numPartitions, boolean compactTopic) {
	boolean topicNotReady = true;

	while (topicNotReady) {
	Map<Integer, List<Integer>> topicMetadata = getTopicMetadata(topic);

	if (topicMetadata == null) {
	try {
	createTopic(topic, numPartitions, replicationFactor, compactTopic);
	} catch (ZkNodeExistsException e) {
	// ignore and continue
	}
	} else {
	if (topicMetadata.size() > numPartitions) {
	// else if topic exists with more #.partitions than needed, delete in order to re-create it
	try {
	deleteTopic(topic);
	} catch (ZkNodeExistsException e) {
	// ignore and continue
	}
	} else if (topicMetadata.size() < numPartitions) {
	// else if topic exists with less #.partitions than needed, add partitions
	try {
	addPartitions(topic, numPartitions - topicMetadata.size(), replicationFactor, topicMetadata);
	} catch (ZkNoNodeException e) {
	// ignore and continue
	}
	} else {
	topicNotReady = false;
	}
	}
	}
	}

	private List<Integer> getBrokers() {
	List<Integer> brokers = new ArrayList<>();
	for (String broker: zkClient.getChildren(ZK_BROKER_PATH)) {
	brokers.add(Integer.parseInt(broker));
	}
	Collections.sort(brokers);

	log.debug("Read brokers {} from ZK in partition assignor.", brokers);

	return brokers;
	}

	@SuppressWarnings("unchecked")
	private Map<Integer, List<Integer>> getTopicMetadata(String topic) {
	String data = zkClient.readData(ZK_TOPIC_PATH + "/" + topic, true);

	if (data == null) return null;

	try {
	ObjectMapper mapper = new ObjectMapper();

	Map<String, Object> dataMap = mapper.readValue(data, new TypeReference<Map<String, Object>>() {

	});

	Map<Integer, List<Integer>> partitions = (Map<Integer, List<Integer>>) dataMap.get("partitions");

	log.debug("Read partitions {} for topic {} from ZK in partition assignor.", partitions, topic);

	return partitions;
	} catch (IOException e) {
	throw new StreamsException("Error while reading topic metadata from ZK for internal topic " + topic, e);
	}
	}

	private void createTopic(String topic, int numPartitions, int replicationFactor, boolean compactTopic) throws ZkNodeExistsException {
	log.debug("Creating topic {} with {} partitions from ZK in partition assignor.", topic, numPartitions);
	Properties prop = new Properties();
	ObjectMapper mapper = new ObjectMapper();
	List<Integer> brokers = getBrokers();
	int numBrokers = brokers.size();
	if (numBrokers < replicationFactor) {
	log.warn("Not enough brokers found. The replication factor is reduced from " + replicationFactor + " to " + numBrokers);
	replicationFactor = numBrokers;
	}

	Map<Integer, List<Integer>> assignment = new HashMap<>();

	for (int i = 0; i < numPartitions; i++) {
	ArrayList<Integer> brokerList = new ArrayList<>();
	for (int r = 0; r < replicationFactor; r++) {
	int shift = r * numBrokers / replicationFactor;
	brokerList.add(brokers.get((i + shift) % numBrokers));
	}
	assignment.put(i, brokerList);
	}
	// write out config first just like in AdminUtils.scala createOrUpdateTopicPartitionAssignmentPathInZK()
	if (compactTopic) {
	prop.put(CLEANUP_POLICY_PROP, COMPACT);
	try {
	Map<String, Object> dataMap = new HashMap<>();
	dataMap.put("version", 1);
	dataMap.put("config", prop);
	String data = mapper.writeValueAsString(dataMap);
	zkClient.createPersistent(ZK_ENTITY_CONFIG_PATH + "/" + topic, data, ZooDefs.Ids.OPEN_ACL_UNSAFE);
	} catch (JsonProcessingException e) {
	throw new StreamsException("Error while creating topic config in ZK for internal topic " + topic, e);
	}
	}

	// try to write to ZK with open ACL
	try {
	Map<String, Object> dataMap = new HashMap<>();
	dataMap.put("version", 1);
	dataMap.put("partitions", assignment);
	String data = mapper.writeValueAsString(dataMap);

	zkClient.createPersistent(ZK_TOPIC_PATH + "/" + topic, data, ZooDefs.Ids.OPEN_ACL_UNSAFE);
	} catch (JsonProcessingException e) {
	throw new StreamsException("Error while creating topic metadata in ZK for internal topic " + topic, e);
	}
	}

	private void deleteTopic(String topic) throws ZkNodeExistsException {
	log.debug("Deleting topic {} from ZK in partition assignor.", topic);

	zkClient.createPersistent(ZK_DELETE_TOPIC_PATH + "/" + topic, "", ZooDefs.Ids.OPEN_ACL_UNSAFE);
	}

	private void addPartitions(String topic, int numPartitions, int replicationFactor, Map<Integer, List<Integer>> existingAssignment) {
	log.debug("Adding {} partitions topic {} from ZK with existing partitions assigned as {} in partition assignor.", topic, numPartitions, existingAssignment);

	List<Integer> brokers = getBrokers();
	int numBrokers = brokers.size();
	if (numBrokers < replicationFactor) {
	log.warn("Not enough brokers found. The replication factor is reduced from " + replicationFactor + " to " + numBrokers);
	replicationFactor = numBrokers;
	}

	int startIndex = existingAssignment.size();

	Map<Integer, List<Integer>> newAssignment = new HashMap<>(existingAssignment);

	for (int i = 0; i < numPartitions; i++) {
	ArrayList<Integer> brokerList = new ArrayList<>();
	for (int r = 0; r < replicationFactor; r++) {
	int shift = r * numBrokers / replicationFactor;
	brokerList.add(brokers.get((i + shift) % numBrokers));
	}
	newAssignment.put(i + startIndex, brokerList);
	}

	// try to write to ZK with open ACL
	try {
	Map<String, Object> dataMap = new HashMap<>();
	dataMap.put("version", 1);
	dataMap.put("partitions", newAssignment);

	ObjectMapper mapper = new ObjectMapper();
	String data = mapper.writeValueAsString(dataMap);

	zkClient.writeData(ZK_TOPIC_PATH + "/" + topic, data);
	} catch (JsonProcessingException e) {
	throw new StreamsException("Error while updating topic metadata in ZK for internal topic " + topic, e);
	}
	}

	}