From 9cbd0cd2a9268ff2e8fed0af335e9c4f91c5f601 Mon Sep 17 00:00:00 2001
From: Inigo Goiri <inigoiri@apache.org>
Date: Tue, 12 May 2020 08:20:03 -0700
Subject: [PATCH] YARN-9301. Too many InvalidStateTransitionException with SLS.
 Contributed by Bilwa S T.

---
 .../sls/resourcemanager/MockAMLauncher.java   | 57 ++++++++++---------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/resourcemanager/MockAMLauncher.java b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/resourcemanager/MockAMLauncher.java
index 208629afdb4..24c795b5600 100644
--- a/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/resourcemanager/MockAMLauncher.java
+++ b/hadoop-tools/hadoop-sls/src/main/java/org/apache/hadoop/yarn/sls/resourcemanager/MockAMLauncher.java
@@ -82,36 +82,39 @@ private void setupAMRMToken(RMAppAttempt appAttempt) {
   @Override
   @SuppressWarnings("unchecked")
   public void handle(AMLauncherEvent event) {
-    if (AMLauncherEventType.LAUNCH == event.getType()) {
-      ApplicationId appId =
-          event.getAppAttempt().getAppAttemptId().getApplicationId();
+    ApplicationId appId =
+        event.getAppAttempt().getAppAttemptId().getApplicationId();
+    // find AMSimulator
+    AMSimulator ams = appIdAMSim.get(appId);
+    if (ams == null) {
+      throw new YarnRuntimeException(
+          "Didn't find any AMSimulator for applicationId=" + appId);
+    }
+    Container amContainer = event.getAppAttempt().getMasterContainer();
+    switch (event.getType()) {
+    case LAUNCH:
+      try {
+        setupAMRMToken(event.getAppAttempt());
+        // Notify RMAppAttempt to change state
+        super.context.getDispatcher().getEventHandler().handle(
+            new RMAppAttemptEvent(event.getAppAttempt().getAppAttemptId(),
+                RMAppAttemptEventType.LAUNCHED));
 
-      // find AMSimulator
-      AMSimulator ams = appIdAMSim.get(appId);
-      if (ams != null) {
-        try {
-          Container amContainer = event.getAppAttempt().getMasterContainer();
+        ams.notifyAMContainerLaunched(
+            event.getAppAttempt().getMasterContainer());
+        LOG.info("Notify AM launcher launched:" + amContainer.getId());
 
-          setupAMRMToken(event.getAppAttempt());
-
-          // Notify RMAppAttempt to change state
-          super.context.getDispatcher().getEventHandler().handle(
-              new RMAppAttemptEvent(event.getAppAttempt().getAppAttemptId(),
-                  RMAppAttemptEventType.LAUNCHED));
-
-          ams.notifyAMContainerLaunched(
-              event.getAppAttempt().getMasterContainer());
-          LOG.info("Notify AM launcher launched:" + amContainer.getId());
-
-          se.getNmMap().get(amContainer.getNodeId())
-              .addNewContainer(amContainer, 100000000L);
-
-          return;
-        } catch (Exception e) {
-          throw new YarnRuntimeException(e);
-        }
+        se.getNmMap().get(amContainer.getNodeId())
+            .addNewContainer(amContainer, -1);
+        return;
+      } catch (Exception e) {
+        throw new YarnRuntimeException(e);
       }
-
+    case CLEANUP:
+      se.getNmMap().get(amContainer.getNodeId())
+          .cleanupContainer(amContainer.getId());
+      break;
+    default:
       throw new YarnRuntimeException(
           "Didn't find any AMSimulator for applicationId=" + appId);
     }