From 84f763884021980c456e2ebc21c1a1c1b18fec6c Mon Sep 17 00:00:00 2001
From: Ayush Saxena <ayushsaxena@apache.org>
Date: Tue, 18 Feb 2020 00:29:21 +0530
Subject: [PATCH] HADOOP-13666. Supporting rack exclusion in
 countNumOfAvailableNodes in NetworkTopology. Contributed by Inigo Goiri.

---
 .../org/apache/hadoop/net/NetworkTopology.java |  6 +++++-
 .../apache/hadoop/net/TestNetworkTopology.java | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/net/NetworkTopology.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/net/NetworkTopology.java
index 66799f5e2c9..aae56dd98da 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/net/NetworkTopology.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/net/NetworkTopology.java
@@ -670,7 +670,11 @@ public class NetworkTopology {
           }
           if ((NodeBase.getPath(node) + NodeBase.PATH_SEPARATOR_STR)
               .startsWith(scope + NodeBase.PATH_SEPARATOR_STR)) {
-            excludedCountInScope++;
+            if (node instanceof InnerNode) {
+              excludedCountInScope += ((InnerNode) node).getNumOfLeaves();
+            } else {
+              excludedCountInScope++;
+            }
           } else {
             excludedCountOffScope++;
           }
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/net/TestNetworkTopology.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/net/TestNetworkTopology.java
index f16bfb7069e..74c3f046ff0 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/net/TestNetworkTopology.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/net/TestNetworkTopology.java
@@ -614,4 +614,22 @@ public class TestNetworkTopology {
           frequency.get(dataNodes[i]) > 0);
     }
   }
+
+  @Test
+  public void testCountNumOfAvailableNodes() {
+    int numNodes = cluster.countNumOfAvailableNodes(NodeBase.ROOT, null);
+    assertEquals(20, numNodes);
+
+    // Excluding a single node
+    Collection<Node> excludedNodes = new HashSet<Node>();
+    excludedNodes.add(dataNodes[0]);
+    numNodes = cluster.countNumOfAvailableNodes(NodeBase.ROOT, excludedNodes);
+    assertEquals(19, numNodes);
+
+    // Excluding a full rack
+    Node d4r1 = cluster.getNode("/d4/r1");
+    excludedNodes.add(d4r1);
+    numNodes = cluster.countNumOfAvailableNodes(NodeBase.ROOT, excludedNodes);
+    assertEquals(12, numNodes);
+  }
 }