From c108ed027e05262d81da6e470f3eed80fd396e77 Mon Sep 17 00:00:00 2001 From: guojialiang Date: Mon, 24 Aug 2026 23:45:34 +0800 Subject: [PATCH 1/5] Make primary shard rebalance filter-aware Signed-off-by: guojialiang --- .../allocator/BalancedShardsAllocator.java | 25 +- .../allocator/LocalShardsBalancer.java | 18 +- .../common/settings/ClusterSettings.java | 1 + .../allocation/BalanceConfigurationTests.java | 229 ++++++++++++++++++ .../allocator/LocalShardsBalancerTests.java | 3 + 5 files changed, 272 insertions(+), 4 deletions(-) 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..3babec281d899 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 @@ -189,6 +189,20 @@ public class BalancedShardsAllocator implements ShardsAllocator { Property.Dynamic, Property.NodeScope ); + /** + * Dynamic switch controlling whether primary-shard rebalance weight + * calculation ({@link LocalShardsBalancer#avgPrimaryShardsPerNode()}) should + * respect {@link org.opensearch.cluster.routing.allocation.decider.FilterAllocationDecider} + * exclusions. When {@code true}, filter-excluded nodes are removed from the node + * count so the average is computed over eligible nodes only. Defaults to {@code false} + * to preserve legacy behaviour. + */ + public static final Setting PREFER_PRIMARY_FILTER_AWARE = Setting.boolSetting( + "cluster.routing.allocation.balance.prefer_primary.filter_aware", + false, + Property.Dynamic, + Property.NodeScope + ); public static final Setting ALLOCATOR_TIMEOUT_SETTING = Setting.timeSetting( "cluster.routing.allocation.balanced_shards_allocator.allocator_timeout", @@ -238,6 +252,7 @@ private static Priority parseReroutePriority(String priorityString) { private volatile boolean preferPrimaryShardBalance; private volatile boolean preferPrimaryShardRebalance; + private volatile boolean preferPrimaryFilterAware; private volatile float preferPrimaryShardRebalanceBuffer; private volatile float indexBalanceFactor; private volatile float shardBalanceFactor; @@ -266,6 +281,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)); + setPreferPrimaryFilterAware(PREFER_PRIMARY_FILTER_AWARE.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 +292,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_PRIMARY_FILTER_AWARE, this::setPreferPrimaryFilterAware); 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 +385,10 @@ private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance) this.weightFunction.updateRebalanceConstraint(CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID, preferPrimaryShardRebalance); } + private void setPreferPrimaryFilterAware(boolean preferPrimaryFilterAware) { + this.preferPrimaryFilterAware = preferPrimaryFilterAware; + } + private void setThreshold(float threshold) { this.threshold = threshold; } @@ -409,6 +430,7 @@ public void allocate(RoutingAllocation allocation) { threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, + preferPrimaryFilterAware, ignoreThrottleInRestore, this::allocatorTimedOut ); @@ -435,6 +457,7 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, + preferPrimaryFilterAware, ignoreThrottleInRestore, () -> false // as we don't need to check if timed out or not while just understanding ShardAllocationDecision ); @@ -795,7 +818,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..9e683fee6563b 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 @@ -85,6 +85,7 @@ public LocalShardsBalancer( float threshold, boolean preferPrimaryBalance, boolean preferPrimaryRebalance, + boolean preferPrimaryFilterAware, boolean ignoreThrottleInRestore, Supplier timedOutFunc ) { @@ -94,9 +95,20 @@ 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 nodeCount = routingNodes.size(); + if (preferPrimaryFilterAware) { + int eligible = 0; + for (RoutingNode rn : routingNodes) { + if (allocation.deciders().canAllocateAnyShardToNode(rn, allocation).type() != Decision.Type.NO) { + eligible++; + } + } + if (eligible > 0) { + nodeCount = eligible; + } + } + avgPrimaryShardsPerNode = ((float) primarySum) / nodeCount; nodes = Collections.unmodifiableMap(buildModelFromAssigned()); sorter = newNodeSorter(); inEligibleTargetNode = new HashSet<>(); 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..ede763de8e42d 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_PRIMARY_FILTER_AWARE, 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..417e191c85daf 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,233 @@ public ShardAllocationDecision decideShardAllocation(ShardRouting shard, Routing } } } + + /** + * Reproduces the bug where shard-rebalance weight calculation does not consider + * {@link org.opensearch.cluster.routing.allocation.decider.FilterAllocationDecider} exclusions. + *

+ * Setup (crafted so that regular total-shard balance is already satisfied and cannot mask + * the primary-only imbalance): + *

    + *
  • 8 nodes (node0..node7), the second half node4..node7 is excluded via + * {@code cluster.routing.allocation.exclude._id}.
    • + *
  • 1 index with 12 primary shards + 1 replica each = 24 shards total.
    • + *
  • Every eligible node holds exactly 6 shards (24 / 4 = 6), so the generic + * {@code shardBalance * (node.numShards - avg)} term is neutral: it does not drive any + * relocation on its own.
    • + *
  • Primary shard layout on the 4 eligible nodes is 6 / 2 / 2 / 2 — obviously not balanced + * (ideal would be 3 / 3 / 3 / 3).
    • + *
  • {@link BalancedShardsAllocator#PREFER_PRIMARY_SHARD_BALANCE} and + * {@link BalancedShardsAllocator#PREFER_PRIMARY_SHARD_REBALANCE} are both enabled so that + * the primary-only constraints participate in the weight function.
    • + *
+ *

+ * Root cause being reproduced: + *
+ * {@code LocalShardsBalancer#avgPrimaryShardsPerNode()} divides the total primary count by + * {@code nodes.size()} — which counts filter-excluded nodes too — giving + * {@code 12 / 8 = 1.5}. With the default {@code PRIMARY_SHARD_REBALANCE_BUFFER = 0.10} the + * threshold {@code allowedPrimaryShardCount = ceil(1.5 * 1.10) = 2}. Every eligible node + * except node0 already holds exactly 2 primaries, so + * {@link org.opensearch.cluster.routing.allocation.ConstraintTypes#isPrimaryShardsPerNodeBreached} + * marks each of node1/node2/node3 as "breached" (>= allowed) and penalises them as rebalance + * targets. Node0 (6 primaries) is over the threshold too and would like to shed primaries, but + * every candidate target is penalised the same way, so no primary can be relocated. + *

+ * Expected (ideal) primary distribution across the 4 eligible nodes: 3 / 3 / 3 / 3. + *
+ * Actual buggy distribution after reroute: 6 / 2 / 2 / 2. + *

+ * When the underlying bug is fixed (weight/avg calculation ignores filter-excluded nodes so + * that {@code avgPrimaryShardsPerNode = 12 / 4 = 3}) the final assertion at the bottom must be + * updated to require {@code max - min == 0}. + */ + 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()); + // Enable the new dynamic switch so the primary-rebalance weight calculation + // respects the FilterAllocationDecider exclusions. Default is false; turning it on + // is the fix path being validated by this test. + settingsBuilder.put("cluster.routing.allocation.balance.prefer_primary.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); + + // Craft an initial layout that is imbalanced on primaries but perfectly balanced on the + // total shard count of eligible nodes (6/6/6/6). This isolates the primary-rebalance path + // from the generic shard-balance path. + // + // Shard 0..5 : P -> node0 ; R -> node1/2/3/1/2/3 + // Shard 6,7 : P -> node1 ; R -> node2, node3 + // Shard 8,9 : P -> node2 ; R -> node1, node3 + // Shard 10,11 : P -> node3 ; R -> node1, node2 + // + // Resulting per-node counts on eligible nodes: + // node0: 6P + 0R = 6 + // node1: 2P + 4R = 6 + // node2: 2P + 4R = 6 + // node3: 2P + 4R = 6 + 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]); + + // Drive reroute repeatedly to allow rebalance to converge. + 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); + + // With the fix (dynamic switch above set to true), avgPrimaryShardsPerNode is computed + // over eligible nodes only: 12 / 4 = 3, so allowed = ceil(3 * 1.10) = 4. Eligible nodes + // with 2 primaries are no longer treated as "breached" rebalance targets, and node0 (6P) + // can migrate primaries to node1/2/3 until the layout converges to 3/3/3/3. + assertEquals("primary count on eligible node0 should converge to 3", 3, finalPrimaries[0]); + assertEquals("primary count on eligible node1 should converge to 3", 3, finalPrimaries[1]); + assertEquals("primary count on eligible node2 should converge to 3", 3, finalPrimaries[2]); + assertEquals("primary count on eligible node3 should converge to 3", 3, finalPrimaries[3]); + assertEquals( + "Fix applied: primaries fully balanced across eligible nodes -- final distribution: " + + "node0=" + + finalPrimaries[0] + + ", node1=" + + finalPrimaries[1] + + ", node2=" + + finalPrimaries[2] + + ", node3=" + + finalPrimaries[3], + 0, + maxEligible - minEligible + ); + } + + /** + * Helper: returns primary-shard counts for the given index on each node, indexed the same as {@code nodesList}. + */ + 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..e401d813dc078 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 @@ -84,6 +84,7 @@ public void testAllocateUnassignedWhenAllShardsCanBeAllocated() { false, false, false, + false, null ); @@ -131,6 +132,7 @@ public void testAllocateUnassignedWhenSearchShardsCannotBeAllocated() { false, false, false, + false, null ); @@ -178,6 +180,7 @@ public void testAllocateUnassignedWhenRegularReplicaShardsCannotBeAllocated() { false, false, false, + false, null ); From 6dbbc58dd7f559777d400206f52c57859020e1f1 Mon Sep 17 00:00:00 2001 From: guojialiang Date: Tue, 25 Aug 2026 14:56:02 +0800 Subject: [PATCH 2/5] refactor code Signed-off-by: guojialiang --- .../allocator/BalancedShardsAllocator.java | 9 +- .../allocator/LocalShardsBalancer.java | 4 +- .../allocation/BalanceConfigurationTests.java | 87 ++++--------------- 3 files changed, 21 insertions(+), 79 deletions(-) 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 3babec281d899..dfd1bde21dd23 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 @@ -189,14 +189,7 @@ public class BalancedShardsAllocator implements ShardsAllocator { Property.Dynamic, Property.NodeScope ); - /** - * Dynamic switch controlling whether primary-shard rebalance weight - * calculation ({@link LocalShardsBalancer#avgPrimaryShardsPerNode()}) should - * respect {@link org.opensearch.cluster.routing.allocation.decider.FilterAllocationDecider} - * exclusions. When {@code true}, filter-excluded nodes are removed from the node - * count so the average is computed over eligible nodes only. Defaults to {@code false} - * to preserve legacy behaviour. - */ + public static final Setting PREFER_PRIMARY_FILTER_AWARE = Setting.boolSetting( "cluster.routing.allocation.balance.prefer_primary.filter_aware", 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 9e683fee6563b..d257d4177be14 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 @@ -99,8 +99,8 @@ public LocalShardsBalancer( int nodeCount = routingNodes.size(); if (preferPrimaryFilterAware) { int eligible = 0; - for (RoutingNode rn : routingNodes) { - if (allocation.deciders().canAllocateAnyShardToNode(rn, allocation).type() != Decision.Type.NO) { + for (RoutingNode routingNode : routingNodes) { + if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { eligible++; } } 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 417e191c85daf..ef3a0af6351db 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 @@ -1114,44 +1114,23 @@ public ShardAllocationDecision decideShardAllocation(ShardRouting shard, Routing } /** - * Reproduces the bug where shard-rebalance weight calculation does not consider - * {@link org.opensearch.cluster.routing.allocation.decider.FilterAllocationDecider} exclusions. - *

- * Setup (crafted so that regular total-shard balance is already satisfied and cannot mask - * the primary-only imbalance): - *

    - *
  • 8 nodes (node0..node7), the second half node4..node7 is excluded via - * {@code cluster.routing.allocation.exclude._id}.
    • - *
  • 1 index with 12 primary shards + 1 replica each = 24 shards total.
    • - *
  • Every eligible node holds exactly 6 shards (24 / 4 = 6), so the generic - * {@code shardBalance * (node.numShards - avg)} term is neutral: it does not drive any - * relocation on its own.
    • - *
  • Primary shard layout on the 4 eligible nodes is 6 / 2 / 2 / 2 — obviously not balanced - * (ideal would be 3 / 3 / 3 / 3).
    • - *
  • {@link BalancedShardsAllocator#PREFER_PRIMARY_SHARD_BALANCE} and - * {@link BalancedShardsAllocator#PREFER_PRIMARY_SHARD_REBALANCE} are both enabled so that - * the primary-only constraints participate in the weight function.
    • - *
- *

- * Root cause being reproduced: - *
- * {@code LocalShardsBalancer#avgPrimaryShardsPerNode()} divides the total primary count by - * {@code nodes.size()} — which counts filter-excluded nodes too — giving - * {@code 12 / 8 = 1.5}. With the default {@code PRIMARY_SHARD_REBALANCE_BUFFER = 0.10} the - * threshold {@code allowedPrimaryShardCount = ceil(1.5 * 1.10) = 2}. Every eligible node - * except node0 already holds exactly 2 primaries, so - * {@link org.opensearch.cluster.routing.allocation.ConstraintTypes#isPrimaryShardsPerNodeBreached} - * marks each of node1/node2/node3 as "breached" (>= allowed) and penalises them as rebalance - * targets. Node0 (6 primaries) is over the threshold too and would like to shed primaries, but - * every candidate target is penalised the same way, so no primary can be relocated. - *

- * Expected (ideal) primary distribution across the 4 eligible nodes: 3 / 3 / 3 / 3. - *
- * Actual buggy distribution after reroute: 6 / 2 / 2 / 2. - *

- * When the underlying bug is fixed (weight/avg calculation ignores filter-excluded nodes so - * that {@code avgPrimaryShardsPerNode = 12 / 4 = 3}) the final assertion at the bottom must be - * updated to require {@code max - min == 0}. + * 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; @@ -1170,9 +1149,6 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { excludeList.append("node").append(i); } settingsBuilder.put("cluster.routing.allocation.exclude._id", excludeList.toString()); - // Enable the new dynamic switch so the primary-rebalance weight calculation - // respects the FilterAllocationDecider exclusions. Default is false; turning it on - // is the fix path being validated by this test. settingsBuilder.put("cluster.routing.allocation.balance.prefer_primary.filter_aware", true); AllocationService strategy = createAllocationService(settingsBuilder.build(), new TestGatewayAllocator()); @@ -1190,20 +1166,6 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { IndexMetadata indexMetadata = getIndexMetadata(indexName, numberOfShards, numberOfReplicas); - // Craft an initial layout that is imbalanced on primaries but perfectly balanced on the - // total shard count of eligible nodes (6/6/6/6). This isolates the primary-rebalance path - // from the generic shard-balance path. - // - // Shard 0..5 : P -> node0 ; R -> node1/2/3/1/2/3 - // Shard 6,7 : P -> node1 ; R -> node2, node3 - // Shard 8,9 : P -> node2 ; R -> node1, node3 - // Shard 10,11 : P -> node3 ; R -> node1, node2 - // - // Resulting per-node counts on eligible nodes: - // node0: 6P + 0R = 6 - // node1: 2P + 4R = 6 - // node2: 2P + 4R = 6 - // node3: 2P + 4R = 6 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 }; @@ -1255,7 +1217,6 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { assertEquals("eligible node2 total shards", 6, initialTotals[2]); assertEquals("eligible node3 total shards", 6, initialTotals[3]); - // Drive reroute repeatedly to allow rebalance to converge. clusterState = strategy.reroute(clusterState, "reroute"); clusterState = applyStartedShardsUntilNoChange(clusterState, strategy); for (int i = 0; i < 20; i++) { @@ -1285,17 +1246,8 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { minEligible = Math.min(minEligible, finalPrimaries[i]); } assertEquals("all primaries should stay on eligible nodes", numberOfShards, totalEligiblePrimaries); - - // With the fix (dynamic switch above set to true), avgPrimaryShardsPerNode is computed - // over eligible nodes only: 12 / 4 = 3, so allowed = ceil(3 * 1.10) = 4. Eligible nodes - // with 2 primaries are no longer treated as "breached" rebalance targets, and node0 (6P) - // can migrate primaries to node1/2/3 until the layout converges to 3/3/3/3. - assertEquals("primary count on eligible node0 should converge to 3", 3, finalPrimaries[0]); - assertEquals("primary count on eligible node1 should converge to 3", 3, finalPrimaries[1]); - assertEquals("primary count on eligible node2 should converge to 3", 3, finalPrimaries[2]); - assertEquals("primary count on eligible node3 should converge to 3", 3, finalPrimaries[3]); assertEquals( - "Fix applied: primaries fully balanced across eligible nodes -- final distribution: " + "Primaries fully balanced across eligible nodes -- final distribution: " + "node0=" + finalPrimaries[0] + ", node1=" @@ -1309,9 +1261,6 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { ); } - /** - * Helper: returns primary-shard counts for the given index on each node, indexed the same as {@code nodesList}. - */ private static int[] countPrimariesPerNode(ClusterState state, List nodesList, String indexName) { int[] counts = new int[nodesList.size()]; RoutingNodes routingNodes = state.getRoutingNodes(); From db34ac138f72e143d2303c5373ed938633794f36 Mon Sep 17 00:00:00 2001 From: guojialiang Date: Tue, 25 Aug 2026 22:16:25 +0800 Subject: [PATCH 3/5] refactor code Signed-off-by: guojialiang --- .../allocator/LocalShardsBalancer.java | 16 +-- .../allocator/LocalShardsBalancerTests.java | 103 ++++++++++++++++++ 2 files changed, 112 insertions(+), 7 deletions(-) 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 d257d4177be14..8622a05119cdd 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 primaryBalanceNodeCount; private final BalancedShardsAllocator.NodeSorter sorter; private final Set inEligibleTargetNode; private final Supplier timedOutFunc; @@ -96,19 +97,20 @@ public LocalShardsBalancer( this.routingNodes = allocation.routingNodes(); this.metadata = allocation.metadata(); int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); - int nodeCount = routingNodes.size(); + int primaryBalanceNodeCount = routingNodes.size(); if (preferPrimaryFilterAware) { - int eligible = 0; + int eligibleNodeCount = 0; for (RoutingNode routingNode : routingNodes) { if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { - eligible++; + eligibleNodeCount++; } } - if (eligible > 0) { - nodeCount = eligible; + if (eligibleNodeCount > 0) { + primaryBalanceNodeCount = eligibleNodeCount; } } - avgPrimaryShardsPerNode = ((float) primarySum) / nodeCount; + this.primaryBalanceNodeCount = primaryBalanceNodeCount; + avgPrimaryShardsPerNode = ((float) primarySum) / primaryBalanceNodeCount; nodes = Collections.unmodifiableMap(buildModelFromAssigned()); sorter = newNodeSorter(); inEligibleTargetNode = new HashSet<>(); @@ -136,7 +138,7 @@ public float avgShardsPerNode(String index) { @Override public float avgPrimaryShardsPerNode(String index) { - return ((float) metadata.index(index).getNumberOfShards()) / nodes.size(); + return ((float) metadata.index(index).getNumberOfShards()) / primaryBalanceNodeCount; } @Override 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 e401d813dc078..c4ca6981f737e 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 @@ -32,8 +32,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; @@ -261,4 +263,105 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat return decider.apply(shardRouting); } } + + public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { + 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(); + + // ---- Round 1: exclude {node4,node5,node6} -> 3 eligible -> 6 / 3 = 2.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, + /*preferPrimaryFilterAware*/ true, + false, + null + ); + assertEquals("round1 per-index", 2.0f, balancer1.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round1 cluster-level", 2.0f, balancer1.avgPrimaryShardsPerNode(), 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, + /*preferPrimaryFilterAware*/ true, + false, + null + ); + assertEquals("round2 per-index", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round2 cluster-level", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode(), 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, + /*preferPrimaryFilterAware*/ true, + false, + null + ); + assertEquals("round3 per-index", 1.0f, balancer3.avgPrimaryShardsPerNode("test_0"), 0.0001f); + assertEquals("round3 cluster-level", 1.0f, balancer3.avgPrimaryShardsPerNode(), 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; + } + } + } From 3741cbe1e32240dc03d9a8918adc5a0104c5882d Mon Sep 17 00:00:00 2001 From: guojialiang Date: Wed, 26 Aug 2026 13:43:20 +0800 Subject: [PATCH 4/5] refactor code Signed-off-by: guojialiang --- .../allocation/allocator/LocalShardsBalancerTests.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 c4ca6981f737e..d49f15b244c47 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 @@ -293,7 +293,7 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { 0, false, false, - /*preferPrimaryFilterAware*/ true, + true, false, null ); @@ -319,7 +319,7 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { 0, false, false, - /*preferPrimaryFilterAware*/ true, + true, false, null ); @@ -343,7 +343,7 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { 0, false, false, - /*preferPrimaryFilterAware*/ true, + true, false, null ); From 74368b81649d75a088b9bfc61ee31bdd94d60e22 Mon Sep 17 00:00:00 2001 From: guojialiang Date: Wed, 26 Aug 2026 18:08:26 +0800 Subject: [PATCH 5/5] refactor code Signed-off-by: guojialiang --- .../allocator/BalancedShardsAllocator.java | 18 ++++++------ .../allocator/LocalShardsBalancer.java | 20 ++++++------- .../common/settings/ClusterSettings.java | 2 +- .../allocation/BalanceConfigurationTests.java | 2 +- .../allocator/LocalShardsBalancerTests.java | 29 ++++++++++++++----- 5 files changed, 42 insertions(+), 29 deletions(-) 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 dfd1bde21dd23..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,8 +190,8 @@ public class BalancedShardsAllocator implements ShardsAllocator { Property.NodeScope ); - public static final Setting PREFER_PRIMARY_FILTER_AWARE = Setting.boolSetting( - "cluster.routing.allocation.balance.prefer_primary.filter_aware", + public static final Setting PREFER_FILTER_AWARE_BALANCE = Setting.boolSetting( + "cluster.routing.allocation.balance.filter_aware", false, Property.Dynamic, Property.NodeScope @@ -245,7 +245,7 @@ private static Priority parseReroutePriority(String priorityString) { private volatile boolean preferPrimaryShardBalance; private volatile boolean preferPrimaryShardRebalance; - private volatile boolean preferPrimaryFilterAware; + private volatile boolean preferFilterAwareBalance; private volatile float preferPrimaryShardRebalanceBuffer; private volatile float indexBalanceFactor; private volatile float shardBalanceFactor; @@ -274,7 +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)); - setPreferPrimaryFilterAware(PREFER_PRIMARY_FILTER_AWARE.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)); @@ -285,7 +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_PRIMARY_FILTER_AWARE, this::setPreferPrimaryFilterAware); + 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); @@ -378,8 +378,8 @@ private void setPreferPrimaryShardRebalance(boolean preferPrimaryShardRebalance) this.weightFunction.updateRebalanceConstraint(CLUSTER_PRIMARY_SHARD_REBALANCE_CONSTRAINT_ID, preferPrimaryShardRebalance); } - private void setPreferPrimaryFilterAware(boolean preferPrimaryFilterAware) { - this.preferPrimaryFilterAware = preferPrimaryFilterAware; + private void setPreferFilterAwareBalance(boolean preferFilterAwareBalance) { + this.preferFilterAwareBalance = preferFilterAwareBalance; } private void setThreshold(float threshold) { @@ -423,7 +423,7 @@ public void allocate(RoutingAllocation allocation) { threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, - preferPrimaryFilterAware, + preferFilterAwareBalance, ignoreThrottleInRestore, this::allocatorTimedOut ); @@ -450,7 +450,7 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f threshold, preferPrimaryShardBalance, preferPrimaryShardRebalance, - preferPrimaryFilterAware, + preferFilterAwareBalance, ignoreThrottleInRestore, () -> false // as we don't need to check if timed out or not while just understanding ShardAllocationDecision ); 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 8622a05119cdd..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,7 +72,7 @@ public class LocalShardsBalancer extends ShardsBalancer { private final Metadata metadata; private final float avgPrimaryShardsPerNode; - private final int primaryBalanceNodeCount; + private final int balanceNodeCount; private final BalancedShardsAllocator.NodeSorter sorter; private final Set inEligibleTargetNode; private final Supplier timedOutFunc; @@ -86,7 +86,7 @@ public LocalShardsBalancer( float threshold, boolean preferPrimaryBalance, boolean preferPrimaryRebalance, - boolean preferPrimaryFilterAware, + boolean preferFilterAwareBalance, boolean ignoreThrottleInRestore, Supplier timedOutFunc ) { @@ -97,8 +97,8 @@ public LocalShardsBalancer( this.routingNodes = allocation.routingNodes(); this.metadata = allocation.metadata(); int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); - int primaryBalanceNodeCount = routingNodes.size(); - if (preferPrimaryFilterAware) { + int balanceNodeCount = routingNodes.size(); + if (preferFilterAwareBalance) { int eligibleNodeCount = 0; for (RoutingNode routingNode : routingNodes) { if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { @@ -106,11 +106,11 @@ public LocalShardsBalancer( } } if (eligibleNodeCount > 0) { - primaryBalanceNodeCount = eligibleNodeCount; + balanceNodeCount = eligibleNodeCount; } } - this.primaryBalanceNodeCount = primaryBalanceNodeCount; - avgPrimaryShardsPerNode = ((float) primarySum) / primaryBalanceNodeCount; + this.balanceNodeCount = balanceNodeCount; + avgPrimaryShardsPerNode = ((float) primarySum) / balanceNodeCount; nodes = Collections.unmodifiableMap(buildModelFromAssigned()); sorter = newNodeSorter(); inEligibleTargetNode = new HashSet<>(); @@ -133,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()) / primaryBalanceNodeCount; + return ((float) metadata.index(index).getNumberOfShards()) / balanceNodeCount; } @Override @@ -151,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 ede763de8e42d..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,7 +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_PRIMARY_FILTER_AWARE, + 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 ef3a0af6351db..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 @@ -1149,7 +1149,7 @@ public void testPrimaryRebalanceIgnoresAllocationFilter() { excludeList.append("node").append(i); } settingsBuilder.put("cluster.routing.allocation.exclude._id", excludeList.toString()); - settingsBuilder.put("cluster.routing.allocation.balance.prefer_primary.filter_aware", true); + settingsBuilder.put("cluster.routing.allocation.balance.filter_aware", true); AllocationService strategy = createAllocationService(settingsBuilder.build(), new TestGatewayAllocator()); 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 d49f15b244c47..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; @@ -264,7 +265,7 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingAllocation allocat } } - public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { + public void testAvgShardsPerNodeRefreshesWhenFilterChanges() { int numberOfShards = 6; Metadata metadata = buildMetadata(Metadata.builder(), 1, numberOfShards, 1, 0); RoutingTable routingTable = buildRoutingTable(metadata); @@ -274,7 +275,13 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { .nodes(DiscoveryNodes.builder().add(node1).add(node2).add(node3).add(node4).add(node5).add(node6)) .build(); - // ---- Round 1: exclude {node4,node5,node6} -> 3 eligible -> 6 / 3 = 2.0 + // 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")))) @@ -297,8 +304,10 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { false, null ); - assertEquals("round1 per-index", 2.0f, balancer1.avgPrimaryShardsPerNode("test_0"), 0.0001f); - assertEquals("round1 cluster-level", 2.0f, balancer1.avgPrimaryShardsPerNode(), 0.0001f); + 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( @@ -323,8 +332,10 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { false, null ); - assertEquals("round2 per-index", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode("test_0"), 0.0001f); - assertEquals("round2 cluster-level", 6.0f / 5.0f, balancer2.avgPrimaryShardsPerNode(), 0.0001f); + 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( @@ -347,8 +358,10 @@ public void testAvgPrimaryShardsPerNodeRefreshesWhenFilterChanges() { false, null ); - assertEquals("round3 per-index", 1.0f, balancer3.avgPrimaryShardsPerNode("test_0"), 0.0001f); - assertEquals("round3 cluster-level", 1.0f, balancer3.avgPrimaryShardsPerNode(), 0.0001f); + 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 {