YARN-9206. RMServerUtils does not count SHUTDOWN as an accepted state. Contributed by Kuhu Shukla.

(cherry picked from commit 6ffe6ea8999c1216b6e510e350535754431bac2d)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeState.java

@@ -55,4 +55,16 @@ public enum NodeState {
     return (this == UNHEALTHY || this == DECOMMISSIONED
         || this == LOST || this == SHUTDOWN);
   }
+
+  public boolean isInactiveState() {
+    return this == NodeState.DECOMMISSIONED ||
+      this == NodeState.LOST || this == NodeState.REBOOTED ||
+      this == NodeState.SHUTDOWN;
+  }
+
+  public boolean isActiveState() {
+    return this == NodeState.NEW ||
+        this == NodeState.RUNNING || this == NodeState.UNHEALTHY ||
+        this == NodeState.DECOMMISSIONING;
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/RMServerUtils.java

@@ -106,10 +106,20 @@ public class RMServerUtils {
       EnumSet<NodeState> acceptedStates) {
     // nodes contains nodes that are NEW, RUNNING, UNHEALTHY or DECOMMISSIONING.
     ArrayList<RMNode> results = new ArrayList<RMNode>();
-    if (acceptedStates.contains(NodeState.NEW) ||
-        acceptedStates.contains(NodeState.RUNNING) ||
-        acceptedStates.contains(NodeState.DECOMMISSIONING) ||
-        acceptedStates.contains(NodeState.UNHEALTHY)) {
+    boolean hasActive = false;
+    boolean hasInactive = false;
+    for (NodeState nodeState : acceptedStates) {
+      if (!hasInactive && nodeState.isInactiveState()) {
+        hasInactive = true;
+      }
+      if (!hasActive && nodeState.isActiveState()) {
+        hasActive = true;
+      }
+      if (hasActive && hasInactive) {
+        break;
+      }
+    }
+    if (hasActive) {
       for (RMNode rmNode : context.getRMNodes().values()) {
         if (acceptedStates.contains(rmNode.getState())) {
           results.add(rmNode);
@@ -118,9 +128,7 @@ public class RMServerUtils {
     }
 
     // inactiveNodes contains nodes that are DECOMMISSIONED, LOST, OR REBOOTED
-    if (acceptedStates.contains(NodeState.DECOMMISSIONED) ||
-        acceptedStates.contains(NodeState.LOST) ||
-        acceptedStates.contains(NodeState.REBOOTED)) {
+    if (hasInactive) {
       for (RMNode rmNode : context.getInactiveRMNodes().values()) {
         if ((rmNode != null) && acceptedStates.contains(rmNode.getState())) {
           results.add(rmNode);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/RMWebServices.java

@@ -382,9 +382,7 @@ public class RMWebServices extends WebServices implements RMWebServiceProtocol {
     NodesInfo nodesInfo = new NodesInfo();
     for (RMNode rmNode : rmNodes) {
       NodeInfo nodeInfo = new NodeInfo(rmNode, sched);
-      if (EnumSet
-          .of(NodeState.LOST, NodeState.DECOMMISSIONED, NodeState.REBOOTED)
-          .contains(rmNode.getState())) {
+      if (rmNode.getState().isInactiveState()) {
         nodeInfo.setNodeHTTPAddress(RMWSConsts.EMPTY);
       }
       nodesInfo.add(nodeInfo);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMServerUtils.java

@@ -19,11 +19,13 @@
 package org.apache.hadoop.yarn.server.resourcemanager;
 
 import org.apache.hadoop.yarn.api.records.NodeId;
+import org.apache.hadoop.yarn.api.records.NodeState;
 import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager;
+import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
 import org.junit.Assert;
 import org.junit.Test;
@@ -31,6 +33,9 @@ import org.mockito.Mockito;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -107,6 +112,37 @@ public class TestRMServerUtils {
         RMServerUtils.getApplicableNodeCountForAM(rmContext, conf, reqs));
   }
 
+  @Test
+  public void testQueryRMNodes() throws Exception {
+    RMContext rmContext = Mockito.mock(RMContext.class);
+    NodeId node1 = NodeId.newInstance("node1", 1234);
+    RMNode rmNode1 = Mockito.mock(RMNode.class);
+    ConcurrentMap<NodeId, RMNode> inactiveList =
+        new ConcurrentHashMap<NodeId, RMNode>();
+    Mockito.when(rmNode1.getState()).thenReturn(NodeState.SHUTDOWN);
+    inactiveList.put(node1, rmNode1);
+    Mockito.when(rmContext.getInactiveRMNodes()).thenReturn(inactiveList);
+    List<RMNode> result = RMServerUtils.queryRMNodes(rmContext,
+        EnumSet.of(NodeState.SHUTDOWN));
+    Assert.assertTrue(result.size() != 0);
+    Assert.assertEquals(result.get(0), rmNode1);
+    Mockito.when(rmNode1.getState()).thenReturn(NodeState.DECOMMISSIONED);
+    result = RMServerUtils.queryRMNodes(rmContext,
+        EnumSet.of(NodeState.DECOMMISSIONED));
+    Assert.assertTrue(result.size() != 0);
+    Assert.assertEquals(result.get(0), rmNode1);
+    Mockito.when(rmNode1.getState()).thenReturn(NodeState.LOST);
+    result = RMServerUtils.queryRMNodes(rmContext,
+        EnumSet.of(NodeState.LOST));
+    Assert.assertTrue(result.size() != 0);
+    Assert.assertEquals(result.get(0), rmNode1);
+    Mockito.when(rmNode1.getState()).thenReturn(NodeState.REBOOTED);
+    result = RMServerUtils.queryRMNodes(rmContext,
+        EnumSet.of(NodeState.REBOOTED));
+    Assert.assertTrue(result.size() != 0);
+    Assert.assertEquals(result.get(0), rmNode1);
+  }
+
   @Test
   public void testGetApplicableNodeCountForAMLabels() throws Exception {
     Set<NodeId> noLabelNodes = new HashSet<>();