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

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;
@@ -37,6 +38,7 @@ import org.apache.hadoop.security.HttpCrossOriginFilterInitializer;
 import org.apache.hadoop.security.SecurityUtil;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.service.CompositeService;
+import org.apache.hadoop.util.GenericOptionsParser;
 import org.apache.hadoop.util.JvmPauseMonitor;
 import org.apache.hadoop.util.ShutdownHookManager;
 import org.apache.hadoop.util.StringUtils;
@@ -292,7 +294,19 @@ public class GlobalPolicyGenerator extends CompositeService {
 
   public static void main(String[] argv) {
     try {
-      startGPG(argv, new YarnConfiguration());
+      YarnConfiguration conf = new YarnConfiguration();
+      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.
+          System.err.println("format-policy-store is not yet supported.");
+        } else {
+          printUsage(System.err);
+        }
+      } else {
+        startGPG(argv, conf);
+      }
     } catch (Throwable t) {
       LOG.error("Error starting global policy generator", t);
       System.exit(-1);
@@ -307,4 +321,8 @@ public class GlobalPolicyGenerator extends CompositeService {
   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

@@ -26,9 +26,13 @@ import org.apache.hadoop.test.LambdaTestUtils;
 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.
  */
@@ -58,6 +62,17 @@ public class TestGlobalPolicyGenerator {
     }, 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]"));
+  }
+
   @Test
   public void testUserProvidedUGIConf() throws Exception {
     String errMsg = "Invalid attribute value for " +

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.URL;
 import java.net.InetSocketAddress;
@@ -36,6 +37,7 @@ import org.apache.hadoop.security.HttpCrossOriginFilterInitializer;
 import org.apache.hadoop.security.SecurityUtil;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.service.CompositeService;
+import org.apache.hadoop.util.GenericOptionsParser;
 import org.apache.hadoop.util.JvmPauseMonitor;
 import org.apache.hadoop.util.ShutdownHookManager;
 import org.apache.hadoop.util.StringUtils;
@@ -295,18 +297,29 @@ public class Router extends CompositeService {
     StringUtils.startupShutdownMessage(Router.class, argv, LOG);
     Router router = new Router();
     try {
-
-      // 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.
+          System.err.println("format-state-store is not yet supported.");
+        } 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.
+          System.err.println("remove-application-from-state-store is not yet supported.");
+        } 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);
@@ -348,4 +361,9 @@ public class Router extends CompositeService {
   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;
@@ -43,7 +44,9 @@ import javax.servlet.ServletOutputStream;
 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;
@@ -390,4 +393,15 @@ public class TestRouter {
     }
   }
 
+  @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>]"));
+  }
 }