diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java index be90ceab94d..d0574e6562a 100644 --- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java +++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImpl.java @@ -359,32 +359,31 @@ public void uninitalize() { */ protected Set addDefaultRackBookiesIfMinNumRacksIsEnforced( Set excludeBookies) { - Set comprehensiveExclusionBookiesSet; - if (enforceMinNumRacksPerWriteQuorum) { - Set bookiesInDefaultRack = null; - Set defaultRackLeaves = topology.getLeaves(getDefaultRack()); - for (Node node : defaultRackLeaves) { - if (node instanceof BookieNode) { - if (bookiesInDefaultRack == null) { - bookiesInDefaultRack = new HashSet(excludeBookies); - } - bookiesInDefaultRack.add(((BookieNode) node).getAddr()); - } else { - LOG.error("found non-BookieNode: {} as leaf of defaultrack: {}", node, getDefaultRack()); + if (!enforceMinNumRacksPerWriteQuorum) { + return excludeBookies; + } + + Set bookiesInDefaultRack = null; + Set defaultRackLeaves = topology.getLeaves(getDefaultRack()); + for (Node node : defaultRackLeaves) { + if (node instanceof BookieNode) { + if (bookiesInDefaultRack == null) { + bookiesInDefaultRack = new HashSet<>(); } - } - if ((bookiesInDefaultRack == null) || bookiesInDefaultRack.isEmpty()) { - comprehensiveExclusionBookiesSet = excludeBookies; + bookiesInDefaultRack.add(((BookieNode) node).getAddr()); } else { - comprehensiveExclusionBookiesSet = new HashSet(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()); } + } + if ((bookiesInDefaultRack == null) || bookiesInDefaultRack.isEmpty()) { + return excludeBookies; } else { - comprehensiveExclusionBookiesSet = excludeBookies; + Set comprehensiveExclusionBookiesSet = new HashSet<>(excludeBookies); + comprehensiveExclusionBookiesSet.addAll(bookiesInDefaultRack); + LOG.info("enforceMinNumRacksPerWriteQuorum is enabled, so Excluding bookies of defaultRack: {}", + bookiesInDefaultRack); + return comprehensiveExclusionBookiesSet; } - return comprehensiveExclusionBookiesSet; } @Override diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImplTest.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImplTest.java new file mode 100644 index 00000000000..8cbba944b7b --- /dev/null +++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/RackawareEnsemblePlacementPolicyImplTest.java @@ -0,0 +1,143 @@ +/* + * 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.apache.commons.lang3.reflect.FieldUtils; +import org.junit.Before; +import org.junit.Test; + +public class RackawareEnsemblePlacementPolicyImplTest { + + private RackawareEnsemblePlacementPolicyImpl policy; + private NetworkTopology topology; + private Set excludeBookies; + private BookieId bookie1; + private BookieId bookie2; + + String defaultRack = "/default-rack"; + + @Before + public void setUp() throws IllegalAccessException { + 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 + FieldUtils.writeField(policy, "topology", topology, true); + } + + @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 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 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 defaultRackNodes = new HashSet<>(); + defaultRackNodes.add(bookieNode); + defaultRackNodes.add(nonBookieNode); + + when(topology.getLeaves(defaultRack)) + .thenReturn(defaultRackNodes); + + Set 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 emptySet = Collections.emptySet(); + when(topology.getLeaves(defaultRack)) + .thenReturn(emptySet); + + // Add some bookies to be excluded + excludeBookies.add(BookieId.parse("other:3181")); + Set result = policy.addDefaultRackBookiesIfMinNumRacksIsEnforced(excludeBookies); + + // Verify: should return the original exclusion set + assertSame(excludeBookies, result); + } +} \ No newline at end of file