Java Code Examples for kafka.utils.ZkUtils#getPartitionAssignmentForTopics()

@SuppressWarnings({ "rawtypes", "unchecked" })
private static void verifyTopic(ZkUtils zkUtils, String topic) {
  Set topics = new HashSet();
  topics.add(topic);

  // check # partition and the replication factor
  scala.collection.mutable.Map partitionAssignmentForTopics = zkUtils
      .getPartitionAssignmentForTopics(JavaConversions.asScalaSet(topics).toSeq());
  scala.collection.Map partitionAssignment = (scala.collection.Map) partitionAssignmentForTopics.get(topic).get();

  if (partitionAssignment.size() != 1) {
    throw new RuntimeException(String.format("The schema topic %s should have only 1 partition.", topic));
  }

  // check the retention policy
  Properties prop = AdminUtils.fetchEntityConfig(zkUtils, ConfigType.Topic(), topic);
  String retentionPolicy = prop.getProperty(LogConfig.CleanupPolicyProp());
  if (retentionPolicy == null || "compact".compareTo(retentionPolicy) != 0) {
    throw new RuntimeException(String.format("The retention policy of the schema topic %s must be compact.", topic));
  }
}
 

private scala.collection.Map<Object, Seq<Object>> getReplicaAssignmentForTopic(
    ZkUtils zkUtils, String topic) {
  if (topicPartitionAssignments.containsKey(topic)) {
    return topicPartitionAssignments.get(topic);
  }
  List<String> topics = new ArrayList<>();
  topics.add(topic);
  Seq<String> topicsSeq = scala.collection.JavaConverters.asScalaBuffer(topics).toSeq();

  scala.collection.mutable.Map<String, scala.collection.Map<Object, Seq<Object>>> assignments;
  assignments = zkUtils.getPartitionAssignmentForTopics(topicsSeq);

  scala.collection.Map<Object, Seq<Object>> partitionAssignment = assignments.get(topic).get();
  topicPartitionAssignments.put(topic, partitionAssignment);
  return partitionAssignment;
}
 

public static void main(String[] args) throws Exception {
  CommandLine commandLine = parseCommandLine(args);
  String zookeeper = commandLine.getOptionValue(ZOOKEEPER);

  ZkUtils zkUtils = KafkaUtils.getZkUtils(zookeeper);
  Seq<String> topicsSeq = zkUtils.getAllTopics();
  List<String> topics = scala.collection.JavaConverters.seqAsJavaList(topicsSeq);

  scala.collection.mutable.Map<String, scala.collection.Map<Object, Seq<Object>>>
      partitionAssignments = zkUtils.getPartitionAssignmentForTopics(topicsSeq);

  Map<String, Integer> replicationFactors = new HashMap<>();
  Map<String, Integer> partitionCounts = new HashMap<>();

  topics.stream().forEach(topic -> {
    int partitionCount = partitionAssignments.get(topic).get().size();
    int factor = partitionAssignments.get(topic).get().head()._2().size();
    partitionCounts.put(topic, partitionCount);
    replicationFactors.put(topic, factor);
  });

  List<PartitionInfo> urps = KafkaClusterManager.getUnderReplicatedPartitions(
      zookeeper, SecurityProtocol.PLAINTEXT, null, topics, partitionAssignments, replicationFactors, partitionCounts);

  for (PartitionInfo partitionInfo : urps) {
    LOG.info("under-replicated : {}", partitionInfo);
  }
}
 

/**
 *  KafkaClusterManager periodically check the health of the cluster. If it finds
 *  an under-replicated partitions, it will perform partition reassignment. It will also
 *  do partition reassignment for workload balancing.
 *
 *  If partitions are under-replicated in the middle of work-load balancing due to
 *  broker failure, it will send out an alert. Human intervention is needed in this case.
 */
@Override
public void run() {
  long checkIntervalInMs = clusterConfig.getCheckIntervalInSeconds() * 1000L;
  stopped = false;
  boolean foundUrps = false;
  long firstSeenUrpsTimestamp = 0L;

  while (!stopped) {
    try {
      Thread.sleep(checkIntervalInMs);
      if (maintenanceMode.get()) {
        LOG.debug("Cluster:" + clusterConfig.getClusterName() + " is in maintenace mode");
        continue;
      }
      ZkUtils zkUtils = KafkaUtils.getZkUtils(zkUrl);

      // check if there is any broker that do not have stats.
      List<Broker> noStatsBrokers = getNoStatsBrokers();
      if (!noStatsBrokers.isEmpty()) {
        Email.alertOnNoStatsBrokers(
            drkafkaConfig.getAlertEmails(), clusterConfig.getClusterName(), noStatsBrokers);
        continue;
      }

      Seq<String> topicsSeq = zkUtils.getAllTopics();
      List<String> topics = scala.collection.JavaConverters.seqAsJavaList(topicsSeq);
      scala.collection.mutable.Map<String, scala.collection.Map<Object, Seq<Object>>>
          partitionAssignments = zkUtils.getPartitionAssignmentForTopics(topicsSeq);

      Map<String, Integer> replicationFactors = new HashMap<>();
      Map<String, Integer> partitionCounts = new HashMap<>();
      topics.stream().forEach(topic -> {
        int partitionCount = partitionAssignments.get(topic).get().size();
        int factor = partitionAssignments.get(topic).get().head()._2().size();
        partitionCounts.put(topic, partitionCount);
        replicationFactors.put(topic, factor);
      });

      underReplicatedPartitions = getUnderReplicatedPartitions(zkUrl, securityProtocol, consumerConfigs,
          topics, partitionAssignments, replicationFactors, partitionCounts);
      LOG.info("Under-replicated partitions: {}", underReplicatedPartitions.size());

      for (PartitionInfo partitionInfo : underReplicatedPartitions) {
        LOG.info("under-replicated : {}", partitionInfo);
      }

      kafkaCluster.clearResourceAllocationCounters();
      if (underReplicatedPartitions.size() > 0) {
        // handle under-replicated partitions
        if (!foundUrps) {
          foundUrps = true;
          firstSeenUrpsTimestamp = System.currentTimeMillis();
        } else {
          // send out an alert if the cluster has been under-replicated for a while
          long underReplicatedTimeMills = System.currentTimeMillis() - firstSeenUrpsTimestamp;
          if (underReplicatedTimeMills > clusterConfig.getUnderReplicatedAlertTimeInMs()) {

            Email.alertOnProlongedUnderReplicatedPartitions(drkafkaConfig.getAlertEmails(),
                clusterConfig.getClusterName(),
                clusterConfig.getUnderReplicatedAlertTimeInSeconds(),
                underReplicatedPartitions);
          }
        }
        LOG.info("Under-replicated partitions in cluster {} : {}",
            clusterConfig.getClusterName(), underReplicatedPartitions.size());

        handleUnderReplicatedPartitions(underReplicatedPartitions, replicationFactors);
      } else {
        foundUrps = false;
        firstSeenUrpsTimestamp = Long.MAX_VALUE;
        if (clusterConfig.enabledWorloadBalancing()) {
          preferredLeaders.clear();
          reassignmentMap.clear();
          balanceWorkload();
        }
      }
      if (clusterConfig.enabledDeadbrokerReplacement()) {
        // replace the brokers that do not have kafkastats update for a while
        checkAndReplaceDeadBrokers();
      }
    } catch (Exception e) {
      LOG.error("Unexpected failure in cluster manager for {}: ", zkUrl, e);
    }
  }
}