Refactor addDefaultRackBookiesIfMinNumRacksIsEnforced to Remove Redundant Logic

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

@@ -359,30 +359,22 @@ public void uninitalize() {
      */
     protected Set<BookieId> addDefaultRackBookiesIfMinNumRacksIsEnforced(
             Set<BookieId> excludeBookies) {
-        Set<BookieId> comprehensiveExclusionBookiesSet;
-        if (enforceMinNumRacksPerWriteQuorum) {
-            Set<BookieId> bookiesInDefaultRack = null;
-            Set<Node> defaultRackLeaves = topology.getLeaves(getDefaultRack());
-            for (Node node : defaultRackLeaves) {
-                if (node instanceof BookieNode) {
-                    if (bookiesInDefaultRack == null) {
-                        bookiesInDefaultRack = new HashSet<BookieId>(excludeBookies);
-                    }
-                    bookiesInDefaultRack.add(((BookieNode) node).getAddr());
-                } else {
-                    LOG.error("found non-BookieNode: {} as leaf of defaultrack: {}", node, getDefaultRack());
-                }
-            }
-            if ((bookiesInDefaultRack == null) || bookiesInDefaultRack.isEmpty()) {
-                comprehensiveExclusionBookiesSet = excludeBookies;
+        if (!enforceMinNumRacksPerWriteQuorum) {
+            return excludeBookies;
+        }
+
+        Set<BookieId> comprehensiveExclusionBookiesSet = new HashSet<>(excludeBookies);
+        Set<Node> defaultRackLeaves = topology.getLeaves(getDefaultRack());
+        for (Node node : defaultRackLeaves) {
+            if (node instanceof BookieNode) {
+                comprehensiveExclusionBookiesSet.add(((BookieNode) node).getAddr());
             } else {
-                comprehensiveExclusionBookiesSet = new HashSet<BookieId>(excludeBookies);
-                comprehensiveExclusionBookiesSet.addAll(bookiesInDefaultRack);
-                LOG.info("enforceMinNumRacksPerWriteQuorum is enabled, so Excluding bookies of defaultRack: {}",
-                        bookiesInDefaultRack);
+                LOG.error("found non-BookieNode: {} as leaf of defaultrack: {}", node, getDefaultRack());
             }
-        } else {
-            comprehensiveExclusionBookiesSet = excludeBookies;
+        }
+        if (!comprehensiveExclusionBookiesSet.isEmpty()) {
+            LOG.info("enforceMinNumRacksPerWriteQuorum is enabled, so Excluding bookies of defaultRack: {}",
+                    comprehensiveExclusionBookiesSet);
         }
         return comprehensiveExclusionBookiesSet;
     }

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

@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.bookkeeper.client;
+
+import static org.apache.bookkeeper.client.TopologyAwareEnsemblePlacementPolicy.REPP_DNS_RESOLVER_CLASS;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Optional;
+import java.util.Set;
+import org.apache.bookkeeper.conf.ClientConfiguration;
+import org.apache.bookkeeper.feature.FeatureProvider;
+import org.apache.bookkeeper.net.BookieId;
+import org.apache.bookkeeper.net.BookieNode;
+import org.apache.bookkeeper.net.NetworkTopology;
+import org.apache.bookkeeper.net.Node;
+import org.apache.bookkeeper.net.ScriptBasedMapping;
+import org.apache.bookkeeper.proto.BookieAddressResolver;
+import org.apache.bookkeeper.test.TestStatsProvider;
+import org.apache.bookkeeper.util.StaticDNSResolver;
+import org.junit.Before;
+import org.junit.Test;
+
+public class RackawareEnsemblePlacementPolicyImplTest {
+
+    private RackawareEnsemblePlacementPolicyImpl policy;
+    private NetworkTopology topology;
+    private Set<BookieId> excludeBookies;
+    private BookieId bookie1;
+    private BookieId bookie2;
+
+    String defaultRack = "/default-rack";
+
+    @Before
+    public void setUp() {
+        TestStatsProvider provider = new TestStatsProvider();
+        policy = new RackawareEnsemblePlacementPolicyImpl();
+        ClientConfiguration clientConfiguration = mock(ClientConfiguration.class);
+        when(clientConfiguration.getEnforceMinNumRacksPerWriteQuorum()).thenReturn(true);
+        when(clientConfiguration.getString(REPP_DNS_RESOLVER_CLASS,
+                ScriptBasedMapping.class.getName())).thenReturn(StaticDNSResolver.class.getName());
+        policy.initialize(clientConfiguration,
+                Optional.empty(),
+                mock(io.netty.util.HashedWheelTimer.class),
+                mock(FeatureProvider.class),
+                provider.getStatsLogger(""),
+                mock(BookieAddressResolver.class));
+        topology = mock(NetworkTopology.class);
+        excludeBookies = new HashSet<>();
+        bookie1 = BookieId.parse("bookie1:3181");
+        bookie2 = BookieId.parse("bookie2:3181");
+
+        // Set the default rack
+        when(topology.getLeaves(defaultRack))
+                .thenReturn(Collections.emptySet());
+
+        // Use reflection to inject the mock topology into the policy
+        TestUtils.setField(TopologyAwareEnsemblePlacementPolicy.class, policy, "topology", topology);
+    }
+
+    @Test
+    public void testWithDefaultRackBookies() {
+        // Scenario: bookie nodes exist in the default rack
+        BookieNode node1 = new BookieNode(bookie1, defaultRack);
+        BookieNode node2 = new BookieNode(bookie2, defaultRack);
+
+        Set<Node> defaultRackNodes = new HashSet<>();
+        defaultRackNodes.add(node1);
+        defaultRackNodes.add(node2);
+
+        when(topology.getLeaves(defaultRack))
+                .thenReturn(defaultRackNodes);
+
+        // Add some bookies to be excluded
+        excludeBookies.add(BookieId.parse("other:3181"));
+
+        Set<BookieId> result = policy.addDefaultRackBookiesIfMinNumRacksIsEnforced(excludeBookies);
+
+        // Verify: the result should contain the original exclusion set + all bookies from the default rack
+        assertEquals(3, result.size());
+        assertTrue(result.contains(bookie1));
+        assertTrue(result.contains(bookie2));
+        assertTrue(result.containsAll(excludeBookies));
+    }
+
+    @Test
+    public void testMixedNodesInDefaultRack() {
+        // Scenario: default rack contains mixed BookieNode and non-BookieNode nodes
+        BookieNode bookieNode = new BookieNode(bookie1, defaultRack);
+        Node nonBookieNode = mock(Node.class);
+
+        Set<Node> defaultRackNodes = new HashSet<>();
+        defaultRackNodes.add(bookieNode);
+        defaultRackNodes.add(nonBookieNode);
+
+        when(topology.getLeaves(defaultRack))
+                .thenReturn(defaultRackNodes);
+
+        Set<BookieId> result = policy.addDefaultRackBookiesIfMinNumRacksIsEnforced(excludeBookies);
+
+        // Verify: the result should include the BookieNode but not the non-BookieNode
+        assertEquals(1, result.size());
+        assertTrue(result.contains(bookie1));
+    }
+
+    @Test
+    public void testEmptyDefaultRackBookies() {
+        // Scenario: the collection in default rack is empty
+        Set<Node> emptySet = Collections.emptySet();
+        when(topology.getLeaves(defaultRack))
+                .thenReturn(emptySet);
+
+        // Add some bookies to be excluded
+        excludeBookies.add(BookieId.parse("other:3181"));
+        Set<BookieId> result = policy.addDefaultRackBookiesIfMinNumRacksIsEnforced(excludeBookies);
+
+        // Verify: should return the original exclusion set
+        assertSame(excludeBookies, result);
+    }
+}
+
+// Utility class for setting private fields
+class TestUtils {
+    public static void setField(Class<?> tclass , Object target, String fieldName, Object value) {
+        try {
+            java.lang.reflect.Field field = tclass.getDeclaredField(fieldName);
+            field.setAccessible(true);
+            field.set(target, value);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+}