content: /*
   * 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.master;
  
  import alluxio.Constants;
  import alluxio.exception.status.UnavailableException;
  import alluxio.uri.Authority;
  import alluxio.uri.ZookeeperAuthority;
  import alluxio.util.CommonUtils;
  import alluxio.util.io.PathUtils;
  import alluxio.util.network.NetworkAddressUtils;
  
  import org.apache.curator.CuratorZookeeperClient;
  import org.apache.curator.framework.CuratorFramework;
  import org.apache.curator.framework.CuratorFrameworkFactory;
  import org.apache.curator.retry.ExponentialBackoffRetry;
  import org.apache.curator.utils.ZookeeperFactory;
  import org.apache.zookeeper.Watcher;
  import org.apache.zookeeper.ZooKeeper;
  import org.apache.zookeeper.client.ZKClientConfig;
  import org.apache.zookeeper.data.Stat;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.Closeable;
  import java.net.InetSocketAddress;
  import java.nio.charset.StandardCharsets;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Objects;
  import javax.annotation.concurrent.ThreadSafe;
  
  /**
   * Utility to get leader from zookeeper.
   */
  @ThreadSafe
  public final class ZkMasterInquireClient implements MasterInquireClient, Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(ZkMasterInquireClient.class);
  
    /** Map from key spliced by the address for Zookeeper and path of leader to created client. */
    private static final HashMap<ZkMasterConnectDetails, ZkMasterInquireClient> CREATED_CLIENTS =
        new HashMap<>();
  
    private final ZkMasterConnectDetails mConnectDetails;
    private final String mElectionPath;
    private final CuratorFramework mClient;
    private final int mInquireRetryCount;
  
    /**
     * Zookeeper factory for curator that controls enabling/disabling client authentication.
     */
    private static class AlluxioZookeeperFactory implements ZookeeperFactory {
      private final boolean mAuthEnabled;
  
      public AlluxioZookeeperFactory(boolean authEnabled) {
        mAuthEnabled = authEnabled;
      }
  
      @Override
      public ZooKeeper newZooKeeper(String connectString, int sessionTimeout, Watcher watcher,
                                    boolean canBeReadOnly) throws Exception {
        ZKClientConfig zkConfig = new ZKClientConfig();
        zkConfig.setProperty(ZKClientConfig.ENABLE_CLIENT_SASL_KEY,
            Boolean.toString(mAuthEnabled).toLowerCase());
        return new ZooKeeper(connectString, sessionTimeout, watcher, zkConfig);
      }
    }
  
    /**
     * Gets the client.
     *
     * @param zookeeperAddress the address for Zookeeper
     * @param electionPath the path of the master election
     * @param leaderPath the path of the leader
     * @param inquireRetryCount the number of times to retry connections
     * @param authEnabled if Alluxio client-side auth is enabled
     * @return the client
     */
    public static synchronized ZkMasterInquireClient getClient(String zookeeperAddress,
        String electionPath, String leaderPath, int inquireRetryCount, boolean authEnabled) {
      ZkMasterConnectDetails connectDetails =
          new ZkMasterConnectDetails(zookeeperAddress, leaderPath);
      if (!CREATED_CLIENTS.containsKey(connectDetails)) {
        CREATED_CLIENTS.put(connectDetails,
            new ZkMasterInquireClient(connectDetails, electionPath, inquireRetryCount, authEnabled));
      }
      return CREATED_CLIENTS.get(connectDetails);
    }
  
    /**
     * Constructor for {@link ZkMasterInquireClient}.
     *
     * @param connectDetails connect details
     * @param electionPath the path of the master election
     * @param inquireRetryCount the number of times to retry connections
     * @param authEnabled if Alluxio client-side auth is enabled
     */
    private ZkMasterInquireClient(ZkMasterConnectDetails connectDetails, String electionPath,
        int inquireRetryCount, boolean authEnabled) {
      mConnectDetails = connectDetails;
      mElectionPath = electionPath;
  
      LOG.info("Creating new zookeeper client for {}", connectDetails);
      CuratorFrameworkFactory.Builder curatorBuilder = CuratorFrameworkFactory.builder();
      curatorBuilder.connectString(connectDetails.getZkAddress());
      curatorBuilder.retryPolicy(new ExponentialBackoffRetry(Constants.SECOND_MS, 3));
      curatorBuilder.zookeeperFactory(new AlluxioZookeeperFactory(authEnabled));
      // Force compatibility mode to support writing to 3.4.x servers.
      curatorBuilder.zk34CompatibilityMode(true);
      // Prevent using container parents as it breaks compatibility with 3.4.x servers.
      // This is only required if the client is used to write data to zookeeper.
      curatorBuilder.dontUseContainerParents();
  
      mClient = curatorBuilder.build();
  
      mInquireRetryCount = inquireRetryCount;
    }
  
    @Override
    public synchronized InetSocketAddress getPrimaryRpcAddress() throws UnavailableException {
      ensureStarted();
      long startTime = System.currentTimeMillis();
      int tried = 0;
      try {
        CuratorZookeeperClient curatorClient = mClient.getZookeeperClient();
        // #blockUntilConnectedOrTimedOut() will block for at least 1 second, even if the client is
        // connected within a few milliseconds. We improve the latency here by first waiting on the
        // connection status explicitly.
        for (int i = 0; i < 50; i++) {
          if (curatorClient.isConnected()) {
            break;
          }
          CommonUtils.sleepMs(20);
        }
        curatorClient.blockUntilConnectedOrTimedOut();
        String leaderPath = mConnectDetails.getLeaderPath();
        while (tried++ < mInquireRetryCount) {
          ZooKeeper zookeeper = curatorClient.getZooKeeper();
          if (zookeeper.exists(leaderPath, false) != null) {
            List<String> masters = zookeeper.getChildren(leaderPath, null);
            LOG.debug("Master addresses: {}", masters);
            if (masters.size() >= 1) {
              if (masters.size() == 1) {
                return NetworkAddressUtils.parseInetSocketAddress(masters.get(0));
              }
  
              long maxTime = 0;
              String leader = "";
              for (String master : masters) {
                Stat stat = zookeeper.exists(PathUtils.concatPath(leaderPath, master), null);
                if (stat != null && stat.getCtime() > maxTime) {
                  maxTime = stat.getCtime();
                  leader = master;
                }
              }
              LOG.debug("The leader master: {}", leader);
              return NetworkAddressUtils.parseInetSocketAddress(leader);
            }
          } else {
            LOG.info("{} does not exist ({})", leaderPath, tried);
          }
          CommonUtils.sleepMs(LOG, Constants.SECOND_MS);
        }
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      } catch (Exception e) {
        LOG.error("Error getting the leader master address from zookeeper. Zookeeper: {}",
            mConnectDetails, e);
      } finally {
        LOG.debug("Finished getPrimaryRpcAddress() in {}ms", System.currentTimeMillis() - startTime);
      }
  
      throw new UnavailableException("Failed to determine primary master rpc address");
    }
  
    @Override
    public synchronized List<InetSocketAddress> getMasterRpcAddresses() throws UnavailableException {
      ensureStarted();
      int tried = 0;
      try {
        while (tried < mInquireRetryCount) {
          if (mClient.checkExists().forPath(mElectionPath) != null) {
            List<String> children = mClient.getChildren().forPath(mElectionPath);
            List<InetSocketAddress> 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(NetworkAddressUtils.parseInetSocketAddress(new String(data,
                    StandardCharsets.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: {}", mConnectDetails,
            e);
      }
  
      throw new UnavailableException("Failed to query zookeeper for master RPC addresses");
    }
  
    private synchronized void ensureStarted() {
      switch (mClient.getState()) {
        case LATENT:
          mClient.start();
          return;
        case STARTED:
          return;
        case STOPPED:
          throw new IllegalStateException("Client has already been closed");
        default:
          throw new IllegalStateException("Unknown state: " + mClient.getState());
      }
    }
  
    @Override
    public void close() {
      mClient.close();
    }
  
    @Override
    public ConnectDetails getConnectDetails() {
      return mConnectDetails;
    }
  
    /**
     * Details used to connect to the leader Alluxio master via Zookeeper.
     */
    public static class ZkMasterConnectDetails implements ConnectDetails {
      private final String mZkAddress;
      private final String mLeaderPath;
  
      /**
       * @param zkAddress a zookeeper address
       * @param leaderPath a leader path
       */
      public ZkMasterConnectDetails(String zkAddress, String leaderPath) {
        mZkAddress = zkAddress;
        mLeaderPath = leaderPath;
      }
  
      /**
       * @return the zookeeper address
       */
      public String getZkAddress() {
        return mZkAddress;
      }
  
      /**
       * @return the leader path
       */
      public String getLeaderPath() {
        return mLeaderPath;
      }
  
      @Override
      public Authority toAuthority() {
        return new ZookeeperAuthority(mZkAddress);
      }
  
      @Override
      public boolean equals(Object o) {
        if (this == o) {
          return true;
        }
        if (!(o instanceof ZkMasterConnectDetails)) {
          return false;
        }
        ZkMasterConnectDetails that = (ZkMasterConnectDetails) o;
        return mZkAddress.equals(that.mZkAddress)
            && mLeaderPath.equals(that.mLeaderPath);
      }
  
      @Override
      public int hashCode() {
        return Objects.hash(mZkAddress, mLeaderPath);
      }
  
      @Override
      public String toString() {
        return toAuthority() + mLeaderPath;
      }
    }
  }
  