Fix disk weight ensemble infinite loop bug

Author: hangc0276

bookkeeper-server/src/test/java/org/apache/bookkeeper/client/TestRackawareEnsemblePlacementPolicy.java

@@ -17,6 +17,7 @@
  */
 package org.apache.bookkeeper.client;
 
+import static org.apache.bookkeeper.client.RackawareEnsemblePlacementPolicyImpl.RACKNAME_DISTANCE_FROM_LEAVES;
 import static org.apache.bookkeeper.client.RackawareEnsemblePlacementPolicyImpl.REPP_DNS_RESOLVER_CLASS;
 import static org.apache.bookkeeper.client.RackawareEnsemblePlacementPolicyImpl.shuffleWithMask;
 import static org.apache.bookkeeper.client.RoundRobinDistributionSchedule.writeSetFromValues;
@@ -69,6 +70,7 @@
 import org.hamcrest.Matcher;
 import org.hamcrest.TypeSafeMatcher;
 import org.junit.Test;
+import org.junit.jupiter.api.Timeout;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -790,6 +792,55 @@ public void testNewEnsembleWithSingleRack() throws Exception {
         }
     }
 
+    @Test(timeout = 30_000)
+    public void testNewEnsembleWithExcludeBookies() throws Exception {
+        repp.uninitalize();
+        updateMyRack(NetworkTopology.DEFAULT_REGION_AND_RACK);
+
+        repp = new RackawareEnsemblePlacementPolicy();
+        conf.setDiskWeightBasedPlacementEnabled(true);
+        repp.initialize(conf, Optional.<DNSToSwitchMapping>empty(), timer,
+                DISABLE_ALL, NullStatsLogger.INSTANCE, BookieSocketAddress.LEGACY_BOOKIEID_RESOLVER);
+        repp.withDefaultRack(NetworkTopology.DEFAULT_REGION_AND_RACK);
+
+        BookieSocketAddress addr1 = new BookieSocketAddress("127.0.0.2", 3181);
+        BookieSocketAddress addr2 = new BookieSocketAddress("127.0.0.3", 3181);
+        BookieSocketAddress addr3 = new BookieSocketAddress("127.0.0.4", 3181);
+
+        BookieNode addr1Node = new BookieNode(addr1.toBookieId(), repp.resolveNetworkLocation(addr1.toBookieId()));
+        BookieNode addr2Node = new BookieNode(addr2.toBookieId(), repp.resolveNetworkLocation(addr2.toBookieId()));
+        BookieNode addr3Node = new BookieNode(addr3.toBookieId(), repp.resolveNetworkLocation(addr3.toBookieId()));
+        // update dns mapping
+        StaticDNSResolver.addNodeToRack(addr1.getHostName(), "/default-region/r1");
+        StaticDNSResolver.addNodeToRack(addr2.getHostName(), "/default-region/r2");
+        StaticDNSResolver.addNodeToRack(addr3.getHostName(), "/default-region/r1");
+        // Update cluster
+        Set<BookieId> addrs = new HashSet<BookieId>();
+        addrs.add(addr1.toBookieId());
+        addrs.add(addr2.toBookieId());
+        addrs.add(addr3.toBookieId());
+        repp.onClusterChanged(addrs, new HashSet<BookieId>());
+
+        Set<Node> excludeBookies = new HashSet<>();
+        excludeBookies.add(addr2Node);
+        excludeBookies.add(addr3Node);
+
+        TopologyAwareEnsemblePlacementPolicy.RRTopologyAwareCoverageEnsemble ensemble =
+                new TopologyAwareEnsemblePlacementPolicy.RRTopologyAwareCoverageEnsemble(
+                        2, 2, 2,
+                        RACKNAME_DISTANCE_FROM_LEAVES,
+                        null, null, 1);
+        ensemble.addNode(new BookieNode(addr1.toBookieId(), repp.resolveNetworkLocation(addr1.toBookieId())));
+        try {
+           repp.selectRandomInternal(null, 1, excludeBookies, null, ensemble);
+           fail("Should fail with not enough bookies exception");
+        } catch (BKNotEnoughBookiesException ex) {
+            //
+        }
+
+        conf.setDiskWeightBasedPlacementEnabled(false);
+    }
+
     @Test
     public void testSingleRackWithEnforceMinNumRacks() throws Exception {
         repp.uninitalize();