alluxio-master-2.x/shell/src/main/java/alluxio/cli/LogLevel.java
/*
 * 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.cli;

import alluxio.ClientContext;
import alluxio.Constants;
import alluxio.annotation.PublicApi;
import alluxio.client.block.BlockWorkerInfo;
import alluxio.client.file.FileSystemContext;
import alluxio.client.job.JobMasterClient;
import alluxio.conf.AlluxioConfiguration;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.job.wire.JobWorkerHealth;
import alluxio.util.network.HttpUtils;
import alluxio.util.network.NetworkAddressUtils;
import alluxio.util.network.NetworkAddressUtils.ServiceType;
import alluxio.wire.LogInfo;
import alluxio.wire.WorkerNetAddress;
import alluxio.worker.job.JobMasterClientContext;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Objects;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.DefaultParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.client.utils.URIBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.concurrent.NotThreadSafe;

/**
 * Sets or gets the log level for the specified server.
 */
@NotThreadSafe
@PublicApi
public final class LogLevel {
  private static final Logger LOG = LoggerFactory.getLogger(LogLevel.class);

  public static final String LOG_LEVEL = "logLevel";
  public static final String ROLE_WORKERS = "workers";
  public static final String ROLE_MASTER = "master";
  public static final String ROLE_MASTERS = "masters";
  public static final String ROLE_WORKER = "worker";
  public static final String ROLE_JOB_MASTER = "job_master";
  public static final String ROLE_JOB_MASTERS = "job_masters";
  public static final String ROLE_JOB_WORKER = "job_worker";
  public static final String ROLE_JOB_WORKERS = "job_workers";
  public static final String ROLE_FUSE = "fuse";
  public static final String TARGET_SEPARATOR = ",";
  public static final String TARGET_OPTION_NAME = "target";
  private static final Option TARGET_OPTION =
      Option.builder()
          .required(false)
          .longOpt(TARGET_OPTION_NAME)
          .hasArg(true)
          .desc("<master|workers|job_master|job_workers|host:webPort[:role]>."
              + " A list of targets separated by " + TARGET_SEPARATOR + " can be specified."
              + " Default target is master, job master, all workers and all job workers.")
          .build();
  private static final String LOG_NAME_OPTION_NAME = "logName";
  private static final Option LOG_NAME_OPTION =
      Option.builder()
          .required(true)
          .longOpt(LOG_NAME_OPTION_NAME)
          .hasArg(true)
          .desc("The logger's name(e.g. alluxio.master.file.DefaultFileSystemMaster)"
              + " you want to get or set level.")
          .build();
  private static final String LEVEL_OPTION_NAME = "level";
  private static final Option LEVEL_OPTION =
      Option.builder()
          .required(false)
          .longOpt(LEVEL_OPTION_NAME)
          .hasArg(true)
          .desc("The log level to be set.").build();
  private static final Options OPTIONS = new Options()
      .addOption(TARGET_OPTION)
      .addOption(LOG_NAME_OPTION)
      .addOption(LEVEL_OPTION);

  /**
   * Prints the help message.
   *
   * @param message message before standard usage information
   */
  public static void printHelp(String message) {
    System.err.println(message);
    HelpFormatter help = new HelpFormatter();
    help.printHelp(LOG_LEVEL, OPTIONS, true);
  }

  /**
   * Implements log level setting and getting.
   *
   * @param args list of arguments contains target, logName and level
   * @param alluxioConf configurations
   * @exception ParseException if there is an error in parsing
   */
  public static void logLevel(String[] args, AlluxioConfiguration alluxioConf)
      throws ParseException, IOException {
    CommandLineParser parser = new DefaultParser();
    CommandLine cmd = parser.parse(OPTIONS, args, true /* stopAtNonOption */);

    List<TargetInfo> targets = parseOptTarget(cmd, alluxioConf);
    String logName = parseOptLogName(cmd);
    String level = parseOptLevel(cmd);

    for (TargetInfo targetInfo : targets) {
      setLogLevel(targetInfo, logName, level);
    }
  }

  /**
   * Parse the command options to a list of targets endpoints.
   *
   * @param cmd the command
   * @param conf Alluxio configuration
   * @return a list of {@link TargetInfo}
   * */
  public static List<TargetInfo> parseOptTarget(CommandLine cmd, AlluxioConfiguration conf)
      throws IOException {
    String[] targets;
    if (cmd.hasOption(TARGET_OPTION_NAME)) {
      String argTarget = cmd.getOptionValue(TARGET_OPTION_NAME);
      if (StringUtils.isBlank(argTarget)) {
        throw new IOException("Option " + TARGET_OPTION_NAME + " can not be blank.");
      } else if (argTarget.contains(TARGET_SEPARATOR)) {
        targets = argTarget.split(TARGET_SEPARATOR);
      } else {
        targets = new String[]{argTarget};
      }
    } else {
      // By default we set on all targets (master/workers/job_master/job_workers)
      targets = new String[]{ROLE_MASTER, ROLE_JOB_MASTER, ROLE_WORKERS, ROLE_JOB_WORKERS};
    }
    return getTargetInfos(targets, conf);
  }

  private static List<TargetInfo> getTargetInfos(String[] targets, AlluxioConfiguration conf)
      throws IOException {
    // Trim the elements
    Set<String> targetSet = Arrays.stream(targets).map(String::trim).collect(Collectors.toSet());
    List<TargetInfo> targetInfoList = new ArrayList<>();

    // Allow plural form for the master/job_master and print a notice
    if (targetSet.contains(ROLE_MASTERS)) {
      System.out.println("The logLevel command will only take effect on the primary master, "
              + "instead of on all the masters. ");
      targetSet.remove(ROLE_MASTERS);
      targetSet.add(ROLE_MASTER);
      System.out.println("Target `masters` is replaced with `master`.");
    }
    if (targetSet.contains(ROLE_JOB_MASTERS)) {
      System.out.println("The logLevel command will only take effect on the primary job master, "
              + "instead of on all the masters. ");
      targetSet.remove(ROLE_JOB_MASTERS);
      targetSet.add(ROLE_JOB_MASTER);
      System.out.println("Target `job_masters` is replaced with `job_master`.");
    }

    ClientContext clientContext = ClientContext.create(conf);
    // Created only when needed by master and workers
    FileSystemContext fsContext = null;
    // Created only when needed by the job master and job workers
    JobMasterClient jobClient = null;

    // Process each target
    for (String target : targetSet) {
      if (target.isEmpty()) {
        continue;
      } else if (target.equals(ROLE_MASTER)) {
        if (fsContext == null) {
          fsContext = FileSystemContext.create(clientContext);
        }
        String masterHost = fsContext.getMasterAddress().getHostName();
        int masterPort = NetworkAddressUtils.getPort(ServiceType.MASTER_WEB, conf);
        TargetInfo master = new TargetInfo(masterHost, masterPort, ROLE_MASTER);
        targetInfoList.add(master);
      } else if (target.equals(ROLE_JOB_MASTER)) {
        if (jobClient == null) {
          jobClient = JobMasterClient.Factory.create(JobMasterClientContext
                  .newBuilder(clientContext).build());
        }
        String jobMasterHost = jobClient.getRemoteHostName();
        int jobMasterPort = NetworkAddressUtils.getPort(ServiceType.JOB_MASTER_WEB, conf);
        TargetInfo jobMaster = new TargetInfo(jobMasterHost, jobMasterPort, ROLE_JOB_MASTER);
        targetInfoList.add(jobMaster);
      } else if (target.equals(ROLE_WORKERS)) {
        if (fsContext == null) {
          fsContext = FileSystemContext.create(ClientContext.create(conf));
        }
        List<BlockWorkerInfo> workerInfoList = fsContext.getCachedWorkers();
        if (workerInfoList.size() == 0) {
          System.out.println("No workers found");
          System.exit(1);
        }
        for (BlockWorkerInfo workerInfo : workerInfoList) {
          WorkerNetAddress netAddress = workerInfo.getNetAddress();
          TargetInfo worker = new TargetInfo(netAddress.getHost(),
                  netAddress.getWebPort(), ROLE_WORKER);
          targetInfoList.add(worker);
        }
      } else if (target.equals(ROLE_JOB_WORKERS)) {
        if (jobClient == null) {
          jobClient = JobMasterClient.Factory.create(JobMasterClientContext
                  .newBuilder(clientContext).build());
        }
        List<JobWorkerHealth> jobWorkerInfoList = jobClient.getAllWorkerHealth();
        if (jobWorkerInfoList.size() == 0) {
          System.out.println("No job workers found");
          System.exit(1);
        }
        int jobWorkerPort = conf.getInt(PropertyKey.JOB_WORKER_WEB_PORT);
        for (JobWorkerHealth jobWorkerInfo : jobWorkerInfoList) {
          String jobWorkerHost = jobWorkerInfo.getHostname();
          TargetInfo jobWorker = new TargetInfo(jobWorkerHost, jobWorkerPort, ROLE_JOB_WORKER);
          targetInfoList.add(jobWorker);
        }
      } else if (target.contains(":")) {
        String[] targetInfoParts = target.split(":");
        int port = Integer.parseInt(targetInfoParts[1]);
        String role;
        if (targetInfoParts.length > 2) {
          role = targetInfoParts[2];
        } else {
          role = inferRoleFromPort(port, conf);
        }
        LOG.debug("Port {} maps to role {}", port, role);
        TargetInfo unspecifiedTarget = new TargetInfo(targetInfoParts[0], port, role);
        System.out.format("Role inferred from port: %s%n", unspecifiedTarget);
        targetInfoList.add(unspecifiedTarget);
      } else {
        throw new IOException(String.format("Unrecognized target argument: %s. "
                + "Please pass the targets in the form of <host>:<port>, "
                + "with comma as the separator.", target));
      }
    }
    return targetInfoList;
  }

  private static String parseOptLogName(CommandLine cmd) {
    String argName = cmd.getOptionValue(LOG_NAME_OPTION_NAME);
    if (StringUtils.isNotBlank(argName)) {
      return argName;
    }
    return "";
  }

  private static String parseOptLevel(CommandLine cmd) {
    if (cmd.hasOption(LEVEL_OPTION_NAME)) {
      String argLevel = cmd.getOptionValue(LEVEL_OPTION_NAME);
      if (StringUtils.isNotBlank(argLevel)) {
        return argLevel;
      }
    }
    return null;
  }

  private static void setLogLevel(final TargetInfo targetInfo, String logName, String level)
      throws IOException {
    URIBuilder uriBuilder = new URIBuilder();
    uriBuilder.setScheme("http");
    uriBuilder.setHost(targetInfo.getHost());
    uriBuilder.setPort(targetInfo.getPort());
    uriBuilder.setPath(Constants.REST_API_PREFIX + "/" + targetInfo.getRole() + "/" + LOG_LEVEL);
    uriBuilder.addParameter(LOG_NAME_OPTION_NAME, logName);
    if (level != null) {
      uriBuilder.addParameter(LEVEL_OPTION_NAME, level);
    }
    LOG.info("Setting log level on {}", uriBuilder);
    HttpUtils.post(uriBuilder.toString(), "", 5000, inputStream -> {
      ObjectMapper mapper = new ObjectMapper();
      LogInfo logInfo = mapper.readValue(inputStream, LogInfo.class);
      System.out.println(targetInfo + logInfo.toString());
    });
  }

  private static String inferRoleFromPort(int port, AlluxioConfiguration conf) {
    if (port == NetworkAddressUtils.getPort(ServiceType.MASTER_WEB, conf)) {
      return ROLE_MASTER;
    } else if (port == NetworkAddressUtils.getPort(ServiceType.WORKER_WEB, conf)) {
      return ROLE_WORKER;
    } else if (port == NetworkAddressUtils.getPort(ServiceType.JOB_MASTER_WEB, conf)) {
      return ROLE_JOB_MASTER;
    } else if (port == NetworkAddressUtils.getPort(ServiceType.JOB_WORKER_WEB, conf)) {
      return ROLE_JOB_WORKER;
    } else if (port == NetworkAddressUtils.getPort(ServiceType.FUSE_WEB, conf)) {
      return ROLE_FUSE;
    } else {
      throw new IllegalArgumentException(String.format(
              "Unrecognized port in %s. Please make sure the port is in %s",
              port,
              Arrays.toString(new PropertyKey[]{
                PropertyKey.MASTER_WEB_PORT, PropertyKey.WORKER_WEB_PORT,
                PropertyKey.JOB_MASTER_WEB_PORT, PropertyKey.JOB_WORKER_WEB_PORT
              })));
    }
  }

  /**
   * Sets or gets log level of master and worker through their REST API.
   *
   * @param args same arguments as {@link LogLevel}
   */
  public static void main(String[] args) {
    int exitCode = 1;
    try {
      logLevel(args, Configuration.global());
      exitCode = 0;
    } catch (ParseException e) {
      printHelp("Unable to parse input args: " + e.getMessage());
    } catch (IOException e) {
      System.err.println("Failed to set log level:");
      e.printStackTrace();
    }
    System.exit(exitCode);
  }

  private LogLevel() {} // this class is not intended for instantiation

  /**
   * Object that represents a REST endpoint that logLevel sends HTTP request to.
   * */
  public static final class TargetInfo {
    private final String mRole;
    private final String mHost;
    private final int mPort;

    /**
     * Constructor.
     *
     * @param host hostname
     * @param port port
     * @param role the Alluxio component name(part of URL)
     * */
    public TargetInfo(String host, int port, String role) {
      mHost = host;
      mPort = port;
      mRole = role;
    }

    /**
     * @return port
     * */
    public int getPort() {
      return mPort;
    }

    /**
     * @return hostname
     * */
    public String getHost() {
      return mHost;
    }

    /**
     * @return component name
     * */
    public String getRole() {
      return mRole;
    }

    @Override
    public String toString() {
      return mHost + ":" + mPort + "[" + mRole + "]";
    }

    @Override
    public boolean equals(Object other) {
      if (!(other instanceof TargetInfo)) {
        return false;
      }
      TargetInfo otherTarget = (TargetInfo) other;
      return mRole.equals(otherTarget.mRole) && mHost.equals(otherTarget.mHost)
              && mPort == otherTarget.mPort;
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(mRole, mHost, mPort);
    }
  }
}
