YARN-5697. Use CliParser to parse options in RMAdminCLI. Contributed by Tao Jie.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/RMAdminCLI.java

@@ -29,6 +29,12 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.MissingArgumentException;
+import org.apache.commons.cli.Option;
+import org.apache.commons.cli.Options;
+import org.apache.commons.cli.ParseException;
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
 import org.apache.hadoop.conf.Configuration;
@@ -80,7 +86,6 @@ public class RMAdminCLI extends HAAdmin {
 
   private final RecordFactory recordFactory = 
     RecordFactoryProvider.getRecordFactory(null);
-  private boolean directlyAccessNodeLabelStore = false;
   static CommonNodeLabelsManager localNodeLabelsManager = null;
   private static final String NO_LABEL_ERR_MSG =
       "No cluster node-labels are specified";
@@ -130,8 +135,9 @@ public class RMAdminCLI extends HAAdmin {
               new UsageInfo("<label1,label2,label3> (label splitted by \",\")",
                   "remove from cluster node labels"))
           .put("-replaceLabelsOnNode",
-              new UsageInfo("[-failOnUnknownNodes] " +
-                  "<\"node1[:port]=label1,label2 node2[:port]=label1,label2\">",
+              new UsageInfo(
+                  "<\"node1[:port]=label1,label2 node2[:port]=label1,label2\"> "
+                  + "[-failOnUnknownNodes] ",
               "replace labels on nodes"
                   + " (please note that we do not support specifying multiple"
                   + " labels on a single host for now.)\n\t\t"
@@ -248,8 +254,9 @@ public class RMAdminCLI extends HAAdmin {
       " [-addToClusterNodeLabels <\"label1(exclusive=true),"
                   + "label2(exclusive=false),label3\">]" +
       " [-removeFromClusterNodeLabels <label1,label2,label3>]" +
-      " [-replaceLabelsOnNode [-failOnUnknownNodes] "
-          + "<\"node1[:port]=label1,label2 node2[:port]=label1\">]" +
+      " [-replaceLabelsOnNode " +
+            "<\"node1[:port]=label1,label2 node2[:port]=label1\"> " +
+            "[-failOnUnknownNodes]]" +
       " [-directlyAccessNodeLabelStore]" +
       " [-refreshClusterMaxPriority]" +
       " [-updateNodeResource [NodeID] [MemSize] [vCores] ([OvercommitTimeout])");
@@ -564,11 +571,26 @@ public class RMAdminCLI extends HAAdmin {
     return labels;
   }
 
-  private int addToClusterNodeLabels(String args) throws IOException,
-      YarnException {
-    List<NodeLabel> labels = buildNodeLabelsFromStr(args);
+  private int handleAddToClusterNodeLabels(String[] args, String cmd,
+      boolean isHAEnabled) throws IOException, YarnException, ParseException {
+    Options opts = new Options();
+    opts.addOption("addToClusterNodeLabels", true,
+        "Add to cluster node labels.");
+    opts.addOption("directlyAccessNodeLabelStore", false,
+        "Directly access node label store.");
+    int exitCode = -1;
+    CommandLine cliParser = null;
+    try {
+      cliParser = new GnuParser().parse(opts, args);
+    } catch (MissingArgumentException ex) {
+      System.err.println(NO_LABEL_ERR_MSG);
+      printUsage(args[0], isHAEnabled);
+      return exitCode;
+    }
 
-    if (directlyAccessNodeLabelStore) {
+    List<NodeLabel> labels = buildNodeLabelsFromStr(
+        cliParser.getOptionValue("addToClusterNodeLabels"));
+    if (cliParser.hasOption("directlyAccessNodeLabelStore")) {
       getNodeLabelManagerInstance(getConf()).addToCluserNodeLabels(labels);
     } else {
       ResourceManagerAdministrationProtocol adminProtocol =
@@ -580,11 +602,26 @@ public class RMAdminCLI extends HAAdmin {
     return 0;
   }
 
-  private int removeFromClusterNodeLabels(String args) throws IOException,
-      YarnException {
-    Set<String> labels = buildNodeLabelNamesFromStr(args);
+  private int handleRemoveFromClusterNodeLabels(String[] args, String cmd,
+      boolean isHAEnabled) throws IOException, YarnException, ParseException {
+    Options opts = new Options();
+    opts.addOption("removeFromClusterNodeLabels", true,
+        "Remove From cluster node labels.");
+    opts.addOption("directlyAccessNodeLabelStore", false,
+        "Directly access node label store.");
+    int exitCode = -1;
+    CommandLine cliParser = null;
+    try {
+      cliParser = new GnuParser().parse(opts, args);
+    } catch (MissingArgumentException ex) {
+      System.err.println(NO_LABEL_ERR_MSG);
+      printUsage(args[0], isHAEnabled);
+      return exitCode;
+    }
 
-    if (directlyAccessNodeLabelStore) {
+    Set<String> labels = buildNodeLabelNamesFromStr(
+        cliParser.getOptionValue("removeFromClusterNodeLabels"));
+    if (cliParser.hasOption("directlyAccessNodeLabelStore")) {
       getNodeLabelManagerInstance(getConf()).removeFromClusterNodeLabels(
           labels);
     } else {
@@ -647,14 +684,35 @@ public class RMAdminCLI extends HAAdmin {
     return map;
   }
 
-  private int replaceLabelsOnNodes(String args, boolean failOnUnknownNodes)
-      throws IOException, YarnException {
-    Map<NodeId, Set<String>> map = buildNodeLabelsMapFromStr(args);
-    return replaceLabelsOnNodes(map, failOnUnknownNodes);
+  private int handleReplaceLabelsOnNodes(String[] args, String cmd,
+      boolean isHAEnabled) throws IOException, YarnException, ParseException {
+    Options opts = new Options();
+    opts.addOption("replaceLabelsOnNode", true,
+        "Replace label on node.");
+    opts.addOption("failOnUnknownNodes", false,
+        "Fail on unknown nodes.");
+    opts.addOption("directlyAccessNodeLabelStore", false,
+        "Directly access node label store.");
+    int exitCode = -1;
+    CommandLine cliParser = null;
+    try {
+      cliParser = new GnuParser().parse(opts, args);
+    } catch (MissingArgumentException ex) {
+      System.err.println(NO_MAPPING_ERR_MSG);
+      printUsage(args[0], isHAEnabled);
+      return exitCode;
+    }
+
+    Map<NodeId, Set<String>> map = buildNodeLabelsMapFromStr(
+        cliParser.getOptionValue("replaceLabelsOnNode"));
+    return replaceLabelsOnNodes(map,
+        cliParser.hasOption("failOnUnknownNodes"),
+        cliParser.hasOption("directlyAccessNodeLabelStore"));
   }
 
   private int replaceLabelsOnNodes(Map<NodeId, Set<String>> map,
-      boolean failOnUnknownNodes) throws IOException, YarnException {
+      boolean failOnUnknownNodes, boolean directlyAccessNodeLabelStore)
+      throws IOException, YarnException {
     if (directlyAccessNodeLabelStore) {
       getNodeLabelManagerInstance(getConf()).replaceLabelsOnNode(map);
     } else {
@@ -670,18 +728,6 @@ public class RMAdminCLI extends HAAdmin {
 
   @Override
   public int run(String[] args) throws Exception {
-    // -directlyAccessNodeLabelStore is a additional option for node label
-    // access, so just search if we have specified this option, and remove it
-    List<String> argsList = new ArrayList<String>();
-    for (int i = 0; i < args.length; i++) {
-      if (args[i].equals("-directlyAccessNodeLabelStore")) {
-        directlyAccessNodeLabelStore = true;
-      } else {
-        argsList.add(args[i]);
-      }
-    }
-    args = argsList.toArray(new String[0]);
-    
     YarnConfiguration yarnConf =
         getConf() == null ? new YarnConfiguration() : new YarnConfiguration(
             getConf());
@@ -735,28 +781,7 @@ public class RMAdminCLI extends HAAdmin {
       if ("-refreshQueues".equals(cmd)) {
         exitCode = refreshQueues();
       } else if ("-refreshNodes".equals(cmd)) {
-        if (args.length == 1) {
-          exitCode = refreshNodes();
-        } else if (args.length == 3 || args.length == 4) {
-          // if the graceful timeout specified
-          if ("-g".equals(args[1])) {
-            long timeout = -1;
-            String trackingMode;
-            if (args.length == 4) {
-              timeout = validateTimeout(args[2]);
-              trackingMode = validateTrackingMode(args[3]);
-            } else {
-              trackingMode = validateTrackingMode(args[2]);
-            }
-            exitCode = refreshNodes(timeout, trackingMode);
-          } else {
-            printUsage(cmd, isHAEnabled);
-            return -1;
-          }
-        } else {
-          printUsage(cmd, isHAEnabled);
-          return -1;
-        }
+        exitCode = handleRefreshNodes(args, cmd, isHAEnabled);
       } else if ("-refreshNodesResources".equals(cmd)) {
         exitCode = refreshNodesResources();
       } else if ("-refreshUserToGroupsMappings".equals(cmd)) {
@@ -773,54 +798,13 @@ public class RMAdminCLI extends HAAdmin {
         String[] usernames = Arrays.copyOfRange(args, i, args.length);
         exitCode = getGroups(usernames);
       } else if ("-updateNodeResource".equals(cmd)) {
-        if (args.length < 4 || args.length > 5) {
-          System.err.println("Number of parameters specified for " +
-              "updateNodeResource is wrong.");
-          printUsage(cmd, isHAEnabled);
-          exitCode = -1;
-        } else {
-          String nodeID = args[i++];
-          String memSize = args[i++];
-          String cores = args[i++];
-          int overCommitTimeout = ResourceOption.OVER_COMMIT_TIMEOUT_MILLIS_DEFAULT;
-          if (i == args.length - 1) {
-            overCommitTimeout = Integer.parseInt(args[i]);
-          }
-          exitCode = updateNodeResource(nodeID, Integer.parseInt(memSize),
-              Integer.parseInt(cores), overCommitTimeout);
-        }
+        exitCode = handleUpdateNodeResource(args, cmd, isHAEnabled);
       } else if ("-addToClusterNodeLabels".equals(cmd)) {
-        if (i >= args.length) {
-          System.err.println(NO_LABEL_ERR_MSG);
-          printUsage("", isHAEnabled);
-          exitCode = -1;
-        } else {
-          exitCode = addToClusterNodeLabels(args[i]);
-        }
+        exitCode = handleAddToClusterNodeLabels(args, cmd, isHAEnabled);
       } else if ("-removeFromClusterNodeLabels".equals(cmd)) {
-        if (i >= args.length) {
-          System.err.println(NO_LABEL_ERR_MSG);
-          printUsage("", isHAEnabled);
-          exitCode = -1;
-        } else {
-          exitCode = removeFromClusterNodeLabels(args[i]);
-        }
+        exitCode = handleRemoveFromClusterNodeLabels(args, cmd, isHAEnabled);
       } else if ("-replaceLabelsOnNode".equals(cmd)) {
-        if (i >= args.length) {
-          System.err.println(NO_MAPPING_ERR_MSG);
-          printUsage("", isHAEnabled);
-          exitCode = -1;
-        } else if ("-failOnUnknownNodes".equals(args[i])) {
-          if (i + 1 >= args.length) {
-            System.err.println(NO_MAPPING_ERR_MSG);
-            printUsage("", isHAEnabled);
-            exitCode = -1;
-          } else {
-            exitCode = replaceLabelsOnNodes(args[i + 1], true);
-          }
-        } else {
-          exitCode = replaceLabelsOnNodes(args[i], false);
-        }
+        exitCode = handleReplaceLabelsOnNodes(args, cmd, isHAEnabled);
       } else {
         exitCode = -1;
         System.err.println(cmd.substring(1) + ": Unknown command");
@@ -856,6 +840,52 @@ public class RMAdminCLI extends HAAdmin {
     return exitCode;
   }
 
+  // A helper method to reduce the number of lines of run()
+  private int handleRefreshNodes(String[] args, String cmd, boolean isHAEnabled)
+      throws IOException, YarnException, ParseException {
+    Options opts = new Options();
+    opts.addOption("refreshNodes", false,
+        "Refresh the hosts information at the ResourceManager.");
+    Option gracefulOpt = new Option("g", "graceful", true,
+        "Wait for timeout before marking the NodeManager as decommissioned.");
+    gracefulOpt.setOptionalArg(true);
+    opts.addOption(gracefulOpt);
+    opts.addOption("client", false,
+        "Indicates the timeout tracking should be handled by the client.");
+    opts.addOption("server", false,
+        "Indicates the timeout tracking should be handled by the RM.");
+
+    int exitCode = -1;
+    CommandLine cliParser = null;
+    try {
+      cliParser = new GnuParser().parse(opts, args);
+    } catch (MissingArgumentException ex) {
+      System.out.println("Missing argument for options");
+      printUsage(args[0], isHAEnabled);
+      return exitCode;
+    }
+
+    long timeout = -1;
+    if (cliParser.hasOption("g")) {
+      String strTimeout = cliParser.getOptionValue("g");
+      if (strTimeout != null) {
+        timeout = validateTimeout(strTimeout);
+      }
+      String trackingMode = null;
+      if (cliParser.hasOption("client")) {
+        trackingMode = "client";
+      } else if (cliParser.hasOption("server")) {
+        trackingMode = "server";
+      } else {
+        printUsage(cmd, isHAEnabled);
+        return -1;
+      }
+      return refreshNodes(timeout, trackingMode);
+    } else {
+      return refreshNodes();
+    }
+  }
+
   private long validateTimeout(String strTimeout) {
     long timeout;
     try {
@@ -869,6 +899,28 @@ public class RMAdminCLI extends HAAdmin {
     return timeout;
   }
 
+  private int handleUpdateNodeResource(String[] args, String cmd,
+      boolean isHAEnabled) throws NumberFormatException, IOException,
+      YarnException {
+    int i = 1;
+    if (args.length < 4 || args.length > 5) {
+      System.err.println("Number of parameters specified for "
+          + "updateNodeResource is wrong.");
+      printUsage(cmd, isHAEnabled);
+      return -1;
+    } else {
+      String nodeID = args[i++];
+      String memSize = args[i++];
+      String cores = args[i++];
+      int overCommitTimeout = ResourceOption.OVER_COMMIT_TIMEOUT_MILLIS_DEFAULT;
+      if (i == args.length - 1) {
+        overCommitTimeout = Integer.parseInt(args[i]);
+      }
+      return updateNodeResource(nodeID, Integer.parseInt(memSize),
+          Integer.parseInt(cores), overCommitTimeout);
+    }
+  }
+
   private String validateTrackingMode(String mode) {
     if ("-client".equals(mode)) {
       return "client";

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestRMAdminCLI.java

@@ -472,8 +472,9 @@ public class TestRMAdminCLI {
               "[username]] [-addToClusterNodeLabels " +
               "<\"label1(exclusive=true),label2(exclusive=false),label3\">] " +
               "[-removeFromClusterNodeLabels <label1,label2,label3>] " +
-              "[-replaceLabelsOnNode [-failOnUnknownNodes] " +
-              "<\"node1[:port]=label1,label2 node2[:port]=label1\">] " +
+              "[-replaceLabelsOnNode " +
+              "<\"node1[:port]=label1,label2 node2[:port]=label1\"> " +
+              "[-failOnUnknownNodes]] " +
               "[-directlyAccessNodeLabelStore] [-refreshClusterMaxPriority] " +
               "[-updateNodeResource [NodeID] [MemSize] [vCores] " +
               "([OvercommitTimeout]) [-help [cmd]]"));
@@ -566,8 +567,8 @@ public class TestRMAdminCLI {
               + " [username]] [-addToClusterNodeLabels <\"label1(exclusive=true),"
                   + "label2(exclusive=false),label3\">]"
               + " [-removeFromClusterNodeLabels <label1,label2,label3>] [-replaceLabelsOnNode "
-              + "[-failOnUnknownNodes] "
-              + "<\"node1[:port]=label1,label2 node2[:port]=label1\">] [-directlyAccessNodeLabelStore] "
+              + "<\"node1[:port]=label1,label2 node2[:port]=label1\"> "
+              + "[-failOnUnknownNodes]] [-directlyAccessNodeLabelStore] "
               + "[-refreshClusterMaxPriority] [-updateNodeResource [NodeID] [MemSize] [vCores] "
               + "([OvercommitTimeout]) [-transitionToActive [--forceactive] <serviceId>] "
               + "[-transitionToStandby <serviceId>] "
@@ -614,13 +615,11 @@ public class TestRMAdminCLI {
     dummyNodeLabelsManager.removeFromClusterNodeLabels(ImmutableSet.of("x", "y"));
     
     // change the sequence of "-directlyAccessNodeLabelStore" and labels,
-    // should not matter
+    // should fail
     args =
         new String[] { "-addToClusterNodeLabels",
             "-directlyAccessNodeLabelStore", "x,y" };
-    assertEquals(0, rmAdminCLI.run(args));
-    assertTrue(dummyNodeLabelsManager.getClusterNodeLabelNames().containsAll(
-        ImmutableSet.of("x", "y")));
+    assertEquals(-1, rmAdminCLI.run(args));
     
     // local node labels manager will be close after running
     assertTrue(dummyNodeLabelsManager.getServiceState() == STATE.STOPPED);
@@ -769,6 +768,10 @@ public class TestRMAdminCLI {
     args = new String[] { "-replaceLabelsOnNode" };
     assertTrue(0 != rmAdminCLI.run(args));
 
+    // no labels, should fail
+    args = new String[] { "-replaceLabelsOnNode", "-failOnUnknownNodes" };
+    assertTrue(0 != rmAdminCLI.run(args));
+
     // no labels, should fail
     args =
         new String[] { "-replaceLabelsOnNode", "-directlyAccessNodeLabelStore" };