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.HealthCheckClient;
  import alluxio.common.RpcPortHealthCheckClient;
  import alluxio.conf.AlluxioConfiguration;
  import alluxio.grpc.ServiceType;
  import alluxio.retry.RetryPolicy;
  import alluxio.security.user.UserState;
  import alluxio.util.CommonUtils;
  import alluxio.util.ConfigurationUtils;
  import alluxio.util.ShellUtils;
  import alluxio.util.network.NetworkAddressUtils;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.net.InetSocketAddress;
  import java.util.List;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.function.Supplier;
  
  /**
   * MasterHealthCheckClient check whether Alluxio master is serving RPC requests and
   * if the AlluxioMaster process is running in all master hosts.
   */
  public class MasterHealthCheckClient implements HealthCheckClient {
    private static final Logger LOG = LoggerFactory.getLogger(MasterHealthCheckClient.class);
    private MasterType mAlluxioMasterType;
    private boolean mProcessCheck;
    private ExecutorService mExecutorService;
    private Supplier<RetryPolicy> mPolicySupplier;
    private AlluxioConfiguration mConf;
  
    /**
     * An enum mapping master types to fully qualified class names.
     */
    public enum MasterType {
      MASTER("alluxio.master.AlluxioMaster"),
      JOB_MASTER("alluxio.master.AlluxioJobMaster")
      ;
  
      private String mClassName;
  
      /**
       * Creates a new instance of {@link MasterType}.
       *
       * @param className the fully qualified class name of the master
       */
      MasterType(String className) {
        mClassName = className;
      }
  
      /**
       * @return the fully qualified classname associated with the master type
       */
      public String getClassName() {
        return mClassName;
      }
    }
  
    /**
     * Builder for a {@link MasterHealthCheckClient}.
     */
    public static class Builder {
  
      private boolean mProcessCheck;
      private MasterType mAlluxioMasterType;
      private AlluxioConfiguration mConf;
      private Supplier<RetryPolicy> mPolicySupplier;
  
      /**
       * Constructs the builder with default values.
       *
       * @param alluxioConf Alluxio configuration
       */
      public Builder(AlluxioConfiguration alluxioConf) {
        mProcessCheck = true;
        mAlluxioMasterType = MasterType.MASTER;
        mConf = alluxioConf;
        mPolicySupplier = AlluxioMasterMonitor.TWO_MIN_EXP_BACKOFF;
      }
  
      /**
       * @param processCheck whether to check the AlluxioMaster process is alive
       * @return an instance of the builder
       */
      public Builder withProcessCheck(boolean processCheck) {
        mProcessCheck = processCheck;
        return this;
      }
  
      /**
       *
       * @param alluxioConf Alluxio configuration
       * @return a builder which utlizes the given alluxio configuration
       */
      public Builder withConfiguration(AlluxioConfiguration alluxioConf) {
        mConf = alluxioConf;
        return this;
      }
  
      /**
       * @param masterType the Alluxio master type
       * @return an instance of the builder
       */
      public Builder withAlluxioMasterType(MasterType masterType) {
        mAlluxioMasterType = masterType;
        return this;
      }
  
      /**
       * @param policySupplier the retry policy supplier to use for the health check
       * @return a builder utilizing the given retry policy supplier
       */
      public Builder withRetryPolicy(Supplier<RetryPolicy> policySupplier) {
        mPolicySupplier = policySupplier;
        return this;
      }
  
      /**
       * @return a {@link MasterHealthCheckClient} for the current builder values
       */
      public HealthCheckClient build() {
        return new MasterHealthCheckClient(mAlluxioMasterType, mProcessCheck, mPolicySupplier, mConf);
      }
    }
  
    /**
     * Runnable for checking if the AlluxioMaster is serving RPCs.
     */
    public final class MasterServingHealthCheck extends RpcPortHealthCheckClient implements Runnable {
      private AtomicBoolean mIsServing = new AtomicBoolean(false);
  
      /**
       * Creates a new instance of {@link MasterServingHealthCheck} to check for an open RPC port
       * for the given service.
       *
       * @param nodeAddress the address of the node to connect to
       * @param serviceType the gRPC service type
       * @param retryPolicySupplier a retry policy supplier
       * @param alluxioConf Alluxio configuration
       */
      public MasterServingHealthCheck(InetSocketAddress nodeAddress, ServiceType serviceType,
          Supplier<RetryPolicy> retryPolicySupplier, AlluxioConfiguration alluxioConf) {
        super(nodeAddress, serviceType, retryPolicySupplier, alluxioConf);
      }
  
      @Override
      public void run() {
        if (super.isServing()) {
          mIsServing.set(true);
        }
      }
  
      /**
       * @return whether or we were able to ping the service
       */
      public boolean serving() {
        return mIsServing.get();
      }
    }
  
    /**
     * Runnable for checking if the AlluxioMaster process are running in all the masters hosts.
     * This includes both primary and standby masters.
     */
    public final class ProcessCheckRunnable implements Runnable {
      private String mAlluxioMasterName;
  
      /**
       * Creates a new instance of ProcessCheckRunnable.
       *
       * @param alluxioMasterName the Alluxio master process name
       */
      public ProcessCheckRunnable(String alluxioMasterName) {
        mAlluxioMasterName = alluxioMasterName;
      }
  
      @Override
      public void run() {
        UserState userState = UserState.Factory.create(mConf);
        MasterInquireClient client = MasterInquireClient.Factory.create(mConf, userState);
        try {
          while (true) {
            List<InetSocketAddress> addresses = client.getMasterRpcAddresses();
            for (InetSocketAddress address : addresses) {
              String host = address.getHostName();
              int port = address.getPort();
              LOG.debug("Master health check on node {}", host);
              String cmd = String.format("ssh %s %s %s", ShellUtils.COMMON_SSH_OPTS, host,
                  "ps -ef | grep \"" + mAlluxioMasterName + "$\" | "
                  + "grep \"java\" | "
                  + "awk '{ print $2; }'");
              LOG.debug("Executing: {}", cmd);
              String output = ShellUtils.execCommand("bash", "-c", cmd);
              if (output.isEmpty()) {
                throw new IllegalStateException(
                    String.format("Master process is not running on the host %s", host));
              } else if (output.contains("Connection refused")) {
                throw new IllegalStateException(
                    String.format("Connection refused while connecting to the host %s on port %d",
                        host, port));
              }
              LOG.debug("Master running on node {} with pid={}", host, output);
            }
            CommonUtils.sleepMs(Constants.SECOND_MS);
          }
        } catch (Throwable e) {
          LOG.error("Exception thrown in the master process check", e);
          throw new RuntimeException(e);
        }
      }
    }
  
    /**
     * Creates a new instance of MasterHealthCheckClient.
     *
     * @param alluxioMasterType the Alluxio master process type
     * @param processCheck whether to check the AlluxioMaster process is alive
     * @param retryPolicySupplier the policy supplier to utilize for the rpc check
     * @param alluxioConf Alluxio configuration
     */
    public MasterHealthCheckClient(MasterType alluxioMasterType, boolean processCheck,
        Supplier<RetryPolicy> retryPolicySupplier, AlluxioConfiguration alluxioConf) {
      mAlluxioMasterType = alluxioMasterType;
      mProcessCheck = processCheck;
      mExecutorService = Executors.newFixedThreadPool(2);
      mPolicySupplier = retryPolicySupplier;
      mConf = alluxioConf;
    }
  
    @Override
    public boolean isServing() {
      try {
        ServiceType rpcService;
        InetSocketAddress connectAddr;
        switch (mAlluxioMasterType) {
          case MASTER:
            rpcService = ServiceType.FILE_SYSTEM_MASTER_CLIENT_SERVICE;
            connectAddr = NetworkAddressUtils
              .getConnectAddress(NetworkAddressUtils.ServiceType.MASTER_RPC, mConf);
            break;
          case JOB_MASTER:
            rpcService = ServiceType.JOB_MASTER_CLIENT_SERVICE;
            connectAddr = NetworkAddressUtils
                .getConnectAddress(NetworkAddressUtils.ServiceType.JOB_MASTER_RPC, mConf);
            break;
          default:
            throw new IllegalArgumentException(
                String.format("Master type %s is invalid", mAlluxioMasterType));
        }
        MasterServingHealthCheck masterRpcCheck =
            new MasterServingHealthCheck(connectAddr, rpcService,
                mPolicySupplier,
                mConf);
        Future<?> masterServingFuture = mExecutorService.submit(masterRpcCheck);
        if (mProcessCheck) {
          Future<?> processCheckFuture = mExecutorService.submit(
                  new ProcessCheckRunnable(mAlluxioMasterType.getClassName()));
          CommonUtils.sleepMs(Constants.SECOND_MS);
          // If in HA mode, can't check the RPC service, because the service may not have started
          if (!ConfigurationUtils.isHaMode(mConf)) {
            while (!masterServingFuture.isDone()) {
              if (processCheckFuture.isDone()) {
                throw new IllegalStateException("One or more master processes are not running");
              }
              CommonUtils.sleepMs(Constants.SECOND_MS);
            }
          }
          // Check after a 7 second period if the process is is still alive
          CommonUtils.sleepMs(7L * Constants.SECOND_MS);
          LOG.debug("Checking the master processes one more time...");
          return !processCheckFuture.isDone();
        } else {
          masterServingFuture.get();
          return masterRpcCheck.serving();
        }
      } catch (Exception e) {
        LOG.error("Exception thrown in master health check client", e);
      } finally {
        mExecutorService.shutdown();
      }
      return false;
    }
  }
  