diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java index bdbbe726c5c6d..21bcc366c135f 100644 --- a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java +++ b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java @@ -190,6 +190,13 @@ public class BalancedShardsAllocator implements ShardsAllocator { Property.NodeScope ); + public static final Setting PREFER_FILTER_AWARE_BALANCE = Setting.boolSetting( + "cluster.routing.allocation.balance.filter_aware", + false, + Property.Dynamic, + Property.NodeScope + ); + public static final Setting ALLOCATOR_TIMEOUT_SETTING = Setting.timeSetting( "cluster.routing.allocation.balanced_shards_allocator.allocator_timeout", TimeValue.MINUS_ONE, @@ -238,6 +245,7 @@ private static Priority parseReroutePriority(String priorityString) { private volatile boolean preferPrimaryShardBalance; private volatile boolean preferPrimaryShardRebalance; + private volatile boolean preferFilterAwareBalance; private volatile float preferPrimaryShardRebalanceBuffer; private volatile float indexBalanceFactor; private volatile float shardBalanceFactor; @@ -266,6 +274,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting setPrimaryConstraintThresholdSetting(PRIMARY_CONSTRAINT_THRESHOLD_SETTING.get(settings)); setPreferPrimaryShardBalance(PREFER_PRIMARY_SHARD_BALANCE.get(settings)); setPreferPrimaryShardRebalance(PREFER_PRIMARY_SHARD_REBALANCE.get(settings)); + setPreferFilterAwareBalance(PREFER_FILTER_AWARE_BALANCE.get(settings)); setShardMovementStrategy(SHARD_MOVEMENT_STRATEGY_SETTING.get(settings)); setAllocatorTimeout(ALLOCATOR_TIMEOUT_SETTING.get(settings)); setFollowUpRerouteTaskPriority(FOLLOW_UP_REROUTE_PRIORITY_SETTING.get(settings)); @@ -276,6 +285,7 @@ public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSetting clusterSettings.addSettingsUpdateConsumer(SHARD_BALANCE_FACTOR_SETTING, this::updateShardBalanceFactor); clusterSettings.addSettingsUpdateConsumer(PRIMARY_SHARD_REBALANCE_BUFFER, this::updatePreferPrimaryShardBalanceBuffer); clusterSettings.addSettingsUpdateConsumer(PREFER_PRIMARY_SHARD_REBALANCE, this::setPreferPrimaryShardRebalance); + clusterSettings.addSettingsUpdateConsumer(PREFER_FILTER_AWARE_BALANCE, this::setPreferFilterAwareBalance); clusterSettings.addSettingsUpdateConsumer(THRESHOLD_SETTING, this::setThreshold); clusterSettings.addSettingsUpdateConsumer(PRIMARY_CONSTRAINT_THRESHOLD_SETTING, this::setPrimaryConstraintThresholdSetting); clusterSettings.addSettingsUpdateConsumer(IGNORE_THROTTLE_FOR_REMOTE_RESTORE, this::setIgnoreThrottleInRestore); @@ -368,6 +378,10 @@ private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance) this.weightFunction.updateRebalanceConstraint(CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID, preferPrimaryShardRebalance); } + private void setPreferFilterAwareBalance(boolean preferFilterAwareBalance) { + this.preferFilterAwareBalance = preferFilterAwareBalance; + } + private void setThreshold(float threshold) { this.threshold = threshold; } @@ -409,6 +423,7 @@ public void allocate(RoutingAllocation allocation) { threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, + preferFilterAwareBalance, ignoreThrottleInRestore, this::allocatorTimedOut ); @@ -435,6 +450,7 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, + preferFilterAwareBalance, ignoreThrottleInRestore, () -> false // as we don't need to check if timed out or not while just understanding ShardAllocationDecision ); @@ -795,7 +811,7 @@ public Balancer( float threshold, boolean preferPrimaryBalance ) { - super(logger, allocation, shardMovementStrategy, weight, threshold, preferPrimaryBalance, false, false, () -> false); + super(logger, allocation, shardMovementStrategy, weight, threshold, preferPrimaryBalance, false, false, false, () -> false); } } diff --git a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java index 2fcf51f102fd4..b443044bce889 100644 --- a/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java +++ b/server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java @@ -72,6 +72,7 @@ public class LocalShardsBalancer extends ShardsBalancer { private final Metadata metadata; private final float avgPrimaryShardsPerNode; + private final int balanceNodeCount; private final BalancedShardsAllocator.NodeSorter sorter; private final Set inEligibleTargetNode; private final Supplier timedOutFunc; @@ -85,6 +86,7 @@ public LocalShardsBalancer( float threshold, boolean preferPrimaryBalance, boolean preferPrimaryRebalance, + boolean preferFilterAwareBalance, boolean ignoreThrottleInRestore, Supplier timedOutFunc ) { @@ -94,9 +96,21 @@ public LocalShardsBalancer( this.threshold = threshold; this.routingNodes = allocation.routingNodes(); this.metadata = allocation.metadata(); - avgPrimaryShardsPerNode = (float) (StreamSupport.stream(metadata.spliterator(), false) - .mapToInt(IndexMetadata::getNumberOfShards) - .sum()) / routingNodes.size(); + int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); + int balanceNodeCount = routingNodes.size(); + if (preferFilterAwareBalance) { + int eligibleNodeCount = 0; + for (RoutingNode routingNode : routingNodes) { + if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { + eligibleNodeCount++; + } + } + if (eligibleNodeCount > 0) { + balanceNodeCount = eligibleNodeCount; + } + } + this.balanceNodeCount = balanceNodeCount; + avgPrimaryShardsPerNode = ((float) primarySum) / balanceNodeCount; nodes = Collections.unmodifiableMap(buildModelFromAssigned()); sorter = newNodeSorter(); inEligibleTargetNode = new HashSet<>(); @@ -119,12 +133,12 @@ private BalancedShardsAllocator.ModelNode[] nodesArray() { */ @Override public float avgShardsPerNode(String index) { - return ((float) metadata.index(index).getTotalNumberOfShards()) / nodes.size(); + return ((float) metadata.index(index).getTotalNumberOfShards()) / balanceNodeCount; } @Override public float avgPrimaryShardsPerNode(String index) { - return ((float) metadata.index(index).getNumberOfShards()) / nodes.size(); + return ((float) metadata.index(index).getNumberOfShards()) / balanceNodeCount; } @Override @@ -137,7 +151,7 @@ public float avgPrimaryShardsPerNode() { */ @Override public float avgShardsPerNode() { - return totalShardCount / nodes.size(); + return ((float) totalShardCount) / balanceNodeCount; } /** diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index 77108dedb4245..617a02ceecc41 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -286,6 +286,7 @@ public void apply(Settings value, Settings current, Settings previous) { BalancedShardsAllocator.PRIMARY_SHARD_REBALANCE_BUFFER, BalancedShardsAllocator.PREFER_PRIMARY_SHARD_BALANCE, BalancedShardsAllocator.PREFER_PRIMARY_SHARD_REBALANCE, + BalancedShardsAllocator.PREFER_FILTER_AWARE_BALANCE, BalancedShardsAllocator.SHARD_MOVE_PRIMARY_FIRST_SETTING, BalancedShardsAllocator.SHARD_MOVEMENT_STRATEGY_SETTING, BalancedShardsAllocator.THRESHOLD_SETTING, diff --git a/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java b/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java index 93b9c5ea0e94f..18aa8a606dac2 100644 --- a/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java +++ b/server/src/test/java/org/opensearch/cluster/routing/allocation/BalanceConfigurationTests.java @@ -1112,4 +1112,182 @@ public ShardAllocationDecision decideShardAllocation(ShardRouting shard, Routing } } } + + /** + * Cluster: 8 nodes; node4..node7 excluded via cluster.routing.allocation.exclude._id + * Index : 12 primaries x (1 + 1 replica) = 24 shards, all on the 4 eligible nodes + * every eligible node holds exactly 6 shards. + * Shard layout (each column = one node, each row = one shard on that node): + * + * eligible (allowed) | excluded (filter NO) + * node0 node1 node2 node3 | node4 node5 node6 node7 + * ----- ----- ----- ----- | ----- ----- ----- ----- + * P P P P | . . . . + * P P P P | . . . . + * P r r r | . . . . + * P r r r | . . . . + * P r r r | . . . . + * P r r r | . . . . + * 6P/0r 2P/4r 2P/4r 2P/4r | (P = primary, r = replica) + * + * Primary layout on eligible nodes: 6 / 2 / 2 / 2 (ideal = 3 / 3 / 3 / 3) + */ + public void testPrimaryRebalanceIgnoresAllocationFilter() { + final int numberOfNodes = 8; + final int excludedNodeCount = 4; + final int numberOfShards = 12; + final int numberOfReplicas = 1; + final String indexName = "test"; + + // Enable prefer-primary balance & rebalance, exclude the second half of nodes. + Settings.Builder settingsBuilder = getSettingsBuilderForPrimaryReBalance(); + StringBuilder excludeList = new StringBuilder(); + for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) { + if (excludeList.length() > 0) { + excludeList.append(","); + } + excludeList.append("node").append(i); + } + settingsBuilder.put("cluster.routing.allocation.exclude._id", excludeList.toString()); + settingsBuilder.put("cluster.routing.allocation.balance.filter_aware", true); + + AllocationService strategy = createAllocationService(settingsBuilder.build(), new TestGatewayAllocator()); + + // Build 8 discovery nodes. + DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder(); + List nodesList = new ArrayList<>(numberOfNodes); + for (int i = 0; i < numberOfNodes; i++) { + DiscoveryNode node = newNode("node" + i); + discoBuilder.add(node); + nodesList.add(node.getId()); + } + discoBuilder.localNodeId(nodesList.get(0)); + discoBuilder.clusterManagerNodeId(nodesList.get(0)); + + IndexMetadata indexMetadata = getIndexMetadata(indexName, numberOfShards, numberOfReplicas); + + int[] primaryOwners = new int[] { 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 3, 3 }; + int[] replicaOwners = new int[] { 1, 2, 3, 1, 2, 3, 2, 3, 1, 3, 1, 2 }; + + IndexRoutingTable.Builder indexRoutingTable = IndexRoutingTable.builder(indexMetadata.getIndex()); + IndexMetadata.Builder indexMetaDataBuilder = IndexMetadata.builder(indexMetadata); + + for (int shardId = 0; shardId < numberOfShards; shardId++) { + ShardId sid = new ShardId(indexMetadata.getIndex(), shardId); + IndexShardRoutingTable.Builder shardBuilder = new IndexShardRoutingTable.Builder(sid); + shardBuilder.addShard( + TestShardRouting.newShardRouting(sid, nodesList.get(primaryOwners[shardId]), true, ShardRoutingState.STARTED) + ); + shardBuilder.addShard( + TestShardRouting.newShardRouting(sid, nodesList.get(replicaOwners[shardId]), false, ShardRoutingState.STARTED) + ); + IndexShardRoutingTable shardTable = shardBuilder.build(); + indexRoutingTable.addIndexShard(shardTable); + indexMetaDataBuilder.putInSyncAllocationIds(shardId, shardTable.getAllAllocationIds()); + } + + Metadata.Builder metadata = Metadata.builder(); + metadata.persistentSettings(settingsBuilder.build()); + metadata.put(indexMetaDataBuilder.build(), false); + + RoutingTable.Builder routingTable = RoutingTable.builder(); + routingTable.add(indexRoutingTable); + + ClusterState.Builder stateBuilder = ClusterState.builder(new ClusterName("test")); + stateBuilder.nodes(discoBuilder); + stateBuilder.metadata(metadata.generateClusterUuidIfNeeded().build()); + stateBuilder.routingTable(routingTable.build()); + ClusterState clusterState = stateBuilder.build(); + + // Sanity check: initial primary distribution is 6/2/2/2 on eligible nodes and 0 on excluded. + int[] initialPrimaries = countPrimariesPerNode(clusterState, nodesList, indexName); + int[] initialTotals = countTotalShardsPerNode(clusterState, nodesList, indexName); + logger.info("[before reroute]\n{}", ShardAllocations.printShardDistribution(clusterState)); + assertEquals("initial primaries on node0", 6, initialPrimaries[0]); + assertEquals("initial primaries on node1", 2, initialPrimaries[1]); + assertEquals("initial primaries on node2", 2, initialPrimaries[2]); + assertEquals("initial primaries on node3", 2, initialPrimaries[3]); + for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) { + assertEquals("initial primary count on excluded " + nodesList.get(i), 0, initialPrimaries[i]); + } + // Total shard count on every eligible node is identical so the generic shard-balance term is + // neutral -- any relocation observed later must come from the primary-only rebalance logic. + assertEquals("eligible node0 total shards", 6, initialTotals[0]); + assertEquals("eligible node1 total shards", 6, initialTotals[1]); + assertEquals("eligible node2 total shards", 6, initialTotals[2]); + assertEquals("eligible node3 total shards", 6, initialTotals[3]); + + clusterState = strategy.reroute(clusterState, "reroute"); + clusterState = applyStartedShardsUntilNoChange(clusterState, strategy); + for (int i = 0; i < 20; i++) { + ClusterState next = strategy.reroute(clusterState, "reroute-loop-" + i); + next = applyStartedShardsUntilNoChange(next, strategy); + if (next.equals(clusterState)) { + break; + } + clusterState = next; + } + + logger.info("[after reroute]\n{}", ShardAllocations.printShardDistribution(clusterState)); + + int[] finalPrimaries = countPrimariesPerNode(clusterState, nodesList, indexName); + + // Excluded nodes must still be empty. + for (int i = numberOfNodes - excludedNodeCount; i < numberOfNodes; i++) { + assertEquals("excluded node " + nodesList.get(i) + " must remain empty", 0, finalPrimaries[i]); + } + + int totalEligiblePrimaries = 0; + int maxEligible = Integer.MIN_VALUE; + int minEligible = Integer.MAX_VALUE; + for (int i = 0; i < numberOfNodes - excludedNodeCount; i++) { + totalEligiblePrimaries += finalPrimaries[i]; + maxEligible = Math.max(maxEligible, finalPrimaries[i]); + minEligible = Math.min(minEligible, finalPrimaries[i]); + } + assertEquals("all primaries should stay on eligible nodes", numberOfShards, totalEligiblePrimaries); + assertEquals( + "Primaries fully balanced across eligible nodes -- final distribution: " + + "node0=" + + finalPrimaries[0] + + ", node1=" + + finalPrimaries[1] + + ", node2=" + + finalPrimaries[2] + + ", node3=" + + finalPrimaries[3], + 0, + maxEligible - minEligible + ); + } + + private static int[] countPrimariesPerNode(ClusterState state, List nodesList, String indexName) { + int[] counts = new int[nodesList.size()]; + RoutingNodes routingNodes = state.getRoutingNodes(); + for (int i = 0; i < nodesList.size(); i++) { + RoutingNode node = routingNodes.node(nodesList.get(i)); + if (node == null) { + continue; + } + counts[i] = (int) node.shardsWithState(indexName, STARTED).stream().filter(ShardRouting::primary).count(); + } + return counts; + } + + /** + * Helper: returns total shard counts (primary + replica, only STARTED) for the given index on each node. + */ + private static int[] countTotalShardsPerNode(ClusterState state, List nodesList, String indexName) { + int[] counts = new int[nodesList.size()]; + RoutingNodes routingNodes = state.getRoutingNodes(); + for (int i = 0; i < nodesList.size(); i++) { + RoutingNode node = routingNodes.node(nodesList.get(i)); + if (node == null) { + continue; + } + counts[i] = node.shardsWithState(indexName, STARTED).size(); + } + return counts; + } + } diff --git a/server/src/test/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancerTests.java b/server/src/test/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancerTests.java index 14088efae8cbe..a3b1747c76172 100644 --- a/server/src/test/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancerTests.java +++ b/server/src/test/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancerTests.java @@ -22,6 +22,7 @@ import org.opensearch.cluster.routing.RoutingTable; import org.opensearch.cluster.routing.ShardRouting; import org.opensearch.cluster.routing.ShardRoutingState; +import org.opensearch.cluster.routing.allocation.AllocationService; import org.opensearch.cluster.routing.allocation.RoutingAllocation; import org.opensearch.cluster.routing.allocation.decider.AllocationDecider; import org.opensearch.cluster.routing.allocation.decider.AllocationDeciders; @@ -32,8 +33,10 @@ import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Function; import static org.mockito.ArgumentMatchers.any; @@ -84,6 +87,7 @@ public void testAllocateUnassignedWhenAllShardsCanBeAllocated() { false, false, false, + false, null ); @@ -131,6 +135,7 @@ public void testAllocateUnassignedWhenSearchShardsCannotBeAllocated() { false, false, false, + false, null ); @@ -178,6 +183,7 @@ public void testAllocateUnassignedWhenRegularReplicaShardsCannotBeAllocated() { false, false, false, + false, null ); @@ -258,4 +264,117 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat return decider.apply(shardRouting); } } + + public void testAvgShardsPerNodeRefreshesWhenFilterChanges() { + int numberOfShards = 6; + Metadata metadata = buildMetadata(Metadata.builder(), 1, numberOfShards, 1, 0); + RoutingTable routingTable = buildRoutingTable(metadata); + ClusterState state = ClusterState.builder(ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY)) + .metadata(metadata) + .routingTable(routingTable) + .nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3).add(node4).add(node5).add(node6)) + .build(); + + // Drive all shards to STARTED so totalShardCount reflects the assigned shard count (12). + AllocationService allocationService = createAllocationService(Settings.EMPTY); + state = allocationService.reroute(state, "initial-allocation"); + state = applyStartedShardsUntilNoChange(state, allocationService); + assertEquals(12, state.getRoutingNodes().shardsWithState(ShardRoutingState.STARTED).size()); + + // ---- Round 1: exclude {node4,node5,node6} -> 3 eligible -> 6 / 3 = 2.0, total 12 / 3 = 4.0 + RoutingAllocation allocation1 = new RoutingAllocation( + new AllocationDeciders( + Collections.singletonList(new StaticClusterFilterDecider(new HashSet<>(Arrays.asList("node4", "node5", "node6")))) + ), + new RoutingNodes(state, false), + state, + ClusterInfo.EMPTY, + null, + System.nanoTime() + ); + LocalShardsBalancer balancer1 = new LocalShardsBalancer( + logger, + allocation1, + null, + mock(BalancedShardsAllocator.WeightFunction.class), + 0, + false, + false, + true, + false, + null + ); + assertEquals("round1 primary per-index", 2.0f, balancer1.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round1 primary cluster-level", 2.0f, balancer1.avgPrimaryShardsPerNode(), 0.0001f); + assertEquals("round1 shards per-index", 12.0f / 3.0f, balancer1.avgShardsPerNode("test_0"), 0.0001f); + assertEquals("round1 shards cluster-level", 12.0f / 3.0f, balancer1.avgShardsPerNode(), 0.0001f); + + // ---- Round 2: filter shrinks to {node6} -> 5 eligible -> 6 / 5 = 1.2 + RoutingAllocation allocation2 = new RoutingAllocation( + new AllocationDeciders( + Collections.singletonList(new StaticClusterFilterDecider(new HashSet<>(Collections.singletonList("node6")))) + ), + new RoutingNodes(state, false), + state, + ClusterInfo.EMPTY, + null, + System.nanoTime() + ); + LocalShardsBalancer balancer2 = new LocalShardsBalancer( + logger, + allocation2, + null, + mock(BalancedShardsAllocator.WeightFunction.class), + 0, + false, + false, + true, + false, + null + ); + assertEquals("round2 primary per-index", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round2 primary cluster-level", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode(), 0.0001f); + assertEquals("round2 shards per-index", 12.0f / 5.0f, balancer2.avgShardsPerNode("test_0"), 0.0001f); + assertEquals("round2 shards cluster-level", 12.0f / 5.0f, balancer2.avgShardsPerNode(), 0.0001f); + + // ---- Round 3: filter cleared -> 6 eligible -> 6 / 6 = 1.0 + RoutingAllocation allocation3 = new RoutingAllocation( + new AllocationDeciders(Collections.singletonList(new StaticClusterFilterDecider(Collections.emptySet()))), + new RoutingNodes(state, false), + state, + ClusterInfo.EMPTY, + null, + System.nanoTime() + ); + LocalShardsBalancer balancer3 = new LocalShardsBalancer( + logger, + allocation3, + null, + mock(BalancedShardsAllocator.WeightFunction.class), + 0, + false, + false, + true, + false, + null + ); + assertEquals("round3 primary per-index", 1.0f, balancer3.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round3 primary cluster-level", 1.0f, balancer3.avgPrimaryShardsPerNode(), 0.0001f); + assertEquals("round3 shards per-index", 2.0f, balancer3.avgShardsPerNode("test_0"), 0.0001f); + assertEquals("round3 shards cluster-level", 2.0f, balancer3.avgShardsPerNode(), 0.0001f); + } + + public static class StaticClusterFilterDecider extends AllocationDecider { + private final Set excluded; + + public StaticClusterFilterDecider(Set excluded) { + this.excluded = excluded; + } + + @Override + public Decision canAllocateAnyShardToNode(RoutingNode node, RoutingAllocation allocation) { + return excluded.contains(node.nodeId()) ? Decision.NO : Decision.ALWAYS; + } + } + }