YARN-11433. Router's main() should support generic options.

Author: slfan1989

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-globalpolicygenerator/src/main/java/org/apache/hadoop/yarn/server/globalpolicygenerator/GlobalPolicyGenerator.java

@@ -19,6 +19,7 @@
 package org.apache.hadoop.yarn.server.globalpolicygenerator;
 
 import java.io.IOException;
+import java.io.PrintStream;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.ArrayList;
@@ -292,8 +293,17 @@ private String getHostName(Configuration config)
   public static void main(String[] argv) {
     try {
       YarnConfiguration conf = new YarnConfiguration();
-      new GenericOptionsParser(conf, argv);
-      startGPG(argv, conf);
+      GenericOptionsParser hParser = new GenericOptionsParser(conf, argv);
+      argv = hParser.getRemainingArgs();
+      if (argv.length > 1) {
+        if (argv[0].equals("-format-policy-store")) {
+          // TODO: YARN-11561. [Federation] GPG Supports Format PolicyStateStore.
+        } else {
+          printUsage(System.err);
+        }
+      } else {
+        startGPG(argv, conf);
+      }
     } catch (Throwable t) {
       LOG.error("Error starting global policy generator", t);
       System.exit(-1);
@@ -308,4 +318,8 @@ public static long getGPGStartupTime() {
   public WebApp getWebApp() {
     return webApp;
   }
+
+  private static void printUsage(PrintStream out) {
+    out.println("Usage: yarn gpg [-format-policy-store]");
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-globalpolicygenerator/src/test/java/org/apache/hadoop/yarn/server/globalpolicygenerator/TestGlobalPolicyGenerator.java

@@ -24,9 +24,13 @@
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.junit.Test;
 
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
 import java.util.List;
 import java.util.concurrent.TimeoutException;
 
+import static org.junit.Assert.assertTrue;
+
 /**
  * Unit test for GlobalPolicyGenerator.
  */
@@ -55,4 +59,15 @@ public void testGpgWithFederation() throws InterruptedException, TimeoutExceptio
       return (services.size() == 1 && gpg.getWebApp() != null);
     }, 100, 5000);
   }
+
+  @Test
+  public void testGPGCLI() {
+    ByteArrayOutputStream dataOut = new ByteArrayOutputStream();
+    ByteArrayOutputStream dataErr = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(dataOut));
+    System.setErr(new PrintStream(dataErr));
+    GlobalPolicyGenerator.main(new String[]{"-help", "-format-policy-store"});
+    assertTrue(dataErr.toString().contains(
+        "Usage: yarn gpg [-format-policy-store]"));
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-router/src/main/java/org/apache/hadoop/yarn/server/router/Router.java

@@ -19,6 +19,7 @@
 package org.apache.hadoop.yarn.server.router;
 
 import java.io.IOException;
+import java.io.PrintStream;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
@@ -246,20 +247,26 @@ public static void main(String[] argv) {
     StringUtils.startupShutdownMessage(Router.class, argv, LOG);
     Router router = new Router();
     try {
-
-      new GenericOptionsParser(conf, argv);
-
-      // Remove the old hook if we are rebooting.
-      if (null != routerShutdownHook) {
-        ShutdownHookManager.get().removeShutdownHook(routerShutdownHook);
+      GenericOptionsParser hParser = new GenericOptionsParser(conf, argv);
+      argv = hParser.getRemainingArgs();
+      if (argv.length > 1) {
+        if (argv[0].equals("-format-state-store")) {
+          // TODO: YARN-11548. [Federation] Router Supports Format FederationStateStore.
+        } else if (argv[0].equals("-remove-application-from-state-store") && argv.length == 2) {
+          // TODO: YARN-11547. [Federation] Router Supports Remove individual application records from FederationStateStore.
+        } else {
+          printUsage(System.err);
+        }
+      } else {
+        // Remove the old hook if we are rebooting.
+        if (null != routerShutdownHook) {
+          ShutdownHookManager.get().removeShutdownHook(routerShutdownHook);
+        }
+        routerShutdownHook = new CompositeServiceShutdownHook(router);
+        ShutdownHookManager.get().addShutdownHook(routerShutdownHook, SHUTDOWN_HOOK_PRIORITY);
+        router.init(conf);
+        router.start();
       }
-
-      routerShutdownHook = new CompositeServiceShutdownHook(router);
-      ShutdownHookManager.get().addShutdownHook(routerShutdownHook,
-          SHUTDOWN_HOOK_PRIORITY);
-
-      router.init(conf);
-      router.start();
     } catch (Throwable t) {
       LOG.error("Error starting Router", t);
       System.exit(-1);
@@ -301,4 +308,8 @@ public static long getClusterTimeStamp() {
   public FedAppReportFetcher getFetcher() {
     return fetcher;
   }
+
+  private static void printUsage(PrintStream out) {
+    out.println("Usage: yarn router [-format-state-store] | [-remove-application-from-state-store <appId>]");
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-router/src/test/java/org/apache/hadoop/yarn/server/router/TestRouter.java

@@ -18,6 +18,7 @@
 package org.apache.hadoop.yarn.server.router;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import org.apache.hadoop.conf.Configuration;
@@ -42,7 +43,9 @@
 import javax.servlet.http.Cookie;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.PrintStream;
 import java.io.PrintWriter;
 import java.util.Collection;
 import java.util.HashMap;
@@ -377,4 +380,14 @@ public Locale getLocale() {
     }
   }
 
+  @Test
+  public void testRouterCLI() {
+    ByteArrayOutputStream dataOut = new ByteArrayOutputStream();
+    ByteArrayOutputStream dataErr = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(dataOut));
+    System.setErr(new PrintStream(dataErr));
+    Router.main(new String[]{"-help", "-format-state-store"});
+    assertTrue(dataErr.toString().contains(
+        "Usage: yarn router [-format-state-store] | [-remove-application-from-state-store <appId>]"));
+  }
 }