Budget-aware low disk mode table rebalance

[J-HowHuang]

Description

Low disk mode today in table rebalance only guarantees that within a segment, the replica to offload will be done first, then the replica to add.

The scope is limited to a segment, not server. However, when servers have limited disk space, we want to further guarantee that at any step, a server won't load more bytes than they initially did if a server is to net lose some bytes, or a server won't load more bytes than it would eventually load if it's to net gain some bytes.

Current implementation won't give this guarantee (see example), and we need this guarantee to avoid jamming disk spaces. Also when the disk is already jammed such that no segments can be added to it, we need this guarantee to rebalance our way out without segments getting into error state due to no disk space.

Example of the current implementation

  3 servers → 3 servers with 2 in the overlap, replication 2, 6 segments, minAvailableReplicas=1:


  segment   start    step 1   step 2   step 3   step 4 = target
  s0        A,B      B        B,C      B,C      B,C
  s1        A,C      C        C,D      D        B,D
  s2        B,C      C        C,D      C,D      C,D
  s3        A,B      B        B,C      B,C      B,C
  s4        A,C      C        C,D      D        B,D
  s5        B,C      C        C,D      C,D      C,D

  segments hosted
  A         4        0        0        0        0
  B         4        2        2        2        4
  C         4        4      → 6 ←      4        4
  D         0        0        4        4        4

New design and new constraint

Scope: we only change how we derive the next step, i.e. the result returned by getNextAssignment. Nothing change in deriving target assignment, only the intermediate steps. Also, for lowDiskMode=false the algorithm is the same.

The illustration here are shown the segment count, but this PR generalize it to segment bytes.

DiskUsageBudget

We compute the budget of each server on how many bytes of segments they can load at most at the beginning of the rebalance.

The idea is that each server won't load more bytes than they initially did if a server is to net lose some bytes, or won't load more bytes than it would eventually load if it's to net gain some bytes.

StepDiskBudget

Compute for each step. This tells you how many bytes you can still take based on the current assignment, respect to the DiskUsageBudget we initially fixed on.

computeNextAssignment

It will only change the segment's assignment if it fits the StepDiskBudget. Reject the new assignment if it violates, so the segment stay at its places for this step. In the case of strict replica group, the cost is computed on the entire replica group since they'll be moved together.

For this mechanism, most of the cases should be able to resolve a sequence that fits the DiskUsageBudget constraint entirely. The above example would become this:

  segment   start    step 1   step 2   step 3   step 4 = target
  s0        A,B      B        B        B        B,C
  s1        A,C      C        C,D      D        B,D
  s2        B,C      C        C,D      C,D      C,D
  s3        A,B      B        B        B        B,C
  s4        A,C      C        C,D      D        B,D
  s5        B,C      C        C,D      C,D      C,D

  segments hosted
  A         4        0        0        0        0
  B         4        2        2        2        4
  C         4        4      → 4 ←      2        4
  D         0        0        4        4        4

Precheck

No guarantee that such sequence is resolvable under the any algorithm, though. If next assignment couldn't be obtained under the constraint, it will fall back to the original low disk mode implementation.

Therefore we have a pre-check to tell if the rebalance will violate this constraint (since the sequence resolution is deterministic, we can verify that during pre-check).

"preChecksResult": {
    "diskUtilization": {
      "preCheckStatus": "ERROR",
      "message": "UNSAFE. Servers with unsafe disk utilization DURING rebalance (>=50%): Server_1 (52%). lowDiskMode cannot avoid it for this target assignment: the rebalance cannot make
  progress without going over the disk these servers start with, by up to Server_1 (150B). Rebalance to a target assignment that frees up space on them first, or add capacity"
    }
  }

It names the server, how far over it would go, and two concrete remedies.

The case it replaces

Before this PR, that same situation returned PASS, because the check assumed lowDiskMode always removed the transient usage:

  {
    "preCheckStatus": "PASS",
    "message": "Within threshold (<50%) AFTER rebalance. Servers that would go over it DURING the rebalance: Server_1 (52%). lowDiskMode avoids that transient disk usage by deleting segments
  before adding the new ones"
  }

That message is still returned — but now only when the replay confirms the budget actually holds, rather than being asserted unconditionally.

The other outcomes on this path, unchanged

  PASS   Within threshold (<50%)

  ERROR  UNSAFE. Servers with unsafe disk utilization AFTER rebalance (>=50%): Server_1 (52%)

  ERROR  UNSAFE. Servers with unsafe disk utilization DURING rebalance (>=50%): Server_1 (52%).
         Enable lowDiskMode to delete segments before adding the new ones

In case of mid-flight uploaded segments

When a new segment appear in ideal state by external sources (e.g. segment upload, consuming segment committed), there are two cases:

1. They are added to the ideal state as the target assignment. This case the rebalance steps outcome won't change
2. They are added to the ideal state that's different to their target assignments. For example when strict replica routing is in place. This case, it might lead to a different computeNextAssignment result.

Every step we check if there are new segments added into the ideal state that's not seen in the beginning, account for their bytes as if they were in the ideal state when we computed the DiskUsageBudget.

Notice that we don't guarantee the disk would be in the safe shape if the segments are added mid-flight because pre-check won't see that. But we'll still move the segments so that no servers would bare additional bytes.

Testing

LowDiskModeRebalanceSimulatorTest drives the real TableRebalancer.getNextAssignment in a loop, exactly as doRebalance does, tracking every server's bytes at each step — so it exercises
production code, not a re-implementation. Three properties are asserted over 21 fixed scenarios:

1. the rebalance always reaches the target assignment;
2. no server exceeds max(bytes it held that this rebalance did not place, bytes the target places on it), unless the disk-utilization pre-check names it up front;
3. under strict replica group routing, segments of a partition are never assigned different instances.

Scenarios cover balanced and replica-group assignment, both routing modes, batchSizePerServer, skewed segment sizes, segments uploaded mid-rebalance, and the two assignments a randomized
search found hardest. main() additionally runs ~32k randomized scenarios for comparing effectiveness by hand; those are not tests.

Constraint violations, before and after

  ┌─────────────────────────────────────┬────────────────────────────────────────┬───────┐
  │               corpus                │                 before                 │ after │
  ├─────────────────────────────────────┼────────────────────────────────────────┼───────┤
  │ old/new server sets, 648            │ 120 over-allocate (18.5%), worst 2.00× │ 0     │
  ├─────────────────────────────────────┼────────────────────────────────────────┼───────┤
  │ strict replica group, 324           │ 92 over-allocate (28.4%), worst 1.78×  │ 0     │
  ├─────────────────────────────────────┼────────────────────────────────────────┼───────┤
  │ random non-uniform assignments, 400 │ 47 over-allocate (11.8%), worst 2.00×  │ 0     │
  └─────────────────────────────────────┴────────────────────────────────────────┴───────┘

Post-change the wider sweeps — 1,944 server-set shapes and 30,000 random group structures — also report zero. Mean step count moves 3.4 → 3.5, so the bound costs essentially nothing.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 84.25926% with 34 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.59%. Comparing base (4824a5f) to head (c499036).
⚠️ Report is 33 commits behind head on master.

Files with missing lines	Patch %	Lines
...ntroller/helix/core/rebalance/TableRebalancer.java	85.78%	14 Missing and 13 partials ⚠️
...lix/core/rebalance/DefaultRebalancePreChecker.java	73.07%	6 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19386      +/-   ##
============================================
+ Coverage     67.20%   67.59%   +0.38%     
- Complexity     1418     1430      +12     
============================================
  Files          3479     3486       +7     
  Lines        223075   224327    +1252     
  Branches      35135    35420     +285     
============================================
+ Hits         149920   151630    +1710     
+ Misses        61184    60663     -521     
- Partials      11971    12034      +63     

Flag	Coverage Δ
integration	100.00% <ø> (+100.00%)	⬆️
integration1	100.00% <ø> (?)
integration2	0.00% <ø> (ø)
java-25	67.59% <84.25%> (+0.38%)	⬆️
lane-a	100.00% <ø> (+100.00%)	⬆️
lane-b	0.00% <ø> (ø)
temurin	67.59% <84.25%> (+0.38%)	⬆️
unittests	67.59% <84.25%> (+0.38%)	⬆️
unittests1	57.67% <ø> (-0.02%)	⬇️
unittests2	39.36% <84.25%> (+0.36%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Jackie-Jiang]

Non-blocking review. The overall direction looks useful, but I found two safety issues that should be addressed or explicitly documented before relying on the new low-disk guarantee.

Critical

1. Live rebalance discards the disk ceiling when progress stalls (TableRebalancer.java:1718)

   The fallback recomputes the step with stepBudget=null, deliberately allowing servers to exceed the initial/target byte ceiling. Because pre-checks are optional and do not gate execution, a full server can still exhaust disk.

   Please abort with the blocked servers and required bytes, or require a separate explicit, default-off override for unsafe continuation.

2. Incomplete replay can be reported as safe (TableRebalancer.java:2452)

   Exceptions and the 10,000-step limit return the same partial/empty map as a completed safe replay. The pre-check interprets an empty map as PASS, although the actual rebalance may later take the unsafe fallback.

   Please return explicit SAFE, UNSAFE, or INCOMPLETE/FAILED status and add a regression requiring more than 10,000 steps.

Major

3. Synchronous replay may be prohibitively expensive (TableRebalancer.java:2441)

   Up to 10,000 iterations repeatedly scan assignments, create assignment-sized maps/sets, and duplicate budget/group calculations. Please use incremental accounting or a bounded work/time budget that returns inconclusive, and benchmark production-scale tables with batching and strict routing.

4. Pre-check integration is not exercised end-to-end (DefaultRebalancePreCheckerTest.java:183)

   The new failure test overrides getServersForcedOverDiskBudget() with a hard-coded result. A wiring error across PreCheckContext, routing mode, batching, segment sizes, and replay would leave both test suites green. Please add a real known-over-budget fixture through DefaultRebalancePreChecker.

5. The performance claim lacks relevant measurements

   The description says the bound “costs essentially nothing,” but reports mean rebalance step count rather than pre-check latency, CPU, allocation, or GC. Please add production-scale measurements before retaining that claim.

Minor

• Consider extracting the disk-budget and replay machinery into a package-private planner; TableRebalancer grows to roughly 2,700 lines.
• retainInstancesThatFit() documentation says non-fitting instances are retained at the back, but the implementation filters them out.

Verification: reviewed head c4990363; focused controller tests passed (24 tests, 0 failures), and GitHub CI was green. I found no compatibility, locking, or version-checked IdealState-write regressions.