From e91e8b711c68273460b36557fc37fdfc86be097b Mon Sep 17 00:00:00 2001
From: Jason Lowe <jlowe@apache.org>
Date: Fri, 22 Jan 2016 21:36:15 +0000
Subject: [PATCH] YARN-4371. "yarn application -kill" should take multiple
 application ids. Contributed by Sunil G

---
 hadoop-yarn-project/CHANGES.txt               |   3 +
 .../yarn/client/cli/ApplicationCLI.java       |  57 ++++++-
 .../hadoop/yarn/client/cli/TestYarnCLI.java   | 141 ++++++++++++++++--
 3 files changed, 183 insertions(+), 18 deletions(-)

diff --git a/hadoop-yarn-project/CHANGES.txt b/hadoop-yarn-project/CHANGES.txt
index a4f9874781c..0ceb9057f4a 100644
--- a/hadoop-yarn-project/CHANGES.txt
+++ b/hadoop-yarn-project/CHANGES.txt
@@ -733,6 +733,9 @@ Release 2.8.0 - UNRELEASED
 
     YARN-4524. Cleanup AppSchedulingInfo. (Karthik Kambatla via wangda)
 
+    YARN-4371. "yarn application -kill" should take multiple application ids
+    (Sunil G via jlowe)
+
   OPTIMIZATIONS
 
     YARN-3339. TestDockerContainerExecutor should pull a single image and not
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ApplicationCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ApplicationCLI.java
index 77e06885d68..caa4d466444 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ApplicationCLI.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/ApplicationCLI.java
@@ -23,6 +23,7 @@ import java.io.OutputStreamWriter;
 import java.io.PrintWriter;
 import java.nio.charset.Charset;
 import java.text.DecimalFormat;
+import java.util.Collection;
 import java.util.EnumSet;
 import java.util.HashSet;
 import java.util.List;
@@ -103,7 +104,6 @@ public class ApplicationCLI extends YarnCLI {
           + "Supports optional use of -appTypes to filter applications "
           + "based on application type, "
           + "and -appStates to filter applications based on application state.");
-      opts.addOption(KILL_CMD, true, "Kills the application.");
       opts.addOption(MOVE_TO_QUEUE_CMD, true, "Moves the application to a "
           + "different queue.");
       opts.addOption(QUEUE_CMD, true, "Works with the movetoqueue command to"
@@ -127,7 +127,12 @@ public class ApplicationCLI extends YarnCLI {
       opts.addOption(UPDATE_PRIORITY, true,
           "update priority of an application. ApplicationId can be"
               + " passed using 'appId' option.");
-      opts.getOption(KILL_CMD).setArgName("Application ID");
+      Option killOpt = new Option(KILL_CMD, true, "Kills the application. "
+          + "Set of applications can be provided separated with space");
+      killOpt.setValueSeparator(' ');
+      killOpt.setArgs(Option.UNLIMITED_VALUES);
+      killOpt.setArgName("Application ID");
+      opts.addOption(killOpt);
       opts.getOption(MOVE_TO_QUEUE_CMD).setArgName("Application ID");
       opts.getOption(QUEUE_CMD).setArgName("Queue Name");
       opts.getOption(STATUS_CMD).setArgName("Application ID");
@@ -239,15 +244,11 @@ public class ApplicationCLI extends YarnCLI {
         listContainers(cliParser.getOptionValue(LIST_CMD));
       }
     } else if (cliParser.hasOption(KILL_CMD)) {
-      if (args.length != 3) {
+      if (args.length < 3 || hasAnyOtherCLIOptions(cliParser, opts, KILL_CMD)) {
         printUsage(title, opts);
         return exitCode;
       }
-      try{
-        killApplication(cliParser.getOptionValue(KILL_CMD));
-      } catch (ApplicationNotFoundException e) {
-        return exitCode;
-      }
+      return killApplication(cliParser.getOptionValues(KILL_CMD));
     } else if (cliParser.hasOption(MOVE_TO_QUEUE_CMD)) {
       if (!cliParser.hasOption(QUEUE_CMD)) {
         printUsage(title, opts);
@@ -481,6 +482,30 @@ public class ApplicationCLI extends YarnCLI {
     writer.flush();
   }
 
+  /**
+   * Kills applications with the application id as appId
+   *
+   * @param Array of applicationIds
+   * @return errorCode
+   * @throws YarnException
+   * @throws IOException
+   */
+  private int killApplication(String[] applicationIds) throws YarnException,
+      IOException {
+    int returnCode = -1;
+    for (String applicationId : applicationIds) {
+      try {
+        killApplication(applicationId);
+        returnCode = 0;
+      } catch (ApplicationNotFoundException e) {
+        // Suppress all ApplicationNotFoundException for now.
+        continue;
+      }
+    }
+
+    return returnCode;
+  }
+
   /**
    * Kills the application with the application id as appId
    * 
@@ -726,4 +751,20 @@ public class ApplicationCLI extends YarnCLI {
           + " as application is in final states");
     }
   }
+
+  @SuppressWarnings("unchecked")
+  private boolean hasAnyOtherCLIOptions(CommandLine cliParser, Options opts,
+      String excludeOption) {
+    Collection<Option> ops = opts.getOptions();
+    for (Option op : ops) {
+      // Skip exclude option from the option list
+      if (op.getOpt().equals(excludeOption)) {
+        continue;
+      }
+      if (cliParser.hasOption(op.getOpt())) {
+        return true;
+      }
+    }
+    return false;
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java
index 95ccccfd6be..3783aac659e 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestYarnCLI.java
@@ -744,14 +744,6 @@ public class TestYarnCLI {
     Assert.assertEquals(createApplicationCLIHelpMessage(),
         sysOutStream.toString());
 
-    sysOutStream.reset();
-    ApplicationId applicationId = ApplicationId.newInstance(1234, 5);
-    result = cli.run(
-        new String[] {"application", "-kill", applicationId.toString(), "args" });
-    verify(spyCli).printUsage(any(String.class), any(Options.class));
-    Assert.assertEquals(createApplicationCLIHelpMessage(),
-        sysOutStream.toString());
-
     sysOutStream.reset();
     NodeId nodeId = NodeId.newInstance("host0", 0);
     result = cli.run(
@@ -878,7 +870,134 @@ public class TestYarnCLI {
       Assert.fail("Unexpected exception: " + e);
     }
   }
-  
+
+  @Test
+  public void testKillApplications() throws Exception {
+    ApplicationCLI cli = createAndGetAppCLI();
+    ApplicationId applicationId1 = ApplicationId.newInstance(1234, 5);
+    ApplicationId applicationId2 = ApplicationId.newInstance(1234, 6);
+    ApplicationId applicationId3 = ApplicationId.newInstance(1234, 7);
+    ApplicationId applicationId4 = ApplicationId.newInstance(1234, 8);
+
+    // Test Scenario 1: Both applications are FINISHED.
+    ApplicationReport newApplicationReport1 = ApplicationReport.newInstance(
+        applicationId1, ApplicationAttemptId.newInstance(applicationId1, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.FINISHED, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53789f, "YARN", null);
+    ApplicationReport newApplicationReport2 = ApplicationReport.newInstance(
+        applicationId2, ApplicationAttemptId.newInstance(applicationId2, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.FINISHED, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.34344f, "YARN", null);
+    when(client.getApplicationReport(applicationId1)).thenReturn(
+        newApplicationReport1);
+    when(client.getApplicationReport(applicationId2)).thenReturn(
+        newApplicationReport2);
+    int result = cli.run(new String[]{"application", "-kill",
+        applicationId1.toString() + " " + applicationId2.toString()});
+    assertEquals(0, result);
+    verify(client, times(0)).killApplication(applicationId1);
+    verify(client, times(0)).killApplication(applicationId2);
+    verify(sysOut).println(
+        "Application " + applicationId1 + " has already finished ");
+    verify(sysOut).println(
+        "Application " + applicationId2 + " has already finished ");
+
+    // Test Scenario 2: Both applications are RUNNING.
+    ApplicationReport newApplicationReport3 = ApplicationReport.newInstance(
+        applicationId1, ApplicationAttemptId.newInstance(applicationId1, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.RUNNING, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53789f, "YARN", null);
+    ApplicationReport newApplicationReport4 = ApplicationReport.newInstance(
+        applicationId2, ApplicationAttemptId.newInstance(applicationId2, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.RUNNING, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53345f, "YARN", null);
+    when(client.getApplicationReport(applicationId1)).thenReturn(
+        newApplicationReport3);
+    when(client.getApplicationReport(applicationId2)).thenReturn(
+        newApplicationReport4);
+    result = cli.run(new String[]{"application", "-kill",
+        applicationId1.toString() + " " + applicationId2.toString()});
+    assertEquals(0, result);
+    verify(client).killApplication(applicationId1);
+    verify(client).killApplication(applicationId2);
+    verify(sysOut).println(
+        "Killing application application_1234_0005");
+    verify(sysOut).println(
+        "Killing application application_1234_0006");
+
+    // Test Scenario 3: Both applications are not present.
+    doThrow(new ApplicationNotFoundException("Application with id '"
+        + applicationId3 + "' doesn't exist in RM.")).when(client)
+        .getApplicationReport(applicationId3);
+    doThrow(new ApplicationNotFoundException("Application with id '"
+        + applicationId4 + "' doesn't exist in RM.")).when(client)
+        .getApplicationReport(applicationId4);
+    result = cli.run(new String[]{"application", "-kill",
+        applicationId3.toString() + " " + applicationId4.toString()});
+    Assert.assertNotEquals(0, result);
+    verify(sysOut).println(
+        "Application with id 'application_1234_0007' doesn't exist in RM.");
+    verify(sysOut).println(
+        "Application with id 'application_1234_0008' doesn't exist in RM.");
+
+    // Test Scenario 4: one application is not present and other RUNNING
+    doThrow(new ApplicationNotFoundException("Application with id '"
+        + applicationId3 + "' doesn't exist in RM.")).when(client)
+        .getApplicationReport(applicationId3);
+    ApplicationReport newApplicationReport5 = ApplicationReport.newInstance(
+        applicationId1, ApplicationAttemptId.newInstance(applicationId1, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.RUNNING, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53345f, "YARN", null);
+    when(client.getApplicationReport(applicationId1)).thenReturn(
+        newApplicationReport5);
+    result = cli.run(new String[]{"application", "-kill",
+        applicationId3.toString() + " " + applicationId1.toString()});
+    Assert.assertEquals(0, result);
+
+    // Test Scenario 5: kill operation with some other command.
+    sysOutStream.reset();
+    result = cli.run(new String[]{"application", "--appStates", "RUNNING",
+        "-kill", applicationId3.toString() + " " + applicationId1.toString()});
+    Assert.assertEquals(-1, result);
+    Assert.assertEquals(createApplicationCLIHelpMessage(),
+        sysOutStream.toString());
+  }
+
+  @Test
+  public void testKillApplicationsOfDifferentEndStates() throws Exception {
+    ApplicationCLI cli = createAndGetAppCLI();
+    ApplicationId applicationId1 = ApplicationId.newInstance(1234, 5);
+    ApplicationId applicationId2 = ApplicationId.newInstance(1234, 6);
+
+    // Scenario: One application is FINISHED and other is RUNNING.
+    ApplicationReport newApplicationReport5 = ApplicationReport.newInstance(
+        applicationId1, ApplicationAttemptId.newInstance(applicationId1, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.FINISHED, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53789f, "YARN", null);
+    ApplicationReport newApplicationReport6 = ApplicationReport.newInstance(
+        applicationId2, ApplicationAttemptId.newInstance(applicationId2, 1),
+        "user", "queue", "appname", "host", 124, null,
+        YarnApplicationState.RUNNING, "diagnostics", "url", 0, 0,
+        FinalApplicationStatus.SUCCEEDED, null, "N/A", 0.53345f, "YARN", null);
+    when(client.getApplicationReport(applicationId1)).thenReturn(
+        newApplicationReport5);
+    when(client.getApplicationReport(applicationId2)).thenReturn(
+        newApplicationReport6);
+    int result = cli.run(new String[]{"application", "-kill",
+        applicationId1.toString() + " " + applicationId2.toString()});
+    assertEquals(0, result);
+    verify(client, times(1)).killApplication(applicationId2);
+    verify(sysOut).println(
+        "Application " + applicationId1 + " has already finished ");
+    verify(sysOut).println("Killing application application_1234_0006");
+  }
+
   @Test
   public void testMoveApplicationAcrossQueues() throws Exception {
     ApplicationCLI cli = createAndGetAppCLI();
@@ -1694,7 +1813,9 @@ public class TestYarnCLI {
     pw.println("                                 based on input comma-separated list of");
     pw.println("                                 application types.");
     pw.println(" -help                           Displays help for all commands.");
-    pw.println(" -kill <Application ID>          Kills the application.");
+    pw.println(" -kill <Application ID>          Kills the application. Set of");
+    pw.println("                                 applications can be provided separated");
+    pw.println("                                 with space");
     pw.println(" -list                           List applications. Supports optional use");
     pw.println("                                 of -appTypes to filter applications based");
     pw.println("                                 on application type, and -appStates to");