From 3dd40008846a65503408cf80eb09b04e56cb47a3 Mon Sep 17 00:00:00 2001
From: Inigo Goiri <inigoiri@apache.org>
Date: Fri, 4 Jan 2019 09:55:09 -0800
Subject: [PATCH] HADOOP-16028. Fix NetworkTopology chooseRandom function to
 support excluded nodes. Contributed by Sihai Ke.

(cherry picked from commit f4e18242bd8117a5c506ec6d3f25c85011fa82d0)
---
 .../apache/hadoop/net/NetworkTopology.java    |  9 +++--
 .../hadoop/net/TestClusterTopology.java       | 35 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

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 1f077a753bb..f7f60ecc0a6 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
@@ -532,8 +532,13 @@ public class NetworkTopology {
     if (excludedScope == null) {
       availableNodes = countNumOfAvailableNodes(scope, excludedNodes);
     } else {
-      availableNodes =
-          countNumOfAvailableNodes("~" + excludedScope, excludedNodes);
+      netlock.readLock().lock();
+      try {
+        availableNodes = countNumOfAvailableNodes(scope, excludedNodes) -
+            countNumOfAvailableNodes(excludedScope, excludedNodes);
+      } finally {
+        netlock.readLock().unlock();
+      }
     }
     LOG.debug("Choosing random from {} available nodes on node {},"
         + " scope={}, excludedScope={}, excludeNodes={}. numOfDatanodes={}.",
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/net/TestClusterTopology.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/net/TestClusterTopology.java
index 6aad6c5c173..fbed6052a5c 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/net/TestClusterTopology.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/net/TestClusterTopology.java
@@ -20,6 +20,7 @@ package org.apache.hadoop.net;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Arrays;
 
 import org.apache.commons.math3.stat.inference.ChiSquareTest;
 import org.apache.hadoop.conf.Configuration;
@@ -194,6 +195,40 @@ public class TestClusterTopology extends Assert {
         2, histogram.size());
   }
 
+  @Test
+  public void testChooseRandomExcluded() {
+    // create the topology
+    //                        a1
+    //                b1------|--------b2
+    //                 |                |
+    //          c1-----|-----c2         c3
+    //         /  \          |          |
+    //        /    \         |          |
+    //     node1    node2   node3      node4
+
+    NetworkTopology cluster = NetworkTopology.getInstance(new Configuration());
+    NodeElement node1 = getNewNode("node1", "/a1/b1/c1");
+    cluster.add(node1);
+    NodeElement node2 = getNewNode("node2", "/a1/b1/c1");
+    cluster.add(node2);
+    NodeElement node3 = getNewNode("node3", "/a1/b1/c2");
+    cluster.add(node3);
+    NodeElement node4 = getNewNode("node4", "/a1/b2/c3");
+    cluster.add(node4);
+
+    Node node = cluster.chooseRandom("/a1/b1", "/a1/b1/c1", null);
+    assertSame("node3", node.getName());
+
+    node = cluster.chooseRandom("/a1/b1", "/a1/b1/c1", Arrays.asList(node1));
+    assertSame("node3", node.getName());
+
+    node = cluster.chooseRandom("/a1/b1", "/a1/b1/c1", Arrays.asList(node3));
+    assertNull(node);
+
+    node = cluster.chooseRandom("/a1/b1", "/a1/b1/c1", Arrays.asList(node4));
+    assertSame("node3", node.getName());
+  }
+
   private NodeElement getNewNode(String name, String rackLocation) {
     NodeElement node = new NodeElement(name);
     node.setNetworkLocation(rackLocation);