/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio;

import alluxio.util.CommonUtils;
import alluxio.util.io.PathUtils;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import javax.annotation.concurrent.ThreadSafe;

/**
 * Utility to get leader from zookeeper.
 */
@ThreadSafe
public final class MasterInquireClient {
  private static final Logger LOG = LoggerFactory.getLogger(Constants.LOGGER_TYPE);

  /** Map from key spliced by the address for Zookeeper and path of leader to created client. */
  private static HashMap<String, MasterInquireClient> sCreatedClients = new HashMap<>();

  /**
   * Gets the client.
   *
   * @param zookeeperAddress the address for Zookeeper
   * @param electionPath the path of the master election
   * @param leaderPath the path of the leader
   *
   * @return the client
   */
  public static synchronized MasterInquireClient getClient(String zookeeperAddress,
      String electionPath, String leaderPath) {
    String key = zookeeperAddress + leaderPath;
    if (!sCreatedClients.containsKey(key)) {
      sCreatedClients.put(key, new MasterInquireClient(zookeeperAddress, electionPath, leaderPath));
    }
    return sCreatedClients.get(key);
  }

  private final String mZookeeperAddress;
  private final String mElectionPath;
  private final String mLeaderPath;
  private final CuratorFramework mClient;
  private final int mMaxTry;

  private MasterInquireClient(String zookeeperAddress, String electionPath, String leaderPath) {
    mZookeeperAddress = zookeeperAddress;
    mElectionPath = electionPath;
    mLeaderPath = leaderPath;

    LOG.info("Creating new zookeeper client. address: {}", mZookeeperAddress);
    mClient =
        CuratorFrameworkFactory.newClient(mZookeeperAddress, new ExponentialBackoffRetry(
            Constants.SECOND_MS, 3));
    mClient.start();

    mMaxTry = Configuration.getInt(PropertyKey.ZOOKEEPER_LEADER_INQUIRY_RETRY_COUNT);
  }

  /**
   * @return the address of the current leader master
   */
  public synchronized String getLeaderAddress() {
    int tried = 0;
    try {
      while (tried < mMaxTry) {
        if (mClient.checkExists().forPath(mLeaderPath) != null) {
          List<String> masters = mClient.getChildren().forPath(mLeaderPath);
          LOG.info("Master addresses: {}", masters);
          if (masters.size() >= 1) {
            if (masters.size() == 1) {
              return masters.get(0);
            }

            long maxTime = 0;
            String leader = "";
            for (String master : masters) {
              Stat stat = mClient.checkExists().forPath(
                  PathUtils.concatPath(mLeaderPath, master));
              if (stat != null && stat.getCtime() > maxTime) {
                maxTime = stat.getCtime();
                leader = master;
              }
            }
            LOG.info("The leader master: {}", leader);
            return leader;
          }
        } else {
          LOG.info("{} does not exist ({})", mLeaderPath, ++tried);
        }
        CommonUtils.sleepMs(LOG, Constants.SECOND_MS);
      }
    } catch (Exception e) {
      LOG.error("Error getting the leader master address from zookeeper. Zookeeper address: {}",
          mZookeeperAddress, e);
    }

    return null;
  }

  /**
   * @return the list of all the master addresses
   */
  public synchronized List<String> getMasterAddresses() {
    int tried = 0;
    try {
      while (tried < mMaxTry) {
        if (mClient.checkExists().forPath(mElectionPath) != null) {
          List<String> children = mClient.getChildren().forPath(mElectionPath);
          List<String> ret = new ArrayList<>();
          for (String child : children) {
            byte[] data = mClient.getData().forPath(PathUtils.concatPath(mElectionPath, child));
            if (data != null) {
              // The data of the child znode is the corresponding master address for now
              ret.add(new String(data, "utf-8"));
            }
          }
          LOG.info("All masters: {}", ret);
          return ret;
        } else {
          LOG.info("{} does not exist ({})", mElectionPath, ++tried);
        }
      }
    } catch (Exception e) {
      LOG.error("Error getting the master addresses from zookeeper. Zookeeper address: {}",
          mZookeeperAddress, e);
    }

    return null;
  }
}

@@ -761,24 +761,63 @@
 ckoffRetry;%0A
+import org.apache.zookeeper.ZooKeeper;%0A
 import org.a
@@ -2851,32 +2851,154 @@
     
-int tried = 0;%0A    try %7B
+long startTime = System.currentTimeMillis();%0A    int tried = 0;%0A    try %7B%0A      ZooKeeper zookeeper = mClient.getZookeeperClient().getZooKeeper();
 %0A   
@@ -3042,37 +3042,24 @@
 if (
-mClient.checkExists().forPath
+zookeeper.exists
 (mLe
@@ -3062,24 +3062,31 @@
 (mLeaderPath
+, false
 ) != null) %7B
@@ -3119,23 +3119,25 @@
 sters = 
-mClient
+zookeeper
 .getChil
@@ -3141,26 +3141,16 @@
 hildren(
-).forPath(
 mLeaderP
@@ -3148,24 +3148,30 @@
 (mLeaderPath
+, null
 );%0A         
@@ -3478,57 +3478,25 @@
 t = 
-mClient.checkExists().forPath(%0A                  
+zookeeper.exists(
 Path
@@ -3532,16 +3532,22 @@
  master)
+, null
 );%0A     
@@ -4117,32 +4117,144 @@
 dress, e);%0A    %7D
+ finally %7B%0A      LOG.debug(%22Finished getLeaderAddress() in %7B%7Dms%22, System.currentTimeMillis() - startTime);%0A    %7D
 %0A%0A    return nul

Use faster low-level Zookeeper API for inquire client