From e750907d0a7e2fb4b33a7c876eaa4d9122a1deea Mon Sep 17 00:00:00 2001
From: Subru Krishnan <subru@apache.org>
Date: Tue, 30 May 2017 13:37:37 -0700
Subject: [PATCH] YARN-6666. Fix unit test failure in
 TestRouterClientRMService. (Botong Huang via Subru).

---
 .../server/MockResourceManagerFacade.java     | 38 +++++++------------
 1 file changed, 14 insertions(+), 24 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/MockResourceManagerFacade.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/MockResourceManagerFacade.java
index 9535ed5bd6e..65c12c642c7 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/MockResourceManagerFacade.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/MockResourceManagerFacade.java
@@ -21,8 +21,8 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -118,6 +118,7 @@
 import org.apache.hadoop.yarn.api.records.YarnApplicationAttemptState;
 import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.exceptions.ApplicationMasterNotRegisteredException;
+import org.apache.hadoop.yarn.exceptions.ApplicationNotFoundException;
 import org.apache.hadoop.yarn.exceptions.InvalidApplicationMasterRequestException;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
@@ -167,6 +168,7 @@ public class MockResourceManagerFacade implements ApplicationClientProtocol,
   private static final Logger LOG =
       LoggerFactory.getLogger(MockResourceManagerFacade.class);
 
+  private HashSet<ApplicationId> applicationMap = new HashSet<>();
   private HashMap<String, List<ContainerId>> applicationContainerIdMap =
       new HashMap<String, List<ContainerId>>();
   private HashMap<ContainerId, Container> allocatedContainerMap =
@@ -424,38 +426,26 @@ public GetNewApplicationResponse getNewApplication(
   @Override
   public SubmitApplicationResponse submitApplication(
       SubmitApplicationRequest request) throws YarnException, IOException {
+    ApplicationId appId = null;
+    if (request.getApplicationSubmissionContext() != null) {
+      appId = request.getApplicationSubmissionContext().getApplicationId();
+    }
+    LOG.info("Application submitted: " + appId);
+    applicationMap.add(appId);
     return SubmitApplicationResponse.newInstance();
   }
 
   @Override
   public KillApplicationResponse forceKillApplication(
       KillApplicationRequest request) throws YarnException, IOException {
-    String appId = "";
-    boolean foundApp = false;
+    ApplicationId appId = null;
     if (request.getApplicationId() != null) {
-      appId = request.getApplicationId().toString();
-      synchronized (applicationContainerIdMap) {
-        for (Entry<String, List<ContainerId>> entry : applicationContainerIdMap
-            .entrySet()) {
-          ApplicationAttemptId attemptId =
-              ApplicationAttemptId.fromString(entry.getKey());
-          if (attemptId.getApplicationId().equals(request.getApplicationId())) {
-            // Remove the apptempt and the containers that were being tracked
-            List<ContainerId> ids =
-                applicationContainerIdMap.remove(entry.getKey());
-            if (ids != null) {
-              for (ContainerId c : ids) {
-                allocatedContainerMap.remove(c);
-              }
-            }
-            foundApp = true;
-          }
-        }
+      appId = request.getApplicationId();
+      if (!applicationMap.remove(appId)) {
+        throw new ApplicationNotFoundException(
+            "Trying to kill an absent application: " + appId);
       }
     }
-    if (!foundApp) {
-      throw new YarnException("The application id is NOT registered: " + appId);
-    }
     LOG.info("Force killing application: " + appId);
     return KillApplicationResponse.newInstance(true);
   }