From 1b7fd8c5d7ee2361e727a7c9c12b0969ad141df4 Mon Sep 17 00:00:00 2001 From: J-HowHuang Date: Thu, 6 Aug 2026 14:57:09 -0700 Subject: [PATCH 1/5] consolidate two disk util rebalance prechecks into one --- .../rebalance/DefaultRebalancePreChecker.java | 92 ++++++---- .../DefaultRebalancePreCheckerTest.java | 164 ++++++++++++++++++ .../TableRebalancerClusterStatelessTest.java | 47 ++--- .../tests/TableRebalanceIntegrationTest.java | 9 +- 4 files changed, 251 insertions(+), 61 deletions(-) create mode 100644 pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index 50bb56849efb..cafc3888359d 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -56,8 +56,7 @@ public class DefaultRebalancePreChecker implements RebalancePreChecker { public static final String NEEDS_RELOAD_STATUS = "needsReloadStatus"; public static final String IS_MINIMIZE_DATA_MOVEMENT = "isMinimizeDataMovement"; - public static final String DISK_UTILIZATION_DURING_REBALANCE = "diskUtilizationDuringRebalance"; - public static final String DISK_UTILIZATION_AFTER_REBALANCE = "diskUtilizationAfterRebalance"; + public static final String DISK_UTILIZATION = "diskUtilization"; public static final String REBALANCE_CONFIG_OPTIONS = "rebalanceConfigOptions"; public static final String REPLICA_GROUPS_INFO = "replicaGroupsInfo"; @@ -106,13 +105,10 @@ public Map check(PreCheckContext preCheckCont diskUtilizationThreshold = 1.0; } - // Check if all servers involved in the rebalance have enough disk space for rebalance operation. + // Check if all servers involved in the rebalance have enough disk space, both while the rebalance is running and + // once it is done. // Notice this check could have false positives (disk utilization is subject to change by other operations anytime) - preCheckResult.put(DISK_UTILIZATION_DURING_REBALANCE, - checkDiskUtilization(preCheckContext, diskUtilizationThreshold, true)); - // Check if all servers involved in the rebalance will have enough disk space after the rebalance. - preCheckResult.put(DISK_UTILIZATION_AFTER_REBALANCE, - checkDiskUtilization(preCheckContext, diskUtilizationThreshold, false)); + preCheckResult.put(DISK_UTILIZATION, checkDiskUtilization(preCheckContext, diskUtilizationThreshold)); preCheckResult.put(REBALANCE_CONFIG_OPTIONS, checkRebalanceConfig(rebalanceConfig, tableConfig, preCheckContext.getCurrentAssignment(), preCheckContext.getTargetAssignment(), @@ -269,19 +265,26 @@ private RebalancePreCheckerResult checkIsMinimizeDataMovement(TableConfig tableC } /// Estimates whether the servers of the target assignment stay within the disk utilization threshold, based on the - /// average segment size and the number of segments added to (and, unless checking for the worst case, removed from) - /// each server. Every segment is assumed to take up disk space on each server it is assigned to. Downstream projects - /// where that does not hold (e.g. because a segment can be stored outside of the server, as indicated by + /// average segment size and the number of segments added to and removed from each server. Two points in time are + /// estimated: + /// + /// - **After the rebalance**, i.e. once every server has both added and removed all the segments it has to. Going + /// over the threshold there is an error whatever the rebalance config, since no way of running the rebalance + /// brings the end state back within the threshold. + /// - **During the rebalance**, where a server can transiently hold the segments it is gaining on top of the ones it + /// is about to lose. Only `lowDiskMode` rules that peak out, by waiting for the segments to be deleted before + /// adding the new ones, so going over the threshold there is an error unless it is enabled. Note that `downtime` + /// does not help: it hands every segment to Helix at once without ordering the drops before the adds. + /// + /// Every segment is assumed to take up disk space on each server it is assigned to. Downstream projects where that + /// does not hold (e.g. because a segment can be stored outside of the server, as indicated by /// [TierConfig#getTierBackend()]) can override this. - protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preCheckContext, double threshold, - boolean worstCase) { + protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preCheckContext, double threshold) { Map> currentAssignment = preCheckContext.getCurrentAssignment(); Map> targetAssignment = preCheckContext.getTargetAssignment(); TableSizeReader.TableSubTypeSizeDetails tableSubTypeSizeDetails = preCheckContext.getTableSubTypeSizeDetails(); - boolean isDiskUtilSafe = true; - StringBuilder message = - new StringBuilder("UNSAFE. Servers with unsafe disk utilization (>" + (short) (threshold * 100) + "%): "); - String sep = ""; + List serversUnsafeDuringRebalance = new ArrayList<>(); + List serversUnsafeAfterRebalance = new ArrayList<>(); Map> existingServersToSegmentMap = new HashMap<>(); Map> newServersToSegmentMap = new HashMap<>(); @@ -327,22 +330,47 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec long diskUtilizationGain = newSegmentSet.size() * avgSegmentSize; long diskUtilizationLoss = removedSegmentSet.size() * avgSegmentSize; - long diskUtilizationFootprint = - diskUsage.getUsedSpaceBytes() + diskUtilizationGain - (worstCase ? 0 : diskUtilizationLoss); - double diskUtilizationFootprintRatio = - (double) diskUtilizationFootprint / diskUsage.getTotalSpaceBytes(); - - if (diskUtilizationFootprintRatio >= threshold) { - isDiskUtilSafe = false; - message.append(sep) - .append(server) - .append(String.format(" (%d%%)", (short) (diskUtilizationFootprintRatio * 100))); - sep = ", "; - } + // While the rebalance is running, the segments being added can co-exist with the ones being removed + addIfOverThreshold(serversUnsafeDuringRebalance, server, + (double) (diskUsage.getUsedSpaceBytes() + diskUtilizationGain) / diskUsage.getTotalSpaceBytes(), threshold); + addIfOverThreshold(serversUnsafeAfterRebalance, server, + (double) (diskUsage.getUsedSpaceBytes() + diskUtilizationGain - diskUtilizationLoss) + / diskUsage.getTotalSpaceBytes(), threshold); + } + + // A server over the threshold once the rebalance is done is over it during the rebalance as well, so the end state + // is what to report first: it is both the more severe problem and the one that has to be solved by adding capacity + // rather than by tuning the rebalance config + if (!serversUnsafeAfterRebalance.isEmpty()) { + return RebalancePreCheckerResult.error( + getUnsafeDiskUtilizationMessage("after rebalance", serversUnsafeAfterRebalance, threshold)); + } + String withinThreshold = String.format("Within threshold (<%d%%)", (short) (threshold * 100)); + if (serversUnsafeDuringRebalance.isEmpty()) { + return RebalancePreCheckerResult.pass(withinThreshold); + } + // lowDiskMode is the only way to rule the transient disk usage above out, since it waits for the segments to be + // deleted before adding the new ones. downtime does not: it moves every segment in one shot without ordering the + // drops before the adds, so a server can still end up holding both at once + if (preCheckContext.getRebalanceConfig().isLowDiskMode()) { + return RebalancePreCheckerResult.pass(withinThreshold + " after rebalance. Some servers would go over it during " + + "the rebalance, but lowDiskMode avoids that transient disk usage"); } - return isDiskUtilSafe ? RebalancePreCheckerResult.pass( - String.format("Within threshold (<%d%%)", (short) (threshold * 100))) - : RebalancePreCheckerResult.error(message.toString()); + return RebalancePreCheckerResult.error( + getUnsafeDiskUtilizationMessage("during rebalance", serversUnsafeDuringRebalance, threshold) + + ". Enable lowDiskMode to delete segments before adding the new ones"); + } + + private static void addIfOverThreshold(List servers, String server, double utilizationRatio, + double threshold) { + if (utilizationRatio >= threshold) { + servers.add(server + String.format(" (%d%%)", (short) (utilizationRatio * 100))); + } + } + + private static String getUnsafeDiskUtilizationMessage(String when, List servers, double threshold) { + return String.format("UNSAFE. Servers with unsafe disk utilization %s (>%d%%): %s", when, (short) (threshold * 100), + String.join(", ", servers)); } private RebalancePreCheckerResult checkRebalanceConfig(RebalanceConfig rebalanceConfig, TableConfig tableConfig, diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java new file mode 100644 index 000000000000..59f12254bf23 --- /dev/null +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java @@ -0,0 +1,164 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.helix.core.rebalance; + +import java.util.HashMap; +import java.util.Map; +import org.apache.pinot.common.restlet.resources.DiskUsageInfo; +import org.apache.pinot.common.restlet.resources.RebalanceConfig; +import org.apache.pinot.common.restlet.resources.RebalancePreCheckerResult; +import org.apache.pinot.common.restlet.resources.RebalancePreCheckerResult.PreCheckStatus; +import org.apache.pinot.controller.helix.core.rebalance.RebalancePreChecker.PreCheckContext; +import org.apache.pinot.controller.util.TableSizeReader; +import org.apache.pinot.controller.validation.ResourceUtilizationInfo; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.AfterClass; +import org.testng.annotations.Test; + +import static org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Tests the consolidated disk utilization pre-check of [DefaultRebalancePreChecker]. +/// +/// All the tests share the same assignment: 4 segments of 100 bytes each, all on `Server_0`, out of which `segment_2` +/// and `segment_3` move to `Server_1`. Every server has 1000 bytes of total space and the threshold is 50% of it. +public class DefaultRebalancePreCheckerTest { + private static final String SERVER_0 = "Server_0"; + private static final String SERVER_1 = "Server_1"; + private static final int NUM_SEGMENTS = 4; + private static final long TOTAL_SPACE_BYTES = 1000L; + private static final long TABLE_SIZE_PER_REPLICA_BYTES = 400L; + private static final double THRESHOLD = 0.5; + + private final DefaultRebalancePreChecker _preChecker = new DefaultRebalancePreChecker(); + + /// [ResourceUtilizationInfo] is a mutable static shared by everything running in the same JVM fork. + @AfterClass + public void tearDown() { + ResourceUtilizationInfo.setDiskUsageInfo(Map.of()); + } + + @Test + public void testWithinThresholdBothDuringAndAfterRebalance() { + // Server_0 sheds 200 bytes and Server_1 gains them, neither ever goes over 500 bytes + setDiskUsage(400L, 0L); + RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig()); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); + assertEquals(result.getMessage(), "Within threshold (<50%)"); + } + + @Test + public void testOverThresholdAfterRebalanceIsAnErrorWhateverTheRebalanceConfig() { + // Server_1 ends up at 550 of its 1000 bytes, which no rebalance config can bring back under the threshold + setDiskUsage(400L, 350L); + for (RebalanceConfig rebalanceConfig : new RebalanceConfig[]{ + new RebalanceConfig(), lowDiskMode(), downtime() + }) { + RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertEquals(result.getMessage(), + "UNSAFE. Servers with unsafe disk utilization after rebalance (>50%): " + SERVER_1 + " (55%)"); + } + } + + @Test + public void testOverThresholdOnlyDuringRebalanceIsAnErrorWithoutLowDiskMode() { + // Server_0 is at 550 of its 1000 bytes and only gets back under the threshold once it has shed its 200 bytes. + // downtime does not order the drops before the adds, so it does not rule the transient peak out either. + setDiskUsage(550L, 0L); + for (RebalanceConfig rebalanceConfig : new RebalanceConfig[]{new RebalanceConfig(), downtime()}) { + RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertTrue(result.getMessage() + .startsWith("UNSAFE. Servers with unsafe disk utilization during rebalance (>50%): " + SERVER_0 + + " (55%)"), result.getMessage()); + } + } + + @Test + public void testOverThresholdOnlyDuringRebalanceIsSafeWithLowDiskMode() { + setDiskUsage(550L, 0L); + RebalancePreCheckerResult result = checkDiskUtilization(lowDiskMode()); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); + assertTrue(result.getMessage().startsWith("Within threshold (<50%) after rebalance"), result.getMessage()); + } + + @Test + public void testDiskUsageInfoNotAvailable() { + ResourceUtilizationInfo.setDiskUsageInfo(Map.of()); + RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig()); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.WARN); + assertTrue(result.getMessage().startsWith("Disk usage info has not been updated"), result.getMessage()); + } + + private RebalancePreCheckerResult checkDiskUtilization(RebalanceConfig rebalanceConfig) { + return _preChecker.checkDiskUtilization(getPreCheckContext(rebalanceConfig), THRESHOLD); + } + + private static RebalanceConfig lowDiskMode() { + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setLowDiskMode(true); + return rebalanceConfig; + } + + private static RebalanceConfig downtime() { + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDowntime(true); + return rebalanceConfig; + } + + private static void setDiskUsage(long usedSpaceBytesServer0, long usedSpaceBytesServer1) { + long now = System.currentTimeMillis(); + ResourceUtilizationInfo.setDiskUsageInfo( + Map.of(SERVER_0, new DiskUsageInfo(SERVER_0, "", TOTAL_SPACE_BYTES, usedSpaceBytesServer0, now), SERVER_1, + new DiskUsageInfo(SERVER_1, "", TOTAL_SPACE_BYTES, usedSpaceBytesServer1, now))); + } + + private static PreCheckContext getPreCheckContext(RebalanceConfig rebalanceConfig) { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build(); + return new PreCheckContext("jobId", tableConfig.getTableName(), tableConfig, getCurrentAssignment(), + getTargetAssignment(), getTableSizeDetails(), rebalanceConfig, null, null); + } + + private static Map> getCurrentAssignment() { + Map> currentAssignment = new HashMap<>(); + for (int i = 0; i < NUM_SEGMENTS; i++) { + currentAssignment.put("segment_" + i, Map.of(SERVER_0, ONLINE)); + } + return currentAssignment; + } + + private static Map> getTargetAssignment() { + Map> targetAssignment = new HashMap<>(); + for (int i = 0; i < NUM_SEGMENTS; i++) { + targetAssignment.put("segment_" + i, Map.of(i < NUM_SEGMENTS / 2 ? SERVER_0 : SERVER_1, ONLINE)); + } + return targetAssignment; + } + + private static TableSizeReader.TableSubTypeSizeDetails getTableSizeDetails() { + TableSizeReader.TableSubTypeSizeDetails tableSubTypeSizeDetails = new TableSizeReader.TableSubTypeSizeDetails(); + tableSubTypeSizeDetails._reportedSizePerReplicaInBytes = TABLE_SIZE_PER_REPLICA_BYTES; + return tableSubTypeSizeDetails; + } +} diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java index b1f5c6e9debb..c96e29a6d1f2 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java @@ -277,11 +277,10 @@ public void testRebalance() assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); Map preCheckResult = rebalanceResult.getPreChecksResult(); assertNotNull(preCheckResult); - assertEquals(preCheckResult.size(), 6); + assertEquals(preCheckResult.size(), 5); assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.NEEDS_RELOAD_STATUS)); assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT)); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE)); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE)); + assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION)); assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.REBALANCE_CONFIG_OPTIONS)); assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.REPLICA_GROUPS_INFO)); // Sending request to servers should fail for all, so needsPreprocess should be set to "error" to indicate that a @@ -294,14 +293,9 @@ public void testRebalance() RebalancePreCheckerResult.PreCheckStatus.PASS); assertEquals(preCheckResult.get(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT).getMessage(), "Instance assignment not allowed, no need for minimizeDataMovement"); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE).getPreCheckStatus(), + assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION).getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.PASS); - assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE) - .getMessage() - .startsWith("Within threshold")); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE).getPreCheckStatus(), - RebalancePreCheckerResult.PreCheckStatus.PASS); - assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE) + assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION) .getMessage() .startsWith("Within threshold")); assertEquals(preCheckResult.get(DefaultRebalancePreChecker.REBALANCE_CONFIG_OPTIONS).getPreCheckStatus(), @@ -1172,16 +1166,10 @@ public void testRebalancePreCheckerDiskUtil() assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); Map preCheckResult = rebalanceResult.getPreChecksResult(); assertNotNull(preCheckResult); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE)); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE).getPreCheckStatus(), - RebalancePreCheckerResult.PreCheckStatus.PASS); - assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE) - .getMessage() - .startsWith("Within threshold")); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE)); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE).getPreCheckStatus(), + assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION)); + assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION).getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.PASS); - assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE) + assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION) .getMessage() .startsWith("Within threshold")); @@ -1199,11 +1187,24 @@ public void testRebalancePreCheckerDiskUtil() assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); preCheckResult = rebalanceResult.getPreChecksResult(); assertNotNull(preCheckResult); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE)); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE).getPreCheckStatus(), + assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION)); + assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION).getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.ERROR); - assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE)); - assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE).getPreCheckStatus(), + assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION) + .getMessage() + .startsWith("UNSAFE. Servers with unsafe disk utilization after rebalance")); + + // The servers are over the threshold no matter how the rebalance is run, so lowDiskMode does not help + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setPreChecks(true); + rebalanceConfig.setLowDiskMode(true); + + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + preCheckResult = rebalanceResult.getPreChecksResult(); + assertNotNull(preCheckResult); + assertEquals(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION).getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.ERROR); _helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TableRebalanceIntegrationTest.java index 30eceaba1c3f..d7e7713aaf22 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TableRebalanceIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TableRebalanceIntegrationTest.java @@ -931,11 +931,10 @@ private void checkRebalancePreCheckStatus(RebalanceResult rebalanceResult, Rebal assertEquals(rebalanceResult.getStatus(), expectedStatus); Map preChecksResult = rebalanceResult.getPreChecksResult(); assertNotNull(preChecksResult); - assertEquals(preChecksResult.size(), 6); + assertEquals(preChecksResult.size(), 5); assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT)); assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.NEEDS_RELOAD_STATUS)); - assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE)); - assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE)); + assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION)); assertTrue(preChecksResult.containsKey(DefaultRebalancePreChecker.REBALANCE_CONFIG_OPTIONS)); assertEquals(preChecksResult.get(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT).getPreCheckStatus(), expectedMinimizeDataMovementStatus); @@ -957,9 +956,7 @@ private void checkRebalancePreCheckStatus(RebalanceResult rebalanceResult, Rebal // .RESOURCE_UTILIZATION_CHECKER_INITIAL_DELAY was set to 30000s, see org.apache.pinot.controller.helix // .ControllerTest.getDefaultControllerConfiguration), server's disk util should be unavailable on all servers if // not explicitly set via org.apache.pinot.controller.validation.ResourceUtilizationInfo.setDiskUsageInfo - assertEquals(preChecksResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_DURING_REBALANCE).getPreCheckStatus(), - RebalancePreCheckerResult.PreCheckStatus.WARN); - assertEquals(preChecksResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION_AFTER_REBALANCE).getPreCheckStatus(), + assertEquals(preChecksResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION).getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); } From 2de0216c81dc9dd5c152990edaa43afa7db36c6a Mon Sep 17 00:00:00 2001 From: J-HowHuang Date: Thu, 6 Aug 2026 16:41:10 -0700 Subject: [PATCH 2/5] patch downtime logic --- .../rebalance/DefaultRebalancePreChecker.java | 27 ++++++++++++++----- .../DefaultRebalancePreCheckerTest.java | 16 +++++++++++ .../TableRebalancerClusterStatelessTest.java | 11 ++++++++ 3 files changed, 47 insertions(+), 7 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index cafc3888359d..ece9c50df6df 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -273,8 +273,10 @@ private RebalancePreCheckerResult checkIsMinimizeDataMovement(TableConfig tableC /// brings the end state back within the threshold. /// - **During the rebalance**, where a server can transiently hold the segments it is gaining on top of the ones it /// is about to lose. Only `lowDiskMode` rules that peak out, by waiting for the segments to be deleted before - /// adding the new ones, so going over the threshold there is an error unless it is enabled. Note that `downtime` - /// does not help: it hands every segment to Helix at once without ordering the drops before the adds. + /// adding the new ones, so going over the threshold there is an error unless it is enabled. `downtime` does not + /// help: it replaces the IdealState with the target assignment in one go, without ordering the drops before the + /// adds. Worse, that one-shot path skips the incremental one `lowDiskMode` acts on, so `downtime` cancels + /// `lowDiskMode` out entirely. /// /// Every segment is assumed to take up disk space on each server it is assigned to. Downstream projects where that /// does not hold (e.g. because a segment can be stored outside of the server, as indicated by @@ -350,15 +352,19 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec return RebalancePreCheckerResult.pass(withinThreshold); } // lowDiskMode is the only way to rule the transient disk usage above out, since it waits for the segments to be - // deleted before adding the new ones. downtime does not: it moves every segment in one shot without ordering the - // drops before the adds, so a server can still end up holding both at once - if (preCheckContext.getRebalanceConfig().isLowDiskMode()) { + // deleted before adding the new ones. It is however only honored by the incremental rebalance path, which downtime + // skips altogether by replacing the IdealState with the target assignment in one go + RebalanceConfig rebalanceConfig = preCheckContext.getRebalanceConfig(); + if (rebalanceConfig.isLowDiskMode() && !rebalanceConfig.isDowntime()) { return RebalancePreCheckerResult.pass(withinThreshold + " after rebalance. Some servers would go over it during " + "the rebalance, but lowDiskMode avoids that transient disk usage"); } return RebalancePreCheckerResult.error( - getUnsafeDiskUtilizationMessage("during rebalance", serversUnsafeDuringRebalance, threshold) - + ". Enable lowDiskMode to delete segments before adding the new ones"); + getUnsafeDiskUtilizationMessage("during rebalance", serversUnsafeDuringRebalance, threshold) + ( + rebalanceConfig.isLowDiskMode() + ? ". lowDiskMode has no effect while downtime is enabled, disable downtime for it to delete segments " + + "before adding the new ones" + : ". Enable lowDiskMode to delete segments before adding the new ones")); } private static void addIfOverThreshold(List servers, String server, double utilizationRatio, @@ -400,6 +406,13 @@ private RebalancePreCheckerResult checkRebalanceConfig(RebalanceConfig rebalance pass = false; warnings.add("Number of replicas (" + numReplicas + ") is greater than 1, downtime is not recommended."); } + // Downtime replaces the IdealState with the target assignment in one go, skipping the incremental path that is + // the only one honoring lowDiskMode + if (rebalanceConfig.isLowDiskMode()) { + pass = false; + warnings.add("lowDiskMode has no effect when downtime is enabled, disable downtime for segments to be deleted " + + "before the new ones are added."); + } } // Peer download enabled tables may have data loss during rebalance, when downtime=true or minAvailableReplicas=0. diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java index 59f12254bf23..03857ca79243 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java @@ -103,6 +103,22 @@ public void testOverThresholdOnlyDuringRebalanceIsSafeWithLowDiskMode() { assertTrue(result.getMessage().startsWith("Within threshold (<50%) after rebalance"), result.getMessage()); } + @Test + public void testDowntimeCancelsLowDiskModeOut() { + // Downtime replaces the IdealState with the target assignment in one go, skipping the incremental path that is the + // only one honoring lowDiskMode, so the transient peak stands and the message has to say so + setDiskUsage(550L, 0L); + RebalanceConfig rebalanceConfig = lowDiskMode(); + rebalanceConfig.setDowntime(true); + + RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertEquals(result.getMessage(), + "UNSAFE. Servers with unsafe disk utilization during rebalance (>50%): " + SERVER_0 + " (55%). lowDiskMode has " + + "no effect while downtime is enabled, disable downtime for it to delete segments before adding the new " + + "ones"); + } + @Test public void testDiskUsageInfoNotAvailable() { ResourceUtilizationInfo.setDiskUsageInfo(Map.of()); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java index c96e29a6d1f2..2a61c09677a7 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java @@ -1318,6 +1318,17 @@ public void testRebalancePreCheckerRebalanceConfig() assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.PASS); assertEquals(preCheckerResult.getMessage(), "All rebalance parameters look good"); + // trigger the warning about lowDiskMode being inert with downtime, which replaces the IdealState in one go + rebalanceConfig.setLowDiskMode(true); + rebalanceResult = tableRebalancer.rebalance(newTableConfig, rebalanceConfig, null); + preCheckerResult = rebalanceResult.getPreChecksResult().get(DefaultRebalancePreChecker.REBALANCE_CONFIG_OPTIONS); + assertNotNull(preCheckerResult); + assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); + assertEquals(preCheckerResult.getMessage(), + "lowDiskMode has no effect when downtime is enabled, disable downtime for segments to be deleted before the " + + "new ones are added."); + rebalanceConfig.setLowDiskMode(false); + // trigger peer-download enabled table rebalance warning newTableConfig.getValidationConfig().setPeerSegmentDownloadScheme("http"); From a25d60a9356d43503b9b092e9f7dd51ffcfc7665 Mon Sep 17 00:00:00 2001 From: J-HowHuang Date: Thu, 6 Aug 2026 16:53:56 -0700 Subject: [PATCH 3/5] capitalize keywords --- .../helix/core/rebalance/DefaultRebalancePreChecker.java | 6 +++--- .../core/rebalance/DefaultRebalancePreCheckerTest.java | 8 ++++---- .../rebalance/TableRebalancerClusterStatelessTest.java | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index ece9c50df6df..ce31c8609293 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -345,7 +345,7 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec // rather than by tuning the rebalance config if (!serversUnsafeAfterRebalance.isEmpty()) { return RebalancePreCheckerResult.error( - getUnsafeDiskUtilizationMessage("after rebalance", serversUnsafeAfterRebalance, threshold)); + getUnsafeDiskUtilizationMessage("AFTER rebalance", serversUnsafeAfterRebalance, threshold)); } String withinThreshold = String.format("Within threshold (<%d%%)", (short) (threshold * 100)); if (serversUnsafeDuringRebalance.isEmpty()) { @@ -356,11 +356,11 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec // skips altogether by replacing the IdealState with the target assignment in one go RebalanceConfig rebalanceConfig = preCheckContext.getRebalanceConfig(); if (rebalanceConfig.isLowDiskMode() && !rebalanceConfig.isDowntime()) { - return RebalancePreCheckerResult.pass(withinThreshold + " after rebalance. Some servers would go over it during " + return RebalancePreCheckerResult.pass(withinThreshold + " AFTER rebalance. Some servers would go over it DURING " + "the rebalance, but lowDiskMode avoids that transient disk usage"); } return RebalancePreCheckerResult.error( - getUnsafeDiskUtilizationMessage("during rebalance", serversUnsafeDuringRebalance, threshold) + ( + getUnsafeDiskUtilizationMessage("DURING rebalance", serversUnsafeDuringRebalance, threshold) + ( rebalanceConfig.isLowDiskMode() ? ". lowDiskMode has no effect while downtime is enabled, disable downtime for it to delete segments " + "before adding the new ones" diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java index 03857ca79243..3882848dba63 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java @@ -77,7 +77,7 @@ public void testOverThresholdAfterRebalanceIsAnErrorWhateverTheRebalanceConfig() RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); assertEquals(result.getMessage(), - "UNSAFE. Servers with unsafe disk utilization after rebalance (>50%): " + SERVER_1 + " (55%)"); + "UNSAFE. Servers with unsafe disk utilization AFTER rebalance (>50%): " + SERVER_1 + " (55%)"); } } @@ -90,7 +90,7 @@ public void testOverThresholdOnlyDuringRebalanceIsAnErrorWithoutLowDiskMode() { RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); assertTrue(result.getMessage() - .startsWith("UNSAFE. Servers with unsafe disk utilization during rebalance (>50%): " + SERVER_0 + .startsWith("UNSAFE. Servers with unsafe disk utilization DURING rebalance (>50%): " + SERVER_0 + " (55%)"), result.getMessage()); } } @@ -100,7 +100,7 @@ public void testOverThresholdOnlyDuringRebalanceIsSafeWithLowDiskMode() { setDiskUsage(550L, 0L); RebalancePreCheckerResult result = checkDiskUtilization(lowDiskMode()); assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); - assertTrue(result.getMessage().startsWith("Within threshold (<50%) after rebalance"), result.getMessage()); + assertTrue(result.getMessage().startsWith("Within threshold (<50%) AFTER rebalance"), result.getMessage()); } @Test @@ -114,7 +114,7 @@ public void testDowntimeCancelsLowDiskModeOut() { RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); assertEquals(result.getMessage(), - "UNSAFE. Servers with unsafe disk utilization during rebalance (>50%): " + SERVER_0 + " (55%). lowDiskMode has " + "UNSAFE. Servers with unsafe disk utilization DURING rebalance (>50%): " + SERVER_0 + " (55%). lowDiskMode has " + "no effect while downtime is enabled, disable downtime for it to delete segments before adding the new " + "ones"); } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java index 2a61c09677a7..d010aa7ccb6f 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java @@ -1192,7 +1192,7 @@ public void testRebalancePreCheckerDiskUtil() RebalancePreCheckerResult.PreCheckStatus.ERROR); assertTrue(preCheckResult.get(DefaultRebalancePreChecker.DISK_UTILIZATION) .getMessage() - .startsWith("UNSAFE. Servers with unsafe disk utilization after rebalance")); + .startsWith("UNSAFE. Servers with unsafe disk utilization AFTER rebalance")); // The servers are over the threshold no matter how the rebalance is run, so lowDiskMode does not help rebalanceConfig = new RebalanceConfig(); From c0af72e1dabefd5102524bbe378391442eaaed72 Mon Sep 17 00:00:00 2001 From: J-HowHuang Date: Thu, 13 Aug 2026 14:33:54 -0700 Subject: [PATCH 4/5] address PR comments --- .../rebalance/DefaultRebalancePreChecker.java | 50 ++++--- .../DefaultRebalancePreCheckerTest.java | 134 +++++++++++------- 2 files changed, 116 insertions(+), 68 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index ce31c8609293..e99906b6ebb6 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -272,11 +272,14 @@ private RebalancePreCheckerResult checkIsMinimizeDataMovement(TableConfig tableC /// over the threshold there is an error whatever the rebalance config, since no way of running the rebalance /// brings the end state back within the threshold. /// - **During the rebalance**, where a server can transiently hold the segments it is gaining on top of the ones it - /// is about to lose. Only `lowDiskMode` rules that peak out, by waiting for the segments to be deleted before - /// adding the new ones, so going over the threshold there is an error unless it is enabled. `downtime` does not - /// help: it replaces the IdealState with the target assignment in one go, without ordering the drops before the - /// adds. Worse, that one-shot path skips the incremental one `lowDiskMode` acts on, so `downtime` cancels - /// `lowDiskMode` out entirely. + /// is about to lose. Only servers actually gaining segments are estimated: one that merely sheds them can only be + /// over the threshold because it already was, which the rebalance neither causes nor can be blamed for. Only + /// `lowDiskMode` rules that peak out, by waiting for the segments to be deleted before adding the new ones, so + /// going over the threshold there is an error unless it is enabled. `downtime` does not help: it replaces the + /// IdealState with the target assignment in one go, without ordering the drops before the adds. Worse, that + /// one-shot path skips the incremental one `lowDiskMode` acts on, so `downtime` cancels `lowDiskMode` out + /// entirely. `bestEfforts` weakens it rather than cancelling it — the deletes are still awaited, just no longer + /// unconditionally — so it downgrades the result to a warning instead of an error. /// /// Every segment is assumed to take up disk space on each server it is assigned to. Downstream projects where that /// does not hold (e.g. because a segment can be stored outside of the server, as indicated by @@ -332,9 +335,14 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec long diskUtilizationGain = newSegmentSet.size() * avgSegmentSize; long diskUtilizationLoss = removedSegmentSet.size() * avgSegmentSize; - // While the rebalance is running, the segments being added can co-exist with the ones being removed - addIfOverThreshold(serversUnsafeDuringRebalance, server, - (double) (diskUsage.getUsedSpaceBytes() + diskUtilizationGain) / diskUsage.getTotalSpaceBytes(), threshold); + // While the rebalance is running, the segments being added can co-exist with the ones being removed. A server + // gaining nothing never builds up that transient usage: it is only ever over the threshold because it already + // was, which is not something this rebalance causes nor something lowDiskMode could do anything about. If it is + // still over once the rebalance is done, the estimate below catches it + if (diskUtilizationGain > 0) { + addIfOverThreshold(serversUnsafeDuringRebalance, server, + (double) (diskUsage.getUsedSpaceBytes() + diskUtilizationGain) / diskUsage.getTotalSpaceBytes(), threshold); + } addIfOverThreshold(serversUnsafeAfterRebalance, server, (double) (diskUsage.getUsedSpaceBytes() + diskUtilizationGain - diskUtilizationLoss) / diskUsage.getTotalSpaceBytes(), threshold); @@ -355,16 +363,18 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec // deleted before adding the new ones. It is however only honored by the incremental rebalance path, which downtime // skips altogether by replacing the IdealState with the target assignment in one go RebalanceConfig rebalanceConfig = preCheckContext.getRebalanceConfig(); - if (rebalanceConfig.isLowDiskMode() && !rebalanceConfig.isDowntime()) { - return RebalancePreCheckerResult.pass(withinThreshold + " AFTER rebalance. Some servers would go over it DURING " - + "the rebalance, but lowDiskMode avoids that transient disk usage"); + if (rebalanceConfig.isDowntime() || !rebalanceConfig.isLowDiskMode()) { + return RebalancePreCheckerResult.error( + getUnsafeDiskUtilizationMessage("DURING rebalance", serversUnsafeDuringRebalance, threshold) + + (rebalanceConfig.isDowntime() + ? ". lowDiskMode, which would delete segments before adding the new ones, has no effect while downtime " + + "is enabled" + : ". Enable lowDiskMode to delete segments before adding the new ones")); } - return RebalancePreCheckerResult.error( - getUnsafeDiskUtilizationMessage("DURING rebalance", serversUnsafeDuringRebalance, threshold) + ( - rebalanceConfig.isLowDiskMode() - ? ". lowDiskMode has no effect while downtime is enabled, disable downtime for it to delete segments " - + "before adding the new ones" - : ". Enable lowDiskMode to delete segments before adding the new ones")); + String serversGoingOver = " Servers that would go over it DURING the rebalance: " + String.join(", ", + serversUnsafeDuringRebalance) + "."; + return RebalancePreCheckerResult.pass(withinThreshold + " AFTER rebalance." + serversGoingOver + " lowDiskMode " + + "avoids that transient disk usage by deleting segments before adding the new ones"); } private static void addIfOverThreshold(List servers, String server, double utilizationRatio, @@ -374,9 +384,11 @@ private static void addIfOverThreshold(List servers, String server, doub } } + /// The threshold is rendered as `>=` because [#addIfOverThreshold] flags a server whose utilization reaches it, not + /// only one that exceeds it. private static String getUnsafeDiskUtilizationMessage(String when, List servers, double threshold) { - return String.format("UNSAFE. Servers with unsafe disk utilization %s (>%d%%): %s", when, (short) (threshold * 100), - String.join(", ", servers)); + return String.format("UNSAFE. Servers with unsafe disk utilization %s (>=%d%%): %s", when, + (short) (threshold * 100), String.join(", ", servers)); } private RebalancePreCheckerResult checkRebalanceConfig(RebalanceConfig rebalanceConfig, TableConfig tableConfig, diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java index 3882848dba63..ecf5583965a2 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreCheckerTest.java @@ -18,7 +18,6 @@ */ package org.apache.pinot.controller.helix.core.rebalance; -import java.util.HashMap; import java.util.Map; import org.apache.pinot.common.restlet.resources.DiskUsageInfo; import org.apache.pinot.common.restlet.resources.RebalanceConfig; @@ -40,16 +39,32 @@ /// Tests the consolidated disk utilization pre-check of [DefaultRebalancePreChecker]. /// -/// All the tests share the same assignment: 4 segments of 100 bytes each, all on `Server_0`, out of which `segment_2` -/// and `segment_3` move to `Server_1`. Every server has 1000 bytes of total space and the threshold is 50% of it. +/// Unless stated otherwise, the tests share the same assignment: 4 segments of 100 bytes each, moving so that both +/// servers gain and shed some. `Server_0` gains 100 bytes and sheds 200, `Server_1` gains 200 and sheds 100. Every +/// server has 1000 bytes of total space and the threshold is 50% of it. public class DefaultRebalancePreCheckerTest { private static final String SERVER_0 = "Server_0"; private static final String SERVER_1 = "Server_1"; - private static final int NUM_SEGMENTS = 4; private static final long TOTAL_SPACE_BYTES = 1000L; private static final long TABLE_SIZE_PER_REPLICA_BYTES = 400L; private static final double THRESHOLD = 0.5; + private static final Map> CURRENT_ASSIGNMENT = + Map.of("segment_0", Map.of(SERVER_0, ONLINE), "segment_1", Map.of(SERVER_0, ONLINE), "segment_2", + Map.of(SERVER_0, ONLINE), "segment_3", Map.of(SERVER_1, ONLINE)); + private static final Map> TARGET_ASSIGNMENT = + Map.of("segment_0", Map.of(SERVER_0, ONLINE), "segment_1", Map.of(SERVER_1, ONLINE), "segment_2", + Map.of(SERVER_1, ONLINE), "segment_3", Map.of(SERVER_0, ONLINE)); + + /// An assignment where `Server_0` only sheds segments: it gives its 2 remaining segments to `Server_1` and gains + /// nothing. + private static final Map> SHED_ONLY_CURRENT_ASSIGNMENT = + Map.of("segment_0", Map.of(SERVER_0, ONLINE), "segment_1", Map.of(SERVER_0, ONLINE), "segment_2", + Map.of(SERVER_0, ONLINE), "segment_3", Map.of(SERVER_0, ONLINE)); + private static final Map> SHED_ONLY_TARGET_ASSIGNMENT = + Map.of("segment_0", Map.of(SERVER_0, ONLINE), "segment_1", Map.of(SERVER_0, ONLINE), "segment_2", + Map.of(SERVER_1, ONLINE), "segment_3", Map.of(SERVER_1, ONLINE)); + private final DefaultRebalancePreChecker _preChecker = new DefaultRebalancePreChecker(); /// [ResourceUtilizationInfo] is a mutable static shared by everything running in the same JVM fork. @@ -60,8 +75,8 @@ public void tearDown() { @Test public void testWithinThresholdBothDuringAndAfterRebalance() { - // Server_0 sheds 200 bytes and Server_1 gains them, neither ever goes over 500 bytes - setDiskUsage(400L, 0L); + // Server_0 peaks at 400 of its 1000 bytes and Server_1 at 300, neither reaches 500 + setDiskUsage(300L, 100L); RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig()); assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); assertEquals(result.getMessage(), "Within threshold (<50%)"); @@ -69,54 +84,78 @@ public void testWithinThresholdBothDuringAndAfterRebalance() { @Test public void testOverThresholdAfterRebalanceIsAnErrorWhateverTheRebalanceConfig() { - // Server_1 ends up at 550 of its 1000 bytes, which no rebalance config can bring back under the threshold - setDiskUsage(400L, 350L); + // Server_1 ends up at 520 of its 1000 bytes, which no rebalance config can bring back under the threshold + setDiskUsage(100L, 420L); for (RebalanceConfig rebalanceConfig : new RebalanceConfig[]{ - new RebalanceConfig(), lowDiskMode(), downtime() + new RebalanceConfig(), lowDiskMode(), downtime(), bestEfforts() }) { RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); assertEquals(result.getMessage(), - "UNSAFE. Servers with unsafe disk utilization AFTER rebalance (>50%): " + SERVER_1 + " (55%)"); + "UNSAFE. Servers with unsafe disk utilization AFTER rebalance (>=50%): " + SERVER_1 + " (52%)"); } } @Test public void testOverThresholdOnlyDuringRebalanceIsAnErrorWithoutLowDiskMode() { - // Server_0 is at 550 of its 1000 bytes and only gets back under the threshold once it has shed its 200 bytes. - // downtime does not order the drops before the adds, so it does not rule the transient peak out either. - setDiskUsage(550L, 0L); - for (RebalanceConfig rebalanceConfig : new RebalanceConfig[]{new RebalanceConfig(), downtime()}) { - RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); - assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); - assertTrue(result.getMessage() - .startsWith("UNSAFE. Servers with unsafe disk utilization DURING rebalance (>50%): " + SERVER_0 - + " (55%)"), result.getMessage()); - } + // Server_1 transiently holds the 200 bytes it gains on top of the 100 it is about to shed, peaking at 520 + setDiskUsage(100L, 320L); + RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig()); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertEquals(result.getMessage(), + "UNSAFE. Servers with unsafe disk utilization DURING rebalance (>=50%): " + SERVER_1 + " (52%). Enable " + + "lowDiskMode to delete segments before adding the new ones"); } @Test public void testOverThresholdOnlyDuringRebalanceIsSafeWithLowDiskMode() { - setDiskUsage(550L, 0L); + setDiskUsage(100L, 320L); RebalancePreCheckerResult result = checkDiskUtilization(lowDiskMode()); assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); - assertTrue(result.getMessage().startsWith("Within threshold (<50%) AFTER rebalance"), result.getMessage()); + assertEquals(result.getMessage(), + "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"); } @Test public void testDowntimeCancelsLowDiskModeOut() { // Downtime replaces the IdealState with the target assignment in one go, skipping the incremental path that is the - // only one honoring lowDiskMode, so the transient peak stands and the message has to say so + // only one honoring lowDiskMode, so the transient peak stands whether or not lowDiskMode is set + setDiskUsage(100L, 320L); + RebalanceConfig lowDiskModeAndDowntime = lowDiskMode(); + lowDiskModeAndDowntime.setDowntime(true); + + for (RebalanceConfig rebalanceConfig : new RebalanceConfig[]{downtime(), lowDiskModeAndDowntime}) { + RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertEquals(result.getMessage(), + "UNSAFE. Servers with unsafe disk utilization DURING rebalance (>=50%): " + SERVER_1 + " (52%). lowDiskMode, " + + "which would delete segments before adding the new ones, has no effect while downtime is enabled"); + } + } + + @Test + public void testServerGainingNothingIsNotFlaggedDuringRebalance() { + // Server_0 is already at 550 of its 1000 bytes and only sheds segments. The rebalance cannot push it any higher + // and lowDiskMode would have nothing to delete first, so flagging it would blame the rebalance for a pre-existing + // condition. It drops to 350 once done, so there is nothing to report at all. setDiskUsage(550L, 0L); - RebalanceConfig rebalanceConfig = lowDiskMode(); - rebalanceConfig.setDowntime(true); + RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig(), SHED_ONLY_CURRENT_ASSIGNMENT, + SHED_ONLY_TARGET_ASSIGNMENT); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.PASS); + assertEquals(result.getMessage(), "Within threshold (<50%)"); + } - RebalancePreCheckerResult result = checkDiskUtilization(rebalanceConfig); + @Test + public void testServerGainingNothingAndStayingOverThresholdIsStillFlagged() { + // Server_0 shedding its 2 segments is not enough to bring it back under the threshold, which the AFTER estimate + // catches even though the DURING one skips it + setDiskUsage(750L, 0L); + RebalancePreCheckerResult result = checkDiskUtilization(new RebalanceConfig(), SHED_ONLY_CURRENT_ASSIGNMENT, + SHED_ONLY_TARGET_ASSIGNMENT); assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); assertEquals(result.getMessage(), - "UNSAFE. Servers with unsafe disk utilization DURING rebalance (>50%): " + SERVER_0 + " (55%). lowDiskMode has " - + "no effect while downtime is enabled, disable downtime for it to delete segments before adding the new " - + "ones"); + "UNSAFE. Servers with unsafe disk utilization AFTER rebalance (>=50%): " + SERVER_0 + " (55%)"); } @Test @@ -128,7 +167,13 @@ public void testDiskUsageInfoNotAvailable() { } private RebalancePreCheckerResult checkDiskUtilization(RebalanceConfig rebalanceConfig) { - return _preChecker.checkDiskUtilization(getPreCheckContext(rebalanceConfig), THRESHOLD); + return checkDiskUtilization(rebalanceConfig, CURRENT_ASSIGNMENT, TARGET_ASSIGNMENT); + } + + private RebalancePreCheckerResult checkDiskUtilization(RebalanceConfig rebalanceConfig, + Map> currentAssignment, Map> targetAssignment) { + return _preChecker.checkDiskUtilization(getPreCheckContext(rebalanceConfig, currentAssignment, targetAssignment), + THRESHOLD); } private static RebalanceConfig lowDiskMode() { @@ -143,6 +188,12 @@ private static RebalanceConfig downtime() { return rebalanceConfig; } + private static RebalanceConfig bestEfforts() { + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setBestEfforts(true); + return rebalanceConfig; + } + private static void setDiskUsage(long usedSpaceBytesServer0, long usedSpaceBytesServer1) { long now = System.currentTimeMillis(); ResourceUtilizationInfo.setDiskUsageInfo( @@ -150,26 +201,11 @@ private static void setDiskUsage(long usedSpaceBytesServer0, long usedSpaceBytes new DiskUsageInfo(SERVER_1, "", TOTAL_SPACE_BYTES, usedSpaceBytesServer1, now))); } - private static PreCheckContext getPreCheckContext(RebalanceConfig rebalanceConfig) { + private static PreCheckContext getPreCheckContext(RebalanceConfig rebalanceConfig, + Map> currentAssignment, Map> targetAssignment) { TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build(); - return new PreCheckContext("jobId", tableConfig.getTableName(), tableConfig, getCurrentAssignment(), - getTargetAssignment(), getTableSizeDetails(), rebalanceConfig, null, null); - } - - private static Map> getCurrentAssignment() { - Map> currentAssignment = new HashMap<>(); - for (int i = 0; i < NUM_SEGMENTS; i++) { - currentAssignment.put("segment_" + i, Map.of(SERVER_0, ONLINE)); - } - return currentAssignment; - } - - private static Map> getTargetAssignment() { - Map> targetAssignment = new HashMap<>(); - for (int i = 0; i < NUM_SEGMENTS; i++) { - targetAssignment.put("segment_" + i, Map.of(i < NUM_SEGMENTS / 2 ? SERVER_0 : SERVER_1, ONLINE)); - } - return targetAssignment; + return new PreCheckContext("jobId", tableConfig.getTableName(), tableConfig, currentAssignment, targetAssignment, + getTableSizeDetails(), rebalanceConfig, null, null); } private static TableSizeReader.TableSubTypeSizeDetails getTableSizeDetails() { From f2c9602538b1ce0297595239cfb9845df77ccec0 Mon Sep 17 00:00:00 2001 From: J-HowHuang Date: Thu, 13 Aug 2026 14:41:21 -0700 Subject: [PATCH 5/5] empty dummy commit