[fix][broker] Fix some problems in calculate totalAvailableBookies in method getExcludedBookiesWithIsolationGroups when some bookies belongs to multiple isolation groups. (apache#24091)

hanmz · lhotari · commit 6052e720f9aa · 2025-04-18T00:37:34.000+03:00
(cherry picked from commit 556024d)
diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/bookie/rackawareness/IsolatedBookieEnsemblePlacementPolicy.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/bookie/rackawareness/IsolatedBookieEnsemblePlacementPolicy.java
@@ -206,6 +206,9 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,
                     return excludedBookies;
                 }
                 Set<String> allGroups = allGroupsBookieMapping.keySet();
+                if (allGroups.isEmpty()) {
+                    return excludedBookies;
+                }
                 int totalAvailableBookiesInPrimaryGroup = 0;
                 Set<String> primaryIsolationGroup = Collections.emptySet();
                 Set<String> secondaryIsolationGroup = Collections.emptySet();
@@ -222,9 +225,10 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,
                         }
                     } else {
                         for (String groupBookie : bookiesInGroup) {
-                            totalAvailableBookiesInPrimaryGroup += knownBookies
-                                .containsKey(BookieId.parse(groupBookie)) ? 1 : 0;
-                            primaryGroupBookies.add(BookieId.parse(groupBookie));
+                            BookieId bookieId = BookieId.parse(groupBookie);
+                            if (primaryGroupBookies.add(bookieId)) {
+                                totalAvailableBookiesInPrimaryGroup += knownBookies.containsKey(bookieId) ? 1 : 0;
+                            }
                         }
                     }
                 }
@@ -256,8 +260,9 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,
                         Map<String, BookieInfo> bookieGroup = allGroupsBookieMapping.get(group);
                         if (bookieGroup != null && !bookieGroup.isEmpty()) {
                             for (String bookieAddress : bookieGroup.keySet()) {
-                                excludedBookies.remove(BookieId.parse(bookieAddress));
-                                totalAvailableBookiesFromPrimaryAndSecondary += 1;
+                                if (excludedBookies.remove(BookieId.parse(bookieAddress))) {
+                                    totalAvailableBookiesFromPrimaryAndSecondary += 1;
+                                }
                             }
                         }
                     }
diff --git a/pulsar-broker-common/src/test/java/org/apache/pulsar/bookie/rackawareness/IsolatedBookieEnsemblePlacementPolicyTest.java b/pulsar-broker-common/src/test/java/org/apache/pulsar/bookie/rackawareness/IsolatedBookieEnsemblePlacementPolicyTest.java
@@ -699,4 +699,123 @@ public void testDefaultIsolationPolicyNotCovered() throws Exception {
                 .newEnsemble(2, 2, 2, Collections.emptyMap(), new HashSet<>()).getResult();
         assertEquals(BookieIdGroup1.containsAll(defaultBookieList),true);
     }
+
+    @Test
+    public void testGetExcludedBookiesWithIsolationGroups() throws Exception {
+        Map<String, Map<String, BookieInfo>> bookieMapping = new HashMap<>();
+        final String isolationGroup1 = "Group1";
+        final String isolationGroup2 = "Group2";
+        final String isolationGroup3 = "Group3";
+
+        Map<String, BookieInfo> group1 = new HashMap<>();
+        group1.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+        group1.put(BOOKIE2, BookieInfo.builder().rack("rack0").build());
+
+        Map<String, BookieInfo> group2 = new HashMap<>();
+        group2.put(BOOKIE3, BookieInfo.builder().rack("rack1").build());
+        group2.put(BOOKIE4, BookieInfo.builder().rack("rack1").build());
+
+        bookieMapping.put(isolationGroup1, group1);
+        bookieMapping.put(isolationGroup2, group2);
+
+        store.put(BookieRackAffinityMapping.BOOKIE_INFO_ROOT_PATH, jsonMapper.writeValueAsBytes(bookieMapping),
+                Optional.empty()).join();
+
+        IsolatedBookieEnsemblePlacementPolicy isolationPolicy = new IsolatedBookieEnsemblePlacementPolicy();
+        ClientConfiguration bkClientConf = new ClientConfiguration();
+        bkClientConf.setProperty(BookieRackAffinityMapping.METADATA_STORE_INSTANCE, store);
+        bkClientConf.setProperty(IsolatedBookieEnsemblePlacementPolicy.ISOLATION_BOOKIE_GROUPS, isolationGroup1);
+        bkClientConf.setProperty(IsolatedBookieEnsemblePlacementPolicy.SECONDARY_ISOLATION_BOOKIE_GROUPS, isolationGroup2);
+        isolationPolicy.initialize(bkClientConf, Optional.empty(), timer, SettableFeatureProvider.DISABLE_ALL,
+                NullStatsLogger.INSTANCE, BookieSocketAddress.LEGACY_BOOKIEID_RESOLVER);
+        isolationPolicy.onClusterChanged(writableBookies, readOnlyBookies);
+
+        /* Test common cases */
+        MutablePair<Set<String>, Set<String>> groups = new MutablePair<>();
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(""));
+        Set<BookieId> blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertEquals(blacklist.size(), 2);
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertEquals(blacklist.size(), 2);
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(3, groups);
+        assertTrue(blacklist.isEmpty());
+
+        /* Test a bookie belongs to multiple isolation groups */
+        group1 = new HashMap<>();
+        group1.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+
+        group2 = new HashMap<>();
+        group2.put(BOOKIE2, BookieInfo.builder().rack("rack0").build());
+        group2.put(BOOKIE3, BookieInfo.builder().rack("rack1").build());
+        group2.put(BOOKIE4, BookieInfo.builder().rack("rack1").build());
+
+        Map<String, BookieInfo> group3 = new HashMap<>();
+        group3.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+
+        bookieMapping.put(isolationGroup1, group1);
+        bookieMapping.put(isolationGroup2, group2);
+        bookieMapping.put(isolationGroup3, group3);
+
+        store.put(BookieRackAffinityMapping.BOOKIE_INFO_ROOT_PATH, jsonMapper.writeValueAsBytes(bookieMapping),
+                Optional.empty()).join();
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1, isolationGroup3));
+        groups.setRight(Sets.newHashSet(""));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertEquals(blacklist.size(), 3);
+
+        /* Test a bookie belongs to multiple isolation groups and totalAvailableBookiesInPrimaryGroup < ensembleSize */
+        groups.setLeft(Sets.newHashSet(isolationGroup1, isolationGroup3));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertTrue(blacklist.isEmpty());
+
+        /* Test some bookies not set rack config */
+        group1 = new HashMap<>();
+        group1.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+
+        group2 = new HashMap<>();
+        group2.put(BOOKIE2, BookieInfo.builder().rack("rack0").build());
+
+        bookieMapping.put(isolationGroup1, group1);
+        bookieMapping.put(isolationGroup2, group2);
+
+        store.put(BookieRackAffinityMapping.BOOKIE_INFO_ROOT_PATH, jsonMapper.writeValueAsBytes(bookieMapping),
+                Optional.empty()).join();
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertEquals(blacklist.size(), 2);
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(3, groups);
+        assertTrue(blacklist.isEmpty());
+
+        /* Test some bookies not set rack config and totalAvailableBookiesFromPrimaryAndSecondary < ensembleSize */
+        group1 = new HashMap<>();
+        group1.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+
+        group2 = new HashMap<>();
+        group2.put(BOOKIE1, BookieInfo.builder().rack("rack0").build());
+
+        bookieMapping.put(isolationGroup1, group1);
+        bookieMapping.put(isolationGroup2, group2);
+
+        store.put(BookieRackAffinityMapping.BOOKIE_INFO_ROOT_PATH, jsonMapper.writeValueAsBytes(bookieMapping),
+                Optional.empty()).join();
+
+        groups.setLeft(Sets.newHashSet(isolationGroup1));
+        groups.setRight(Sets.newHashSet(isolationGroup2));
+        blacklist = isolationPolicy.getExcludedBookiesWithIsolationGroups(2, groups);
+        assertTrue(blacklist.isEmpty());
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -206,6 +206,9 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,`
`206`	`206`	`return excludedBookies;`
`207`	`207`	`}`
`208`	`208`	`Set<String> allGroups = allGroupsBookieMapping.keySet();`
	`209`	`+ if (allGroups.isEmpty()) {`
	`210`	`+ return excludedBookies;`
	`211`	`+ }`
`209`	`212`	`int totalAvailableBookiesInPrimaryGroup = 0;`
`210`	`213`	`Set<String> primaryIsolationGroup = Collections.emptySet();`
`211`	`214`	`Set<String> secondaryIsolationGroup = Collections.emptySet();`
`@@ -222,9 +225,10 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,`
`222`	`225`	`}`
`223`	`226`	`} else {`
`224`	`227`	`for (String groupBookie : bookiesInGroup) {`
`225`		`- totalAvailableBookiesInPrimaryGroup += knownBookies`
`226`		`- .containsKey(BookieId.parse(groupBookie)) ? 1 : 0;`
`227`		`- primaryGroupBookies.add(BookieId.parse(groupBookie));`
	`228`	`+ BookieId bookieId = BookieId.parse(groupBookie);`
	`229`	`+ if (primaryGroupBookies.add(bookieId)) {`
	`230`	`+ totalAvailableBookiesInPrimaryGroup += knownBookies.containsKey(bookieId) ? 1 : 0;`
	`231`	`+ }`
`228`	`232`	`}`
`229`	`233`	`}`
`230`	`234`	`}`
`@@ -256,8 +260,9 @@ Set<BookieId> getExcludedBookiesWithIsolationGroups(int ensembleSize,`
`256`	`260`	`Map<String, BookieInfo> bookieGroup = allGroupsBookieMapping.get(group);`
`257`	`261`	`if (bookieGroup != null && !bookieGroup.isEmpty()) {`
`258`	`262`	`for (String bookieAddress : bookieGroup.keySet()) {`
`259`		`- excludedBookies.remove(BookieId.parse(bookieAddress));`
`260`		`- totalAvailableBookiesFromPrimaryAndSecondary += 1;`
	`263`	`+ if (excludedBookies.remove(BookieId.parse(bookieAddress))) {`
	`264`	`+ totalAvailableBookiesFromPrimaryAndSecondary += 1;`
	`265`	`+ }`
`261`	`266`	`}`
`262`	`267`	`}`
`263`	`268`	`}`