Add balance.disk_usage and balance.mode settings for dimensionally-consistent shard rebalancing (#15520)

[AndreKurait]

Summary

Resolves #15520. Adds two opt-in, default-preserving knobs to the shard allocator's weight function, plus a precision fix to the disk-usage term.

Setting	Type	Default	What it does
cluster.routing.allocation.balance.disk_usage	float	0.0f	Adds a disk-byte balance term to the weight function.
cluster.routing.allocation.balance.mode	enum count | ratio	count	Selects the scale on which the three balance terms are expressed.

Both are dynamic, node-scope. Both defaults reproduce pre-PR behavior bit-for-bit (with one exception: the disk term now subtracts in double precision instead of float — a bug fix that only matters at very large cluster sizes; see "Precision fix" below).

Background

BalancedShardsAllocator.WeightFunction scores each (node, index) pair so the allocator can pick moves that reduce imbalance. Before this PR it looked like:

weight = shardBalance * (node.numShards()      - avgShardsPerNode)
       + indexBalance * (node.numShards(index) - avgShardsPerNode(index))

Problem 1 — no disk signal

Two nodes with identical shard counts can carry wildly different byte loads when shard sizes are heterogeneous. The allocator had no signal to correct this.

Problem 2 — the three terms are not on the same scale

The existing shard and per-index terms are raw shard-count deltas. A node with 5 shards above average contributes 5.0. A naively-added disk term expressed as a ratio ((bytes - avgBytes) / avgBytes) has magnitude O(1). Setting disk_usage ≈ shard does not mean "equal influence" — the disk contribution is swamped by two orders of magnitude on any realistic cluster.

What this PR does

1. balance.disk_usage — new disk-byte balance factor

Adds a third term to the weight function:

+ diskUsageBalance * (node.diskUsageInBytes() - avgDiskUsage) / avgDiskUsage

When avgDiskUsage == 0 (empty cluster) the disk term contributes zero. When the factor is 0.0f (default), the term short-circuits and the allocator does not bother tracking per-shard bytes at all (see trackDiskUsage on LocalShardsBalancer). Upgrades are zero-overhead.

2. balance.mode — count (default) vs ratio

Resolves the units mismatch from Problem 2 opt-in, without changing any existing cluster's behavior.

count mode (default):

weight_shard(n)   = shardBalance    * (n.numShards    - avg)
weight_index(n,i) = indexBalance    * (n.numShards(i) - avg(i))
weight_disk(n)    = diskUsageBalance * (n.bytes       - avgBytes) / max(1, avgBytes)

Identical to pre-PR semantics. threshold (default 1.0) keeps its shard-count-delta interpretation ("at least one shard of imbalance").

ratio mode:

weight_shard(n)   = shardBalance    * (n.numShards    - avg)      / max(1, avg)
weight_index(n,i) = indexBalance    * (n.numShards(i) - avg(i))   / max(1, avg(i))
weight_disk(n)    = diskUsageBalance * (n.bytes       - avgBytes) / max(1, avgBytes)

All three terms are relative deviations from the per-axis cluster average. The three balance factors operate on the same dimensionless scale and can be tuned directly relative to each other.

In ratio mode, threshold is interpreted as a relative-deviation fraction (e.g. 0.1 = "at least 10% imbalance"). Operators enabling ratio mode will typically also lower threshold from its default of 1.0. This is documented on BALANCE_MODE_SETTING.

Caveat — small-index amplification. Ratio mode amplifies the weight of indices with few shards: a single misplaced shard of a 3-shard index is a 33% deviation, while the same misplacement on a 30-shard index is a 3% deviation. This is semantically correct (small indices genuinely are more imbalanced per misplaced shard) but changes prioritization relative to count mode. Documented on the setting.

3. Precision fix for the disk term

ShardsBalancer.avgDiskUsageInBytesPerNode() previously returned float. float has ~7 significant digits, so on clusters with total disk usage above ~10 TB the (long bytes - float avgBytes) subtraction lost sub-MB deltas entirely before the divide — the disk term silently stopped seeing imbalances below roughly MB-scale. The return type is now double; the subtraction is done in double (exact for long byte counts up to 2^53 ≈ 9 PB) and the divide is also in double, with a single cast to float for the final weight contribution. This applies to both modes; it is a bug fix, not a behavior change at any realistic cluster size.

No @PublicApi is affected — ShardsBalancer is @opensearch.internal.

Key changes

• BalancedShardsAllocator — new DISK_USAGE_BALANCE_FACTOR_SETTING (dynamic, node-scope, finite-float validator rejecting NaN/±Inf) and new BALANCE_MODE_SETTING (enum). WeightFunction.weight(...) branches on ratioMode and normalizes shard/index terms in ratio mode. Javadoc on both settings and on WeightFunction describes the modes, the threshold interaction, and the small-index amplification caveat.
• BalanceMode (new) — small internal enum with a parse(String) that surfaces a clear error for unknown values.
• ModelNode / LocalShardsBalancer / ShardsBalancer — tracks diskUsageInBytes only when the disk factor is non-zero (see trackDiskUsage); avgDiskUsageInBytesPerNode is double and recomputed from live model state (not a stale snapshot).
• ClusterSettings — one-line registrations alongside existing balance.* settings.

Testing

Unit tests (DiskUsageBalanceTests):

• weight-function algebra includes the normalized disk term iff factor > 0,
• avgDiskUsageInBytesPerNode stays consistent with per-node state across allocateUnassigned / moveShards / tryRelocateShard (regression guard),
• validator rejects NaN and ±Infinity,
• dynamic update path,
• paired correctness tests — testDiskUsageBalanceMovesShardWhenCountBalanced sets up a 3-node / 2-index cluster with counts balanced (3+3+3) but bytes skewed onto node-0 and asserts ≥1 large shard relocates once the factor is enabled; its twin testDiskUsageBalanceNoOpWhenZero runs the same scenario at factor=0.0f and asserts zero relocations,
• ratio-mode unit tests (new) — testRatioModeNormalizesAllThreeTerms asserts the expected algebra for all three terms, testRatioModeHandlesZeroDenominators guards the max(1, avg) divisor paths, testDefaultModeKeepsRawCountDeltas pins count-mode semantics.

Integration tests (DiskUsageBalanceIT, 4 specs, all pass locally in 14.4s):

• testDiskUsageBalanceRebalancesByBytes — end-to-end count-mode disk rebalance: byte spread does not grow when disk_usage balance is enabled.
• testDefaultBehaviorUnchanged — default 0.0f is byte-spread-agnostic (allocator strictly count-balances).
• testRatioModeRebalancesByBytes (new) — with balance.mode=ratio, disk_usage=1.0f, threshold=0.1f: cluster stays green, all shards assigned, byte spread does not exceed count-mode baseline.
• testBalanceModeDynamicToggle (new) — toggles balance.mode between ratio and count four times at runtime; ensures the addSettingsUpdateConsumer path rebuilds the weight function without leaving shards unassigned. Also asserts that balance.mode=bogus is rejected with a clear error.

Suite resets all three settings via @After to prevent state leaking across specs on failure.

Migration

None.

• balance.disk_usage defaults to 0.0f (disabled).
• balance.mode defaults to count (pre-PR behavior).
• The disk-term precision fix is a no-op for any cluster under a few TB of data; above that, it makes the disk term more accurate (fewer below-threshold moves ignored).

All settings are dynamic — enable/tune/disable at runtime without restart.

Unrelated included fix

Commit on this branch removes too-short assertBusy timeouts in IndicesRequestCacheCleanupIT — a known flaky tracked under #21397. The fix is also up standalone as #21494 so it can land on main independently. If #21494 merges first, that commit falls out of this branch on rebase.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 0256021)

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 Stale disk-usage bookkeeping on dynamic toggle trackDiskUsage is fixed at ModelNode construction time based on whether diskUsageBalance != 0f at the moment buildModelFromAssigned runs. If an operator enables balance.disk_usage at runtime, the weight function is rebuilt immediately, but any in-flight balancer using a previously-built node model will still compute disk weights against diskUsageInBytes == 0 for all nodes (the term contributes zero because all nodes look "average"). The disk term only becomes effective on the next reroute that triggers a fresh buildModelFromAssigned. This is a correctness/responsiveness concern worth verifying — particularly because the IT toggles the setting and immediately calls reroute; in production the timing is less guaranteed. void updatePrimaryConstraintThreshold(long primaryConstraintThreshold) { this.primaryConstraintThreshold = primaryConstraintThreshold; } } /** * A model node. * Possibly Incorrect Assertion testNegativeSettingRejected asserts the rejection message contains "must be >=" or "-0.1" or "disk_usage". The standard Setting.floatSetting lower-bound error message wording may not contain any of these substrings exactly (e.g., it commonly says "Failed to parse value [-0.1] for setting [cluster.routing.allocation.balance.disk_usage] must be >= 0.0"). Verify the assertion matches the actual error message; otherwise this test may fail spuriously or hide a real change in validation behavior. String msg = iae.getMessage() == null ? "" : iae.getMessage(); assertTrue( "expected rejection message to mention the lower bound, got: " + msg, msg.contains("must be >=") || msg.contains("-0.1") || msg.toLowerCase(java.util.Locale.ROOT).contains("disk_usage") ); } Test Setup Bug In testRatioModeNormalizesAllThreeTerms, the comment says "node has 4 total shards (avgShards=2)" but makeModelNode("n1", 100L, 3) produces a node with 3 shards. The expected computation uses (3-2)/2 = 0.5 which matches 3 shards, so the comment is just stale — but the comment about index ratio = (3 - 1) / 1 = 2.0 assumes node.numShards(INDEX_SMALL) == 3, which is true only because makeModelNode puts every shard in INDEX_SMALL. This is fragile and the inline comments contradict each other; please reconcile to avoid confusion when the helper is later modified. // 3-shard index on a node that has 2 of the 3 shards (avgShardsForIndex = 1.0) // Node has 4 total shards (avgShards = 2.0) and 300 bytes (avgBytes = 150) // Ratio weights: (4-2)/2 = 1.0, (2-1)/1 = 1.0, (300-150)/150 = 1.0 // All three terms equal → weighted sum = theta0+theta1+theta2 = 1.0 BalancedShardsAllocator.ModelNode node = makeModelNode("n1", 100L, 3); // 300 bytes, 3 shards // Force numShards and numShards(INDEX_SMALL) via direct construction rather than // relying on makeModelNode's default behavior (which puts all shards in INDEX_SMALL). // node.numShards() == 3, node.numShards(INDEX_SMALL) == 3 // For avgShards=2, avgShardsForIndex=1, avgDisk=150: // shard ratio = (3 - 2) / 2 = 0.5 // index ratio = (3 - 1) / 1 = 2.0 // disk ratio = (300 - 150) / 150 = 1.0 ShardsBalancer balancer = new FixedAverageShardsBalancer(2.0f, 1.0f, 150.0); float sum = indexBalance + shardBalance + diskUsageBalance; float theta0 = shardBalance / sum; float theta1 = indexBalance / sum; float theta2 = diskUsageBalance / sum; float expected = theta0 * 0.5f + theta1 * 2.0f + theta2 * 1.0f; assertEquals(expected, ratio.weight(balancer, node, INDEX_SMALL), 1e-4f); Weak Invariant testDiskUsageBalanceRebalancesByBytes and testRatioModeRebalancesByBytes only assert spreadAfter <= spreadBefore. Because the count-balanced 4-shards-per-node layout already happens to have low spread for this fake-size function (and per-index distribution constraints further limit movement), this assertion can pass without any disk-usage-driven rebalancing actually occurring. Consider mirroring the unit test's "count-balanced but byte-skewed" initial layout to make the IT a real positive test for byte-driven rebalancing rather than a no-regression check. public void testDiskUsageBalanceRebalancesByBytes() throws Exception { final MockInternalClusterInfoService clusterInfoService = getMockInternalClusterInfoService(); clusterInfoService.setUpdateFrequency(TimeValue.timeValueMillis(200)); clusterInfoService.setShardSizeFunctionAndRefresh(DiskUsageBalanceIT::fakeShardSize); final List<String> nodeIds = StreamSupport.stream( client().admin().cluster().prepareState().get().getState().getRoutingNodes().spliterator(), false ).map(RoutingNode::nodeId).collect(Collectors.toList()); assertThat("cluster must have 3 data nodes", nodeIds.size(), equalTo(3)); // 6 small + 6 large shards, 0 replicas. With 3 nodes the default shard-count // balancer places 4 shards per node. assertAcked(prepareCreate(INDEX_SMALL).setSettings(Settings.builder().put("number_of_shards", 6).put("number_of_replicas", 0))); assertAcked(prepareCreate(INDEX_LARGE).setSettings(Settings.builder().put("number_of_shards", 6).put("number_of_replicas", 0))); ensureGreen(INDEX_SMALL, INDEX_LARGE); clusterInfoService.refresh(); final Map<String, Integer> initialShardCounts = getShardCountByNodeId(); for (String nodeId : nodeIds) { assertThat( "node " + nodeId + " should have 4 shards under default shard-count balancer", initialShardCounts.get(nodeId), equalTo(4) ); } assertThat("no unassigned shards initially", totalUnassigned(), equalTo(0)); final long spreadBefore = spread(getBytesByNodeId(DiskUsageBalanceIT::fakeShardSize)); logger.info("--> byte spread before enabling disk_usage balance = {}", spreadBefore); // Turn on disk-usage balance alongside shard-count balance. We intentionally leave // SHARD_BALANCE_FACTOR at its default so the allocator's weight function stays well // conditioned (index+shard+disk_usage must sum to > 0, and we don't want to force // arbitrarily-large relocations). assertAcked( client().admin() .cluster() .prepareUpdateSettings() .setPersistentSettings(Settings.builder().put(BalancedShardsAllocator.DISK_USAGE_BALANCE_FACTOR_SETTING.getKey(), 1.0f)) ); clusterInfoService.refresh(); assertAcked(client().admin().cluster().prepareReroute()); // Let any rebalance triggered by the setting update settle, then verify invariants. assertBusy(() -> { ensureGreen(INDEX_SMALL, INDEX_LARGE); assertThat("still no unassigned shards", totalUnassigned(), equalTo(0)); final long spreadAfter = spread(getBytesByNodeId(DiskUsageBalanceIT::fakeShardSize)); logger.info("--> byte spread after enabling disk_usage balance = {}", spreadAfter); assertThat("byte spread should not increase once disk_usage balance is enabled", spreadAfter, lessThanOrEqualTo(spreadBefore)); }, 30, java.util.concurrent.TimeUnit.SECONDS); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 0256021

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Initialize prefer-primary flags before building weight function setPreferPrimaryShardBalance and setPreferPrimaryShardRebalance are not invoked before updateWeightFunction() is called in this constructor, so the initial WeightFunction is built with default false values for preferPrimaryShardBalance/preferPrimaryShardRebalance regardless of cluster settings. Ensure these fields are initialized from settings prior to updateWeightFunction() so the initial weight function reflects user configuration. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [341-349] public BalancedShardsAllocator(Settings settings, ClusterSettings clusterSettings) { setShardBalanceFactor(SHARD_BALANCE_FACTOR_SETTING.get(settings)); setIndexBalanceFactor(INDEX_BALANCE_FACTOR_SETTING.get(settings)); setDiskUsageBalanceFactor(DISK_USAGE_BALANCE_FACTOR_SETTING.get(settings)); setBalanceMode(BALANCE_MODE_SETTING.get(settings)); + setPreferPrimaryShardBalance(PREFER_PRIMARY_SHARD_BALANCE.get(settings)); + setPreferPrimaryShardRebalance(PREFER_PRIMARY_SHARD_REBALANCE.get(settings)); setPreferPrimaryShardRebalanceBuffer(PRIMARY_SHARD_REBALANCE_BUFFER.get(settings)); setIgnoreThrottleInRestore(IGNORE_THROTTLE_FOR_REMOTE_RESTORE.get(settings)); updateWeightFunction(); setThreshold(THRESHOLD_SETTING.get(settings)); Suggestion importance[1-10]: 7 __ Why: Potentially valid concern: if setPreferPrimaryShardBalance/setPreferPrimaryShardRebalance are not called before updateWeightFunction(), the initial WeightFunction may be built with default false values. However, this depends on code outside the visible diff and may already be handled elsewhere.	Medium
	Avoid duplicating internal key-format logic in tests This duplicates the package-private ClusterInfo.shardIdentifierFromRouting logic; if that internal format ever changes, ClusterInfo.getShardSize will return 0L here and these tests will silently pass despite the disk bookkeeping being broken. Prefer invoking the actual ClusterInfo API (e.g. via the shardSizes map keyed using a helper exposed to tests) rather than re-implementing the key format. server/src/test/java/org/opensearch/cluster/routing/allocation/allocator/DiskUsageBalanceTests.java [139-141] +/** Mirrors {@code ClusterInfo.shardIdentifierFromRouting}; keep in sync if that format changes. */ private static String shardIdentifier(ShardRouting shardRouting) { + // TODO: replace with a test-accessible helper from ClusterInfo to avoid silent drift. return shardRouting.shardId().toString() + "[" + (shardRouting.primary() ? "p" : "r") + "]"; } Suggestion importance[1-10]: 4 __ Why: Valid maintainability concern about test fragility if the internal key format changes, but the improved code only adds a comment without solving the problem.	Low
	Remove implicit tracking-enabling constructor overload The two-arg constructor ModelNode(RoutingNode, ClusterInfo) unconditionally enables trackDiskUsage=true, but test code uses it even when the disk-usage factor is zero. More importantly, callers that depend on the zero-overhead default must use the three-arg form; consider making trackDiskUsage derive from a non-null/non-empty ClusterInfo only when explicitly requested, to avoid accidental tracking from intermediate callers. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [827-839] ModelNode(RoutingNode routingNode) { this(routingNode, ClusterInfo.EMPTY, false); } -ModelNode(RoutingNode routingNode, ClusterInfo clusterInfo) { - this(routingNode, clusterInfo, true); +ModelNode(RoutingNode routingNode, ClusterInfo clusterInfo, boolean trackDiskUsage) { + this.routingNode = routingNode; + this.clusterInfo = clusterInfo == null ? ClusterInfo.EMPTY : clusterInfo; + this.trackDiskUsage = trackDiskUsage; } Suggestion importance[1-10]: 3 __ Why: Minor API hygiene suggestion; the existing two-arg overload is used by tests and the overhead is intentional. The risk of accidental enabling is low since callers in production code use the three-arg form.	Low

---

Previous suggestions

Suggestions up to commit fd87cc9

Category	Suggestion	Impact
Possible issue	Cache shard sizes to prevent accounting drift on removal clusterInfo.getShardSize may return a different value on removeShard than it did on addShard if the cluster info is refreshed between the two calls, causing diskUsageInBytes to drift negative or become incorrect. The shard size should be cached at addShard time (e.g., in a per-shard map) and the cached value used in removeShard to ensure accounting consistency. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [900-917] +// In ModelNode fields: +private final Map<String, Long> shardSizeCache = new HashMap<>(); + +public void addShard(ShardRouting shard) { + ... + if (trackDiskUsage) { + long size = clusterInfo.getShardSize(shard, 0L); + shardSizeCache.put(shard.allocationId().getId(), size); + diskUsageInBytes += size; + } + numShards++; +} + public void removeShard(ShardRouting shard) { ... if (trackDiskUsage) { - diskUsageInBytes -= clusterInfo.getShardSize(shard, 0L); + Long cached = shardSizeCache.remove(shard.allocationId().getId()); + diskUsageInBytes -= (cached != null ? cached : clusterInfo.getShardSize(shard, 0L)); } numShards--; } Suggestion importance[1-10]: 6 __ Why: If clusterInfo is refreshed between addShard and removeShard, the shard size returned by getShardSize could differ, causing diskUsageInBytes to drift. However, in practice ModelNode is rebuilt fresh each reroute round from buildModelFromAssigned with a snapshot of clusterInfo, so the same ClusterInfo instance is used throughout a single balance round, making drift unlikely in normal operation.	Low
General	Fix stale disk usage logged after shard removal The disk usage debug log is placed after maxNode.removeShard(shard) and --totalShardCount, so maxNode.diskUsageInBytes() already reflects the post-removal state. Move the debug log to before maxNode.removeShard(shard) to accurately capture the pre-move byte counts for both nodes. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [1147-1157] if (weight.diskUsageBalance() != 0f && logger.isDebugEnabled()) { logger.debug( "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", shard.shardId(), maxNode.getNodeId(), maxNode.diskUsageInBytes(), minNode.getNodeId(), minNode.diskUsageInBytes() ); } +maxNode.removeShard(shard); +--totalShardCount; +... minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); Suggestion importance[1-10]: 5 __ Why: The debug log at line 1147 is placed after maxNode.removeShard(shard) (line 1140), so maxNode.diskUsageInBytes() already reflects the post-removal state. Moving the log before the removal would give more accurate pre-move byte counts, but this is a debug logging issue with no functional impact.	Low
	Remove redundant validator that never triggers The Setting.floatSetting overload already enforces a minimum of 0.0f and maximum of Float.MAX_VALUE, but Float.MAX_VALUE is a finite value, so the custom Validator checking Float.isFinite is redundant and will never trigger (since Float.MAX_VALUE passes isFinite). The validator adds no protection and can be removed to simplify the code. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [144-153] -new Setting.Validator<Float>() { - @Override - public void validate(Float value) { - if (value != null && Float.isFinite(value) == false) { - throw new IllegalArgumentException( - "Illegal value for [cluster.routing.allocation.balance.disk_usage]: must be a finite number, got [" + value + "]" - ); - } - } -}, +public static final Setting<Float> DISK_USAGE_BALANCE_FACTOR_SETTING = Setting.floatSetting( + "cluster.routing.allocation.balance.disk_usage", + 0.0f, + 0.0f, + Float.MAX_VALUE, + Property.Dynamic, + Property.NodeScope +); Suggestion importance[1-10]: 4 __ Why: The custom Validator checking Float.isFinite is indeed redundant since Float.MAX_VALUE is finite and the floatSetting overload already enforces the [0.0f, Float.MAX_VALUE] range. Removing it simplifies the code, but this is a minor cleanup with no functional impact.	Low

Suggestions up to commit 72abdd0

Category	Suggestion	Impact
Possible issue	Ensure symmetric shard size tracking on add/remove When removeShard is called, the diskUsageInBytes is decremented using getShardSize(shard, 0L). However, if the shard size was not available in ClusterInfo at the time addShard was called (returning 0L), but becomes available later, or vice versa, the bookkeeping can go negative or drift. More critically, if removeShard is called for a shard that was never added (or added with a different size), diskUsageInBytes can underflow to a large positive value due to unsigned long arithmetic. Consider caching the shard size at addShard time in a map and using the cached value in removeShard to ensure symmetry. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [913-915] +// In ModelNode, add: private final Map<ShardRouting, Long> trackedShardSizes = new HashMap<>(); +// In addShard: if (trackDiskUsage) { - diskUsageInBytes -= clusterInfo.getShardSize(shard, 0L); + long size = clusterInfo.getShardSize(shard, 0L); + trackedShardSizes.put(shard, size); + diskUsageInBytes += size; +} +// In removeShard: +if (trackDiskUsage) { + Long size = trackedShardSizes.remove(shard); + diskUsageInBytes -= (size != null ? size : 0L); } Suggestion importance[1-10]: 6 __ Why: The concern about asymmetric shard size lookups is valid — if ClusterInfo returns different values for the same shard at addShard vs removeShard time, diskUsageInBytes could drift. However, in practice the ModelNode is rebuilt fresh each reroute via buildModelFromAssigned, so the window for drift is limited. The suggestion is architecturally sound but the risk is low given the current usage pattern.	Low
General	Remove unreachable custom validator on float setting The Setting.floatSetting overload with min and max parameters already rejects values below 0.0f and above Float.MAX_VALUE. However, Float.MAX_VALUE is itself a finite value, so the custom validator's Float.isFinite check will never trigger (since the built-in range check runs first and Float.MAX_VALUE is finite). The validator is therefore dead code and can be removed to reduce confusion, or the max should be changed to something like Float.MAX_VALUE / 2 if overflow is a concern. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [144-153] -new Setting.Validator<Float>() { - @Override - public void validate(Float value) { - if (value != null && Float.isFinite(value) == false) { - throw new IllegalArgumentException( - "Illegal value for [cluster.routing.allocation.balance.disk_usage]: must be a finite number, got [" + value + "]" - ); - } - } -}, +public static final Setting<Float> DISK_USAGE_BALANCE_FACTOR_SETTING = Setting.floatSetting( + "cluster.routing.allocation.balance.disk_usage", + 0.0f, + 0.0f, + Property.Dynamic, + Property.NodeScope +); Suggestion importance[1-10]: 4 __ Why: The analysis is correct — Float.isFinite(Float.MAX_VALUE) returns true, so the custom validator's isFinite check is indeed dead code since the built-in range validation runs first. Removing it simplifies the code, though the impact is minor as it doesn't affect runtime behavior.	Low
	Fix misleading pre-move disk usage logging The debug log captures minNode.diskUsageInBytes() before minNode.addShard(...) is called, so the logged value reflects the pre-relocation state of the target node. This is misleading because the log message says "bytes=" implying the current state after the move. The log should either be placed after addShard or clearly labeled as "before relocation". server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [1147-1157] +minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); +++totalShardCount; if (weight.diskUsageBalance() != 0f && logger.isDebugEnabled()) { logger.debug( - "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", + "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={} after add)", shard.shardId(), maxNode.getNodeId(), maxNode.diskUsageInBytes(), minNode.getNodeId(), minNode.diskUsageInBytes() ); } -minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); +return true; Suggestion importance[1-10]: 3 __ Why: The log captures minNode.diskUsageInBytes() before addShard is called, so the logged "bytes=" value is the pre-relocation state of the target node. The suggestion to move the log after addShard is reasonable for accuracy, but the improved_code restructures the control flow (removing the return true from its original location and adding ++totalShardCount inline), which could introduce bugs and doesn't accurately reflect a minimal fix.	Low

Suggestions up to commit 175a6c6

Category	Suggestion	Impact
Possible issue	Cache shard sizes to prevent byte-tracking drift on removal When removeShard is called, diskUsageInBytes is decremented by the shard's current size from clusterInfo. However, if the shard size has changed between when it was added and when it is removed, the bookkeeping will drift and diskUsageInBytes can become negative or incorrect. The size used during removal should match the size that was recorded during addShard. Consider caching the per-shard byte size in a map at addShard time and using that cached value in removeShard. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [913-915] +// In ModelNode, add a field: +private final Map<String, Long> shardSizeCache = new HashMap<>(); + +// In addShard: if (trackDiskUsage) { - diskUsageInBytes -= clusterInfo.getShardSize(shard, 0L); + long size = clusterInfo.getShardSize(shard, 0L); + shardSizeCache.put(shard.allocationId().getId(), size); + diskUsageInBytes += size; } +// In removeShard: +if (trackDiskUsage) { + Long cached = shardSizeCache.remove(shard.allocationId().getId()); + diskUsageInBytes -= (cached != null ? cached : clusterInfo.getShardSize(shard, 0L)); +} + Suggestion importance[1-10]: 6 __ Why: This is a valid concern: if clusterInfo is updated between addShard and removeShard, the byte counts could drift. However, in practice the clusterInfo passed to ModelNode is a snapshot taken at the start of a balance round and doesn't change during the round, so the risk is low in the current implementation. Still, caching would make the bookkeeping more robust.	Low
General	Remove redundant dead-code validator on disk_usage setting The Setting.floatSetting overload already enforces a minimum of 0.0f and a maximum of Float.MAX_VALUE, but Float.MAX_VALUE is a finite value, so the custom Validator checking Float.isFinite will never trigger (since Float.MAX_VALUE is finite and Float.POSITIVE_INFINITY would already be rejected by the range check). The validator is therefore dead code and can be removed to reduce confusion. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [144-153] -new Setting.Validator<Float>() { - @Override - public void validate(Float value) { - if (value != null && Float.isFinite(value) == false) { - throw new IllegalArgumentException( - "Illegal value for [cluster.routing.allocation.balance.disk_usage]: must be a finite number, got [" + value + "]" - ); - } - } -}, +public static final Setting<Float> DISK_USAGE_BALANCE_FACTOR_SETTING = Setting.floatSetting( + "cluster.routing.allocation.balance.disk_usage", + 0.0f, + 0.0f, + Float.MAX_VALUE, + Property.Dynamic, + Property.NodeScope +); Suggestion importance[1-10]: 4 __ Why: The suggestion is correct that the custom Validator checking Float.isFinite is redundant since Float.MAX_VALUE is finite and the range bounds [0.0f, Float.MAX_VALUE] already exclude Float.POSITIVE_INFINITY. Removing it reduces dead code and confusion, though the impact is minor.	Low
	Log post-relocation byte counts for accuracy The debug log captures minNode.diskUsageInBytes() before minNode.addShard(...) is called, so it logs the pre-relocation byte count for the target node rather than the post-relocation value. This makes the log misleading. Move the log statement after minNode.addShard(...), or capture the post-add value explicitly. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [1147-1157] +minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); +++totalShardCount; if (weight.diskUsageBalance() != 0f && logger.isDebugEnabled()) { logger.debug( "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", shard.shardId(), maxNode.getNodeId(), maxNode.diskUsageInBytes(), minNode.getNodeId(), minNode.diskUsageInBytes() ); } -minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); +return true; Suggestion importance[1-10]: 3 __ Why: The suggestion is technically correct that the log captures pre-relocation bytes for minNode, but this is a debug log and the pre-relocation values are arguably more useful for understanding why the relocation was triggered. The improved_code also restructures the control flow (moving ++totalShardCount and return true) which goes beyond just fixing the log ordering.	Low

Suggestions up to commit fa5cf3a

Category	Suggestion	Impact
Possible issue	Prevent negative disk usage from inconsistent shard size lookups When removeShard is called, the diskUsageInBytes field can go negative if getShardSize returns a different value than what was added (e.g., if shard size changed between add and remove). Consider clamping to zero or storing per-shard sizes in a map to ensure consistency between addShard and removeShard. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [913-915] if (trackDiskUsage) { - diskUsageInBytes -= clusterInfo.getShardSize(shard, 0L); + diskUsageInBytes = Math.max(0L, diskUsageInBytes - clusterInfo.getShardSize(shard, 0L)); } Suggestion importance[1-10]: 4 __ Why: While theoretically possible, in practice ClusterInfo.getShardSize returns the same value for the same shard during a single balance round. The diskUsageInBytes going negative is an edge case that would require shard sizes to change mid-round. The Math.max(0L, ...) fix is a defensive measure but may mask real bugs rather than fix them.	Low
General	Remove redundant validator that can never trigger The Setting.floatSetting with Float.MAX_VALUE as the upper bound already rejects values above Float.MAX_VALUE, but Float.MAX_VALUE itself is finite, so the custom Validator checking Float.isFinite is redundant and never triggered. The validator can be removed to simplify the code. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [139-156] public static final Setting<Float> DISK_USAGE_BALANCE_FACTOR_SETTING = Setting.floatSetting( "cluster.routing.allocation.balance.disk_usage", 0.0f, 0.0f, Float.MAX_VALUE, - new Setting.Validator<Float>() { - @Override - public void validate(Float value) { - if (value != null && Float.isFinite(value) == false) { - throw new IllegalArgumentException( - "Illegal value for [cluster.routing.allocation.balance.disk_usage]: must be a finite number, got [" + value + "]" - ); - } - } - }, Property.Dynamic, Property.NodeScope ); Suggestion importance[1-10]: 4 __ Why: The suggestion is correct that Float.isFinite check is redundant since Float.MAX_VALUE is finite and the built-in range validation already handles the bounds. Removing the custom validator simplifies the code without losing any validation coverage.	Low
	Fix misleading debug log showing post-removal source bytes The debug log is placed after maxNode.removeShard(shard) but before minNode.addShard(...), so maxNode.diskUsageInBytes() already reflects the post-removal state. This makes the log misleading since it shows the source node's bytes after the shard was already removed. Move the log before maxNode.removeShard(shard) or capture the values before removal. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [1147-1156] if (weight.diskUsageBalance() != 0f && logger.isDebugEnabled()) { logger.debug( "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", shard.shardId(), maxNode.getNodeId(), - maxNode.diskUsageInBytes(), + maxNode.diskUsageInBytes() + allocation.clusterInfo().getShardSize(shard, 0L), minNode.getNodeId(), minNode.diskUsageInBytes() ); } minNode.addShard(routingNodes.relocateShard(shard, minNode.getNodeId(), shardSize, allocation.changes()).v1()); Suggestion importance[1-10]: 3 __ Why: The observation is correct that maxNode.diskUsageInBytes() is logged after maxNode.removeShard(shard) has already been called (line 1140), making the log show post-removal bytes. However, the proposed fix using allocation.clusterInfo().getShardSize(shard, 0L) to reconstruct the pre-removal value is a workaround rather than a clean fix, and this is only a debug log with minor impact.	Low

Suggestions up to commit a1fd01e

Category	Suggestion	Impact
General	Fix stale disk-usage value in debug log The debug log captures sourceNode.diskUsageInBytes() and targetNode.diskUsageInBytes() before sourceNode.removeShard(shardRouting) is called, which is correct. However, the same pattern in tryRelocateShard logs the bytes after maxNode.removeShard(shard) has already been called, making the logged value for maxNode stale (already decremented). The log in tryRelocateShard should be moved to before the removeShard call. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [1142-1151] if (weight.diskUsageBalance() != 0f && logger.isDebugEnabled()) { logger.debug( - "disk_usage-aware move: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", - shardRouting.shardId(), - sourceNode.getNodeId(), - sourceNode.diskUsageInBytes(), - targetNode.getNodeId(), - targetNode.diskUsageInBytes() + "disk_usage-aware relocation: shard [{}] from [{}] (bytes={}) to [{}] (bytes={})", + shard.shardId(), + maxNode.getNodeId(), + maxNode.diskUsageInBytes(), + minNode.getNodeId(), + minNode.diskUsageInBytes() ); } -sourceNode.removeShard(shardRouting); +maxNode.removeShard(shard); +--totalShardCount; +long shardSize = allocation.clusterInfo().getShardSize(shard, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE); Suggestion importance[1-10]: 6 __ Why: The observation is accurate: the debug log in tryRelocateShard is placed after maxNode.removeShard(shard) has already been called, so maxNode.diskUsageInBytes() reflects the post-removal value. Moving the log before removeShard would make the logged bytes consistent with the pre-move state, improving debuggability.	Low
	Prevent precision loss in large-byte weight calculation The cast of node.diskUsageInBytes() (a long) to float via subtraction with avgDiskUsage (a float) can silently lose precision for large byte values (e.g., terabyte-scale clusters). The subtraction should be done in double arithmetic before casting to float to preserve meaningful precision in the weight calculation. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java [692] final float avgDiskUsage = balancer.avgDiskUsageInBytesPerNode(); -final float weightDiskUsage = avgDiskUsage > 0.0f ? (node.diskUsageInBytes() - avgDiskUsage) / avgDiskUsage : 0.0f; +final float weightDiskUsage = avgDiskUsage > 0.0f ? (float) ((node.diskUsageInBytes() - (double) avgDiskUsage) / avgDiskUsage) : 0.0f; Suggestion importance[1-10]: 5 __ Why: The suggestion is valid: subtracting a long from a float can lose precision for large values (TB-scale). Using double arithmetic before casting to float is a reasonable improvement, though in practice the normalized ratio keeps values small and the impact is limited.	Low
	Avoid O(N²) recomputation of average disk usage This method is called inside the hot weight() path for every node/index pair during balancing. Computing the average by iterating all nodes on every call is O(N) per weight evaluation, making the overall balancing O(N²). The average should be computed once per balance round (e.g., lazily cached with a dirty flag, or recomputed only when buildModelFromAssigned is called) to avoid quadratic overhead on large clusters. server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/LocalShardsBalancer.java [153-162] +// Compute once and cache; invalidate when nodes map changes (addShard/removeShard). +// Example: compute eagerly after buildModelFromAssigned and update incrementally. public float avgDiskUsageInBytesPerNode() { if (nodes.isEmpty()) { return 0f; } long total = 0L; for (BalancedShardsAllocator.ModelNode node : nodes.values()) { total += node.diskUsageInBytes(); } return ((float) total) / nodes.size(); } +// TODO: cache this value and update it incrementally in addShard/removeShard paths +// to avoid O(N) recomputation on every weight() call during balancing. Suggestion importance[1-10]: 4 __ Why: The concern about O(N) per weight() call is valid in principle, but the improved_code is essentially identical to the existing_code with only a TODO comment added, making it a low-quality suggestion. The actual fix (caching) is not implemented, so the score is limited.	Low

[github-actions]

❌ Gradle check result for 4cd59d5: FAILURE

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?

[github-actions]

Persistent review updated to latest commit 3b98738

[github-actions]

❌ Gradle check result for 3b98738: FAILURE

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?

[github-actions]

Persistent review updated to latest commit b7a0075

[github-actions]

❌ Gradle check result for b7a0075: FAILURE

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?

[github-actions]

Persistent review updated to latest commit feece43

[github-actions]

Persistent review updated to latest commit a397cf1

[github-actions]

Persistent review updated to latest commit b900993

[github-actions]

❌ Gradle check result for b900993: FAILURE

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?

[github-actions]

Persistent review updated to latest commit 32a4a34

[github-actions]

❌ Gradle check result for 32a4a34: FAILURE

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?

[github-actions]

Persistent review updated to latest commit 14a8e61

[github-actions]

❌ Gradle check result for 14a8e61: 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?

[github-actions]

Persistent review updated to latest commit 600bea2

[github-actions]

Persistent review updated to latest commit 0d8fa66

[github-actions]

❌ Gradle check result for 0d8fa66: FAILURE

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?

[github-actions]

❌ Gradle check result for a1fd01e: FAILURE

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?

[andrross]

Can you keep the avgDiskUsageInBytesPerNode() as a long so you can do integer arithmetic with the subtraction, then only convert to floating point with the final ratio? I'm not sure the values ever get large enough that the loss of precision will be significant, but it seems more natural to keep byte values as integers (even for an average).

[andrross]

@shwetathareja FYI

[github-actions]

Persistent review updated to latest commit 5fefcba

[github-actions]

❌ Gradle check result for 5fefcba:

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?

[github-actions]

Persistent review updated to latest commit a1fd01e

[AndreKurait]

The detect-breaking-change (japicmp) failure on this PR is not caused by anything in #21343. Since CI runs against the merged tree (main + this branch), the jar under test includes recent upstream commits \u2014 including #21143 (a70e8ee0, "Refactor WLM settings to use Setting objects and rename field"), which removed a @PublicApi method without deprecation.

Full japicmp output filtered for binary-breaking markers (***! / ---!):

***! MODIFIED CLASS: PUBLIC org.opensearch.cluster.metadata.WorkloadGroup  (not serializable)
        ---! REMOVED METHOD: PUBLIC(-) java.util.Map<java.lang.String,java.lang.String> getSearchSettings()

That is the only incompatibility in the whole report. Everything else is (compatible) new-method/new-field additions \u2014 informational, not blocking.

Verification:

• Pre-merge a1fd01e8 passed Detect Breaking Changes (run 25395196071) at 18:38 UTC.
• Post-merge (same PR SHA, but merge-base advanced to include a70e8ee0): fails (run 25401092289) at 20:40 UTC.
• Ran ./gradlew :server:japicmp locally against merge(origin/main, a1fd01e8) and reproduced; WorkloadGroup.getSearchSettings() removal is the sole ---! line.

PR #21343 itself introduces no @PublicApi removals or signature changes. BalancedShardsAllocator, LocalShardsBalancer, ShardsBalancer, and ClusterSettings changes are all additive (new setting, new concrete method with default return, new weight term) and all classes involved are either @opensearch.internal or unannotated.

Every in-flight PR against main right now will hit the same japicmp failure for the same reason. Unblocking options (up to maintainers):

1. Restore WorkloadGroup.getSearchSettings() as a @Deprecated shim delegating to the new accessor.
2. Apply @DeprecatedApi / @InternalApi on WorkloadGroup if it should no longer be @PublicApi.
3. Roll the japicmp.compare.version baseline forward.

cc @dzane17 @andrross

[github-actions]

✅ Gradle check result for a1fd01e: SUCCESS

[dzane17]

@AndreKurait Thanks for the analysis. I have a PR with the potential fix: #21500

[github-actions]

Persistent review updated to latest commit fa5cf3a

[github-actions]

❌ Gradle check result for fa5cf3a: FAILURE

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?

[github-actions]

Persistent review updated to latest commit 175a6c6

[github-actions]

❌ Gradle check result for 175a6c6: FAILURE

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?

[github-actions]

Persistent review updated to latest commit 72abdd0

[github-actions]

Persistent review updated to latest commit fd87cc9

[github-actions]

✅ Gradle check result for fd87cc9: SUCCESS

[github-actions]

Persistent review updated to latest commit 0256021

[github-actions]

❌ Gradle check result for 0256021: FAILURE

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?

[ztomaszewska]

Hello @AndreKurait thank you for this PR. This feature is really needed! Do you think it is feasible to merge it and have it in 3.8 release?

[kcpicot]

Hello @AndreKurait , yes thank you very much for this PR. I follow it since April as we have numerous issues with storage balancing which causes unassigned shard replicas (currently the storage balancing is by index number which is not really pertinent).

Your pull request would be of great use for us, thank you for your work !