[fix] remove in address2Region while bookie left to get correct rack info

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RegionAwareEnsemblePlacementPolicy.java

@@ -117,6 +117,23 @@ protected String parseBookieRegion(BookieId addr) {
 
     @Override
     public void handleBookiesThatLeft(Set<BookieId> leftBookies) {
+        // case 1: (In some situation, eg, Broker and bookie restart concurrently.)
+        //1. Bookie X join cluster for the first time, encounters a region exception, and `address2Region` record X's
+        // region as default-region.
+        //2. Bookie X left cluster and is removed from knownBookies, but address2Region retains the information of
+        // bookie X.
+        //3. update Bookie X's rack info, and calling `onBookieRackChange` will only update address2Region for
+        // addresses present in knownBookies; therefore, bookie X's region info is not updated.
+        //4. Bookie X join cluster again, since address2Region contains the previous default-region information,
+        // getRegion will directly use cached data, resulting of an incorrect region.
+
+        // The bookie region is initialized to "default-region" in address2Region.
+        // We should ensure that when a bookie leaves the cluster,
+        // we also clean up the corresponding region information for that bookie in address2Region,
+        // so that it can update the correct region for the bookie during onBookieRackChange and
+        // handleBookiesThatJoined.
+        // to avoid traffic skew in ensemble selection.
+        leftBookies.forEach(address2Region::remove);
         super.handleBookiesThatLeft(leftBookies);
 
         for (TopologyAwareEnsemblePlacementPolicy policy: perRegionPlacement.values()) {

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

@@ -29,6 +29,8 @@
 import static org.apache.bookkeeper.feature.SettableFeatureProvider.DISABLE_ALL;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.mock;
@@ -2052,4 +2054,121 @@ public void testNewEnsemblePickLocalRegionBookies()
         LOG.info("Bookie1 Count: {}, Bookie8 Count: {}, Bookie9 Count: {}", bookie1Count, bookie8Count, bookie9Count);
 
     }
+
+    @Test
+    public void testBookieLeftThenJoinWithDNSResolveFailed() throws Exception {
+
+        BookieSocketAddress addr1 = new BookieSocketAddress("127.0.1.1", 3181);
+        BookieSocketAddress addr2 = new BookieSocketAddress("127.0.1.2", 3181);
+        BookieSocketAddress addr3 = new BookieSocketAddress("127.0.1.3", 3181);
+        BookieSocketAddress addr4 = new BookieSocketAddress("127.0.1.4", 3181);
+
+        // init dns mapping
+        // 2. mock dns resolver failed, use default region and rack.
+        // addr1 rack info. /region-1/default-rack -> /default-region/default-rack.
+
+        // 1. mock addr1 dns resolver failed and use default region and rack.
+        StaticDNSResolver.addNodeToRack(addr1.getHostName(), "/default-region/default-rack");
+        StaticDNSResolver.addNodeToRack(addr2.getHostName(), "/region-1/default-rack");
+        StaticDNSResolver.addNodeToRack(addr3.getHostName(), "/region-2/default-rack");
+        StaticDNSResolver.addNodeToRack(addr4.getHostName(), "/region-3/default-rack");
+
+        // init cluster
+        Set<BookieId> addrs = Sets.newHashSet(addr1.toBookieId(),
+                addr2.toBookieId(), addr3.toBookieId(), addr4.toBookieId());
+        repp.onClusterChanged(addrs, new HashSet<>());
+
+        assertEquals(4, repp.knownBookies.size());
+        assertEquals("/default-region/default-rack", repp.knownBookies.get(addr1.toBookieId()).getNetworkLocation());
+        assertEquals("/region-1/default-rack", repp.knownBookies.get(addr2.toBookieId()).getNetworkLocation());
+        assertEquals("/region-2/default-rack", repp.knownBookies.get(addr3.toBookieId()).getNetworkLocation());
+        assertEquals("/region-3/default-rack", repp.knownBookies.get(addr4.toBookieId()).getNetworkLocation());
+
+        assertEquals(4, repp.perRegionPlacement.size());
+        TopologyAwareEnsemblePlacementPolicy unknownRegionPlacement = repp.perRegionPlacement.get("UnknownRegion");
+        assertEquals(1, unknownRegionPlacement.knownBookies.keySet().size());
+        assertEquals("/default-region/default-rack",
+                unknownRegionPlacement.knownBookies.get(addr1.toBookieId()).getNetworkLocation());
+
+        TopologyAwareEnsemblePlacementPolicy region1Placement = repp.perRegionPlacement.get("region-1");
+        assertEquals(1, region1Placement.knownBookies.keySet().size());
+        assertEquals("/region-1/default-rack",
+                region1Placement.knownBookies.get(addr2.toBookieId()).getNetworkLocation());
+
+        TopologyAwareEnsemblePlacementPolicy region2Placement = repp.perRegionPlacement.get("region-2");
+        assertEquals(1, region2Placement.knownBookies.keySet().size());
+        assertEquals("/region-2/default-rack",
+                region2Placement.knownBookies.get(addr3.toBookieId()).getNetworkLocation());
+
+        TopologyAwareEnsemblePlacementPolicy region3Placement = repp.perRegionPlacement.get("region-3");
+        assertEquals(1, region3Placement.knownBookies.keySet().size());
+        assertEquals("/region-3/default-rack",
+                region3Placement.knownBookies.get(addr4.toBookieId()).getNetworkLocation());
+
+        assertEquals("UnknownRegion", repp.address2Region.get(addr1.toBookieId()));
+        assertEquals("region-1", repp.address2Region.get(addr2.toBookieId()));
+        assertEquals("region-2", repp.address2Region.get(addr3.toBookieId()));
+        assertEquals("region-3", repp.address2Region.get(addr4.toBookieId()));
+
+        // 2. addr1 bookie shutdown and decommission
+        addrs.remove(addr1.toBookieId());
+        repp.onClusterChanged(addrs, new HashSet<>());
+
+        assertEquals(3, repp.knownBookies.size());
+        assertNull(repp.knownBookies.get(addr1.toBookieId()));
+        assertEquals("/region-1/default-rack", repp.knownBookies.get(addr2.toBookieId()).getNetworkLocation());
+        assertEquals("/region-2/default-rack", repp.knownBookies.get(addr3.toBookieId()).getNetworkLocation());
+        assertEquals("/region-3/default-rack", repp.knownBookies.get(addr4.toBookieId()).getNetworkLocation());
+
+        // UnknownRegion,region-1,region-2,region-3
+        assertEquals(4, repp.perRegionPlacement.size());
+        // after addr1 bookie left, it should remove from locally address2Region
+        assertNull(repp.address2Region.get(addr1.toBookieId()));
+        assertEquals("region-1", repp.address2Region.get(addr2.toBookieId()));
+        assertEquals("region-2", repp.address2Region.get(addr3.toBookieId()));
+        assertEquals("region-3", repp.address2Region.get(addr4.toBookieId()));
+
+
+        // 3. addr1 bookie start and join
+        addrs.add(addr1.toBookieId());
+        repp.onClusterChanged(addrs, new HashSet<>());
+
+        assertEquals(4, repp.knownBookies.size());
+        assertEquals("/default-region/default-rack", repp.knownBookies.get(addr1.toBookieId()).getNetworkLocation());
+        assertEquals("/region-1/default-rack", repp.knownBookies.get(addr2.toBookieId()).getNetworkLocation());
+        assertEquals("/region-2/default-rack", repp.knownBookies.get(addr3.toBookieId()).getNetworkLocation());
+        assertEquals("/region-3/default-rack", repp.knownBookies.get(addr4.toBookieId()).getNetworkLocation());
+
+        // UnknownRegion,region-1,region-2,region-3
+        assertEquals(4, repp.perRegionPlacement.size());
+        assertEquals("UnknownRegion", repp.address2Region.get(addr1.toBookieId()));
+        // addr1 bookie belongs to unknown region
+        unknownRegionPlacement = repp.perRegionPlacement.get("UnknownRegion");
+        assertEquals(1, unknownRegionPlacement.knownBookies.keySet().size());
+        assertEquals("/default-region/default-rack",
+                unknownRegionPlacement.knownBookies.get(addr1.toBookieId()).getNetworkLocation());
+
+        // 4. Update the correct rack.
+        // change addr1 rack info. /default-region/default-rack -> /region-1/default-rack.
+        List<BookieSocketAddress> bookieAddressList = new ArrayList<>();
+        List<String> rackList = new ArrayList<>();
+        bookieAddressList.add(addr1);
+        rackList.add("/region-1/default-rack");
+        // onBookieRackChange
+        StaticDNSResolver.changeRack(bookieAddressList, rackList);
+
+        assertEquals(4, repp.perRegionPlacement.size());
+        // addr1 bookie, oldRegion=default-region, newRegion=region-1
+        assertEquals("region-1", repp.address2Region.get(addr1.toBookieId()));
+
+        unknownRegionPlacement = repp.perRegionPlacement.get("UnknownRegion");
+        assertEquals(0, unknownRegionPlacement.knownBookies.keySet().size());
+        assertNotNull(unknownRegionPlacement.historyBookies.get(addr1.toBookieId()));
+
+
+        region1Placement = repp.perRegionPlacement.get("region-1");
+        assertEquals(2, region1Placement.knownBookies.keySet().size());
+        assertEquals("/region-1/default-rack",
+                region1Placement.knownBookies.get(addr1.toBookieId()).getNetworkLocation());
+    }
 }