Make shard balance aware of allocation filters

[guojialiang92]

Description

The purpose of this PR is to solve the issue mentioned in #[22832].

• I have introduced a test (BalanceConfigurationTests#testPrimaryRebalanceIgnoresAllocationFilter) that can stably reproduce the primary shard imbalance issue, and achieved the effect of primary shard balancing after optimization.
• I have introduced Configuration cluster.routing.allocation.balance.filter_aware to dynamically enable allocation filter aware.
• To ensure the consistency and predictability of the weight calculation results, I will make sure that all weight calculations take the allocation filter into account.

Related Issues

Resolves #[22832]

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 74368b8)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Inconsistent Denominators When preferFilterAwareBalance is true, balanceNodeCount is set to the number of nodes for which canAllocateAnyShardToNode is not NO, and is used as the denominator in avgShardsPerNode, avgShardsPerNode(index), avgPrimaryShardsPerNode, and avgPrimaryShardsPerNode(index). However, the numerators (totalShardCount, metadata.index(index).getTotalNumberOfShards(), getNumberOfShards(), primarySum) still include shards/indices that are pinned to excluded nodes or have index-level filters. If an index cannot be allocated to any eligible node (e.g., index-level exclude), its shards inflate the average against a smaller denominator, skewing the weight function and potentially triggering unnecessary rebalance attempts. Consider computing eligibility per-index or excluding non-eligible indices from the numerator. 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; Setting Name Mismatch The setting key is cluster.routing.allocation.balance.filter_aware but the PR description advertises cluster.routing.allocation.balance.prefer_primary.filter_aware. Since the setting also influences the generic avgShardsPerNode (non-primary) denominator, the current naming is arguably correct, but this discrepancy with the description/docs should be reconciled before release to avoid user confusion and doc drift. logger.trace("Start balancing cluster"); } if (allocation.hasPendingAsyncFetch()) { /* * see https://github.com/elastic/elasticsearch/issues/14387 * if we allow rebalance operations while we are still fetching shard store data Potential Performance Cost canAllocateAnyShardToNode is invoked for every routing node on every allocator invocation when the flag is on. On large clusters with many deciders this iterates all deciders per node and could add measurable overhead to each reroute. Consider caching or short-circuiting when no filter deciders are active. 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; } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 74368b8

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent divide-by-zero for node count Guard against a zero routingNodes.size() to avoid a divide-by-zero when computing avgPrimaryShardsPerNode. Although allocate() returns early for empty routing nodes, LocalShardsBalancer is also constructed from decideShardAllocation, where this precondition is not enforced. Default balanceNodeCount to at least 1. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [99-113] int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); -int balanceNodeCount = routingNodes.size(); +int balanceNodeCount = Math.max(1, 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; Suggestion importance[1-10]: 4 __ Why: The original code had the same potential divide-by-zero (dividing by routingNodes.size()), so this suggestion identifies a pre-existing edge case rather than a regression. It's a minor defensive improvement, especially since allocate() guards against empty routing nodes but decideShardAllocation does not.	Low

---

Previous suggestions

Suggestions up to commit 3741cbe

Category	Suggestion	Impact
Possible issue	Prevent division by zero on node count Guard against primaryBalanceNodeCount being zero. If routingNodes.size() is 0 (which can happen when the allocator is called without nodes), the division will produce infinity/NaN and corrupt weight calculations. Add an explicit check to fall back to a safe default (e.g. 1) when no nodes are available. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [99-113] int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); int primaryBalanceNodeCount = routingNodes.size(); if (preferPrimaryFilterAware) { int eligibleNodeCount = 0; for (RoutingNode routingNode : routingNodes) { if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { eligibleNodeCount++; } } if (eligibleNodeCount > 0) { primaryBalanceNodeCount = eligibleNodeCount; } } -this.primaryBalanceNodeCount = primaryBalanceNodeCount; -avgPrimaryShardsPerNode = ((float) primarySum) / primaryBalanceNodeCount; +this.primaryBalanceNodeCount = Math.max(primaryBalanceNodeCount, 1); +avgPrimaryShardsPerNode = ((float) primarySum) / this.primaryBalanceNodeCount; Suggestion importance[1-10]: 4 __ Why: Guarding against a zero primaryBalanceNodeCount is a reasonable defensive measure, though in practice allocate() short-circuits when routingNodes.size() == 0, making this an edge-case improvement.	Low
General	Align per-index and cluster-level averages The per-index method now unconditionally divides by primaryBalanceNodeCount, but the cluster-level avgPrimaryShardsPerNode field is only filter-aware when preferPrimaryFilterAware is true. This inconsistency changes behavior for existing users who have not enabled the new setting, since primaryBalanceNodeCount still defaults to routingNodes.size() (fine) but any future divergence would be surprising. Consider mirroring the same conditional to keep the two methods aligned or store both counts explicitly. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [140-142] @Override public float avgPrimaryShardsPerNode(String index) { - return ((float) metadata.index(index).getNumberOfShards()) / primaryBalanceNodeCount; + return ((float) metadata.index(index).getNumberOfShards()) / Math.max(primaryBalanceNodeCount, 1); } Suggestion importance[1-10]: 3 __ Why: The suggestion notes a minor consistency concern, but since primaryBalanceNodeCount defaults to routingNodes.size() when the setting is off, behavior is unchanged. The improved code only adds a safety Math.max, which is a marginal improvement.	Low

Suggestions up to commit db34ac1

Category	Suggestion	Impact
Possible issue	Guard divisor against zero nodes Guard against division-by-zero when routingNodes.size() is 0. Although allocate() short-circuits on an empty routing table, LocalShardsBalancer may also be invoked via decideShardAllocation; a zero denominator would produce Infinity/NaN weights and corrupt subsequent balancing decisions. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [99-113] int primarySum = StreamSupport.stream(metadata.spliterator(), false).mapToInt(IndexMetadata::getNumberOfShards).sum(); int primaryBalanceNodeCount = routingNodes.size(); if (preferPrimaryFilterAware) { int eligibleNodeCount = 0; for (RoutingNode routingNode : routingNodes) { if (allocation.deciders().canAllocateAnyShardToNode(routingNode, allocation).type() != Decision.Type.NO) { eligibleNodeCount++; } } if (eligibleNodeCount > 0) { primaryBalanceNodeCount = eligibleNodeCount; } } -this.primaryBalanceNodeCount = primaryBalanceNodeCount; -avgPrimaryShardsPerNode = ((float) primarySum) / primaryBalanceNodeCount; +this.primaryBalanceNodeCount = Math.max(primaryBalanceNodeCount, 1); +avgPrimaryShardsPerNode = ((float) primarySum) / this.primaryBalanceNodeCount; Suggestion importance[1-10]: 4 __ Why: The prior code also divided by routingNodes.size() without a guard, and allocate() already short-circuits on empty routing nodes. The suggestion is a minor defensive improvement with limited practical impact.	Low

[github-actions]

✅ Gradle check result for db34ac1: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.62%. Comparing base (f81134c) to head (74368b8).
⚠️ Report is 10 commits behind head on main.

Files with missing lines	Patch %	Lines
.../allocation/allocator/BalancedShardsAllocator.java	83.33%	1 Missing ⚠️
...ting/allocation/allocator/LocalShardsBalancer.java	93.33%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##               main   #22835   +/-   ##
=========================================
  Coverage     71.61%   71.62%           
- Complexity    77315    77374   +59     
=========================================
  Files          6170     6170           
  Lines        359671   359685   +14     
  Branches      52450    52454    +4     
=========================================
+ Hits         257591   257631   +40     
+ Misses        81636    81623   -13     
+ Partials      20444    20431   -13     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 3741cbe

[github-actions]

❌ Gradle check result for 3741cbe: null

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[alchemist51]

Why not we do this for the avgShardsPerNode as well?

[guojialiang92]

@alchemist51
Great suggestion. I've improved the code and also submitted a reply in the RFC.
I'm looking forward to your continuing to review the code.

[github-actions]

Persistent review updated to latest commit 74368b8

[github-actions]

✅ Gradle check result for 74368b8: SUCCESS