From 4866735cde052e920eae689c8dc35577934ed8b5 Mon Sep 17 00:00:00 2001
From: bibinchundatt <bibinchundatt@apache.org>
Date: Tue, 16 Jul 2019 14:03:22 +0530
Subject: [PATCH] YARN-9645. Fix Invalid event FINISHED_CONTAINERS_PULLED_BY_AM
 at NEW on NM restart. Contributed by Bilwa S T.

(cherry picked from commit 7a93be0f6002ebb376c30f25a7d403e853c44280)
---
 .../resourcemanager/rmnode/RMNodeImpl.java     |  8 ++++++++
 .../resourcemanager/TestRMNodeTransitions.java | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
index 65a0c20bf79..93a621f54bc 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/rmnode/RMNodeImpl.java
@@ -205,6 +205,9 @@ public class RMNodeImpl implements RMNode, EventHandler<RMNodeEvent> {
       .addTransition(NodeState.NEW, NodeState.DECOMMISSIONED,
           RMNodeEventType.DECOMMISSION,
           new DeactivateNodeTransition(NodeState.DECOMMISSIONED))
+      .addTransition(NodeState.NEW, NodeState.NEW,
+          RMNodeEventType.FINISHED_CONTAINERS_PULLED_BY_AM,
+          new AddContainersToBeRemovedFromNMTransition())
 
       //Transitions from RUNNING state
       .addTransition(NodeState.RUNNING,
@@ -1557,4 +1560,9 @@ public class RMNodeImpl implements RMNode, EventHandler<RMNodeEvent> {
         attrMgr.getAttributesForNode(hostName);
     return nodeattrs.keySet();
   }
+
+  @VisibleForTesting
+  public Set<ContainerId> getContainersToBeRemovedFromNM() {
+    return containersToBeRemovedFromNM;
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMNodeTransitions.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMNodeTransitions.java
index 487d226072d..67decaf9316 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMNodeTransitions.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMNodeTransitions.java
@@ -195,6 +195,13 @@ public class TestRMNodeTransitions {
     return appIdList;
   }
 
+  private List<ContainerId> getContainerIdList() {
+    List<ContainerId> containerIdList = new ArrayList<ContainerId>();
+    containerIdList.add(BuilderUtils.newContainerId(BuilderUtils
+        .newApplicationAttemptId(BuilderUtils.newApplicationId(0, 0), 0), 0));
+    return containerIdList;
+  }
+
   private RMNodeStatusEvent getMockRMNodeStatusEventWithoutRunningApps() {
     NodeHealthStatus healthStatus = mock(NodeHealthStatus.class);
     Boolean yes = new Boolean(true);
@@ -1104,4 +1111,15 @@ public class TestRMNodeTransitions {
     Assert.assertEquals(1, hbrsp.getContainersToBeRemovedFromNM().size());
     Assert.assertEquals(0, node.getCompletedContainers().size());
   }
+
+  @Test
+  public void testFinishedContainersPulledByAMOnNewNode() {
+    RMNodeImpl rmNode = getNewNode();
+    NodeId nodeId = BuilderUtils.newNodeId("localhost", 0);
+
+    rmNode.handle(new RMNodeFinishedContainersPulledByAMEvent(nodeId,
+        getContainerIdList()));
+    Assert.assertEquals(1, rmNode.getContainersToBeRemovedFromNM().size());
+
+  }
 }