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 905690e1ed8d..7690e288c8ff 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 @@ -17,6 +17,7 @@ * under the License. */ package org.apache.pinot.controller.helix.core.rebalance; +import com.google.common.annotations.VisibleForTesting; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; @@ -25,6 +26,7 @@ import java.util.Map; import java.util.Set; import java.util.concurrent.ExecutorService; +import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils; @@ -36,12 +38,14 @@ import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils; import org.apache.pinot.controller.util.TableSizeReader; import org.apache.pinot.controller.validation.ResourceUtilizationInfo; +import org.apache.pinot.spi.config.table.RoutingConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.table.TierConfig; import org.apache.pinot.spi.config.table.assignment.InstanceAssignmentConfig; import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType; import org.apache.pinot.spi.config.table.assignment.InstanceReplicaGroupPartitionConfig; +import org.apache.pinot.spi.utils.DataSizeUtils; import org.apache.pinot.spi.utils.Enablement; import org.apache.pinot.spi.utils.StringUtil; import org.slf4j.Logger; @@ -329,10 +333,62 @@ protected RebalancePreCheckerResult checkDiskUtilization(PreCheckContext preChec } String serversGoingOver = " Servers that would go over it DURING the rebalance: " + String.join(", ", serversUnsafeDuringRebalance) + "."; + // lowDiskMode bounds the disk each server takes on to what it starts the rebalance with, but when the rebalance + // cannot progress at all within those bounds it relaxes them for a step rather than stalling. Replay the rebalance + // to find out whether it has to, instead of assuming lowDiskMode always avoids the transient usage + Map serversForcedOverBudget = getServersForcedOverDiskBudget(preCheckContext); + if (!serversForcedOverBudget.isEmpty()) { + return RebalancePreCheckerResult.error( + getUnsafeDiskUtilizationMessage("DURING rebalance", serversUnsafeDuringRebalance, threshold) + + ". 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 " + formatBytesOverBudget( + serversForcedOverBudget) + ". Rebalance to a target assignment that frees up space on them first, or " + + "add capacity"); + } return RebalancePreCheckerResult.pass(withinThreshold + " AFTER rebalance." + serversGoingOver + " lowDiskMode " + "avoids that transient disk usage by deleting segments before adding the new ones"); } + /// Replays the rebalance to find the servers that lowDiskMode cannot keep within the disk they start with. Returns + /// an empty map when the replay cannot be run, so that a missing input never turns into a failed pre-check. + /// + /// No reachable rebalance has been found that ends up here with a non-empty result: at the first step a server is + /// only at its ceiling when the target assignment places no more on it than it already hosts, and a rebalance that + /// stalls needs every segment to be waiting to gain a replica, which means the target places more in total than the + /// current assignment does, so at least one server has room. This is kept as a guard rather than as an assertion + /// because the argument does not cover every later step, and going over a server's disk silently is worse than + /// reporting a rebalance as unsafe. + @VisibleForTesting + protected Map getServersForcedOverDiskBudget(PreCheckContext preCheckContext) { + Map> currentAssignment = preCheckContext.getCurrentAssignment(); + Map> targetAssignment = preCheckContext.getTargetAssignment(); + TableConfig tableConfig = preCheckContext.getTableConfig(); + RebalanceConfig rebalanceConfig = preCheckContext.getRebalanceConfig(); + Logger tableRebalanceLogger = LoggerFactory.getLogger(getClass().getSimpleName() + '-' + + preCheckContext.getTableNameWithType() + '-' + preCheckContext.getRebalanceJobId()); + List segmentsToMove = SegmentAssignmentUtils.getSegmentsToMove(currentAssignment, targetAssignment); + int minAvailableReplicas = TableRebalancer.getMinAvailableReplicas(currentAssignment, targetAssignment, + segmentsToMove, rebalanceConfig.getMinAvailableReplicas(), tableRebalanceLogger); + if (minAvailableReplicas == TableRebalancer.ILLEGAL_MIN_AVAILABLE_REPLICAS) { + // The rebalance is going to fail on the config before it moves anything + return Map.of(); + } + boolean enableStrictReplicaGroup = tableConfig.getRoutingConfig() != null + && RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE.equalsIgnoreCase( + tableConfig.getRoutingConfig().getInstanceSelectorType()); + return TableRebalancer.getServersForcedOverDiskBudget(currentAssignment, targetAssignment, minAvailableReplicas, + enableStrictReplicaGroup, rebalanceConfig.getBatchSizePerServer(), + preCheckContext.getTableSubTypeSizeDetails(), tableRebalanceLogger); + } + + /// Renders the servers and how far over their disk they would be pushed, largest first. + private static String formatBytesOverBudget(Map serverToBytesOverBudget) { + return serverToBytesOverBudget.entrySet().stream() + .sorted(Map.Entry.comparingByValue().reversed()) + .map(entry -> entry.getKey() + " (" + DataSizeUtils.fromBytes(entry.getValue()) + ")") + .collect(Collectors.joining(", ")); + } + private static void addIfOverThreshold(List servers, String server, double utilizationRatio, double threshold) { if (utilizationRatio >= threshold) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java index d264fb320315..b4478e3cf420 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java @@ -143,6 +143,14 @@ public class TableRebalancer { private static final int TABLE_SIZE_READER_TIMEOUT_MS = 30_000; private static final int STREAM_PARTITION_OFFSET_READ_TIMEOUT_MS = 10_000; private static final AtomicInteger REBALANCE_JOB_COUNTER = new AtomicInteger(0); + /// Returned by [#getMinAvailableReplicas] when `minReplicasToKeepUpForNoDowntime` is not less than the replication + @VisibleForTesting + static final int ILLEGAL_MIN_AVAILABLE_REPLICAS = -1; + /// Backstop for [#getServersForcedOverDiskBudget]. Every step brings at least one segment replica closer to the + /// target assignment, so the replay terminates well within this on any real assignment + private static final int MAX_DISK_BUDGET_REPLAY_STEPS = 10_000; + /// Used by [#getServersForcedOverDiskBudget], which does not read the segment partition ids + private static final PartitionIdFetcher DEFAULT_PARTITION_ID_FETCHER = segmentName -> 0; private final HelixManager _helixManager; private final HelixDataAccessor _helixDataAccessor; private final TableRebalanceObserver _tableRebalanceObserver; @@ -515,36 +523,15 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb // current instances as this is the best we can do, and can help the table get out of this state. // 2. Only check the segments to be moved because we don't need to maintain available replicas for segments not // being moved, including segments with all replicas OFFLINE (error segments during consumption). - int numReplicas = Integer.MAX_VALUE; - for (String segment : segmentsToMove) { - numReplicas = Math.min(targetAssignment.get(segment).size(), numReplicas); - } - int minAvailableReplicas; - if (minReplicasToKeepUpForNoDowntime >= 0) { - // For non-negative value, use it as min available replicas - if (minReplicasToKeepUpForNoDowntime >= numReplicas) { - onReturnFailure("Illegal config for minReplicasToKeepUpForNoDowntime: " + minReplicasToKeepUpForNoDowntime - + ", must be less than number of replicas: " + numReplicas + ", aborting the rebalance", null, - tableRebalanceLogger); - return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, - "Illegal min available replicas config", instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult, summaryResult); - } - minAvailableReplicas = minReplicasToKeepUpForNoDowntime; - } else { - // For negative value, use it as max unavailable replicas - minAvailableReplicas = Math.max(numReplicas + minReplicasToKeepUpForNoDowntime, 0); - } - - int numCurrentAssignmentReplicas = Integer.MAX_VALUE; - for (String segment : segmentsToMove) { - numCurrentAssignmentReplicas = Math.min(currentAssignment.get(segment).size(), numCurrentAssignmentReplicas); - } - if (minAvailableReplicas > numCurrentAssignmentReplicas) { - tableRebalanceLogger.warn("minAvailableReplicas: {} larger than existing number of replicas: {}, " - + "resetting minAvailableReplicas to {}", minAvailableReplicas, numCurrentAssignmentReplicas, - numCurrentAssignmentReplicas); - minAvailableReplicas = numCurrentAssignmentReplicas; + int minAvailableReplicas = getMinAvailableReplicas(currentAssignment, targetAssignment, segmentsToMove, + minReplicasToKeepUpForNoDowntime, tableRebalanceLogger); + if (minAvailableReplicas == ILLEGAL_MIN_AVAILABLE_REPLICAS) { + onReturnFailure("Illegal config for minReplicasToKeepUpForNoDowntime: " + minReplicasToKeepUpForNoDowntime + + ", must be less than number of replicas: " + getMinNumReplicas(targetAssignment, segmentsToMove) + + ", aborting the rebalance", null, tableRebalanceLogger); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, + "Illegal min available replicas config", instancePartitionsMap, tierToInstancePartitionsMap, + targetAssignment, preChecksResult, summaryResult); } DataLossRiskAssessor dataLossRiskAssessor; @@ -592,6 +579,12 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb // // NOTE: Monitor the segments to be moved from both the previous round and this round to ensure the moved segments // in the previous round are also converged. + + // For low disk mode, capture per-server hosted bytes before we start assignments. It will be used to derive the + // ceiling of how many bytes each server can only add on, per target assignment. + DiskUsageBudget diskUsageBudget = + lowDiskMode ? DiskUsageBudget.create(currentAssignment, tableSubTypeSizeDetails) : null; + List oldSegmentsToMove = segmentsToMove; while (true) { boolean segmentsToMoveChanged = oldSegmentsToMove != segmentsToMove; @@ -710,7 +703,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb nextAssignment = getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, - tableRebalanceLogger); + tableRebalanceLogger, diskUsageBudget); } catch (Exception e) { String errorMsg = "Caught exception while calculating the next assignment, aborting the rebalance: " + e.getMessage(); @@ -775,7 +768,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb nextAssignment = getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, - tableRebalanceLogger); + tableRebalanceLogger, diskUsageBudget); } catch (Exception e) { String errorMsg = "Caught exception while calculating the next assignment, aborting the rebalance: " + e.getMessage(); @@ -1661,8 +1654,23 @@ static Map> getNextAssignment(Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor) { + // Anchor the budget to the assignment this call starts from, which is never above the ceiling anchored at the + // start of the rebalance. Without per-segment sizes it bounds the number of segments hosted instead of the bytes return getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, - lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, LOGGER); + lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, LOGGER, + lowDiskMode ? DiskUsageBudget.create(currentAssignment, null) : null); + } + + /// Uses the default LOGGER, with an explicit disk usage budget + @VisibleForTesting + static Map> getNextAssignment(Map> currentAssignment, + Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, + boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, + PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor, + @Nullable DiskUsageBudget diskUsageBudget) { + return getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, + lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, LOGGER, + diskUsageBudget); } /// Returns the next assignment for the table based on the current assignment and the target assignment with regard to @@ -1687,19 +1695,266 @@ static Map> getNextAssignment(Map> getNextAssignment(Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, - PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) { + PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger, + @Nullable DiskUsageBudget diskUsageBudget) { + if (!lowDiskMode) { + return computeNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, + false, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, + tableRebalanceLogger, null); + } + + // We use DiskUsageBudget and StepDiskBudget here to keep track of per-server byte gain/loss and to guarantee the + // invariant of lowDiskMode where each server won't get more bytes than its net gain at any intermediate steps. + Preconditions.checkState(diskUsageBudget != null, "Low disk mode requires a disk usage budget"); + StepDiskBudget stepBudget = diskUsageBudget.forStep(currentAssignment, targetAssignment); + Map> nextAssignment = + computeNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, true, + batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, tableRebalanceLogger, + stepBudget); + if (!nextAssignment.equals(currentAssignment)) { + return nextAssignment; + } + + // Can't progress while respecting the budget. This could happen if servers circular wait for others to drop first. + tableRebalanceLogger.warn("Cannot make progress in low disk mode without exceeding the disk usage a server " + + "started the rebalance with. Allowing the additions for this step, which temporarily increases the disk " + + "usage of these servers. Bytes each server could still take on: {}", stepBudget.getRemainingBytes()); + return computeNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, + true, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, + tableRebalanceLogger, null); + } + + /// @param stepBudget bounds the disk each server may take on in this step, `null` to apply no bound + private static Map> computeNextAssignment( + Map> currentAssignment, Map> targetAssignment, + int minAvailableReplicas, boolean enableStrictReplicaGroup, boolean lowDiskMode, int batchSizePerServer, + Object2IntOpenHashMap segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher, + DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger, + @Nullable StepDiskBudget stepBudget) { return enableStrictReplicaGroup ? getNextStrictReplicaGroupAssignment(currentAssignment, targetAssignment, minAvailableReplicas, lowDiskMode, - batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, tableRebalanceLogger) + batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, tableRebalanceLogger, + stepBudget) : getNextNonStrictReplicaGroupAssignment(currentAssignment, targetAssignment, minAvailableReplicas, - lowDiskMode, batchSizePerServer, dataLossRiskAssessor); + lowDiskMode, batchSizePerServer, dataLossRiskAssessor, stepBudget); + } + + /// Bounds the disk a server may use while a low disk mode rebalance is in flight. + /// + /// If a server results in net byte loss, its ceiling is as many bytes as it originally hosted. + /// If a server results in net byte gain, its ceiling is as many bytes as it will eventually host. + /// + /// This invariant helps us guarantee that a rebalance won't introduce bytes more than necessary on the disk at any + /// moment during the rebalance. Thus, rebalance can safely run as long as we know the target assignment won't go + /// beyond their disk utilization (this can be guard by either pre-check or resource utilization checker etc.) + /// + /// In the case when new segments 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. + /// + /// When no per-segment size is known every segment counts as one byte, which degenerates to bounding the number of + /// segments hosted rather than the bytes. + @VisibleForTesting + static class DiskUsageBudget { + private final Map _segmentSizeBytes; + private final long _defaultSegmentSizeBytes; + private final Map _anchorServerBytes; + private final Set _accountedSegments; + + private DiskUsageBudget(Map segmentSizeBytes, long defaultSegmentSizeBytes, + Map anchorServerBytes, Set accountedSegments) { + _segmentSizeBytes = segmentSizeBytes; + _defaultSegmentSizeBytes = defaultSegmentSizeBytes; + _anchorServerBytes = anchorServerBytes; + _accountedSegments = accountedSegments; + } + + /// @param initialAssignment the assignment the rebalance starts from, used to anchor the per-server ceiling + /// @param tableSizeDetails per-segment sizes, `null` to fall back to counting segments instead of bytes + static DiskUsageBudget create(Map> initialAssignment, + @Nullable TableSizeReader.TableSubTypeSizeDetails tableSizeDetails) { + Map segmentSizeBytes = new HashMap<>(); + long totalKnownBytes = 0; + if (tableSizeDetails != null) { + for (Map.Entry entry : tableSizeDetails._segments.entrySet()) { + // The size a single replica takes up on disk, which is what one server pays for hosting the segment + long sizeBytes = entry.getValue()._maxReportedSizePerReplicaInBytes; + if (sizeBytes > 0) { + segmentSizeBytes.put(entry.getKey(), sizeBytes); + totalKnownBytes += sizeBytes; + } + } + } + // Segments whose size could not be read (missing from all servers, or still consuming) are charged the average + // size of the segments that could be read. When no per-segment size is available at all, fall back to the + // average over the whole table, which is the same estimate the disk utilization pre-check works with. Failing + // that, charge every segment one byte, which turns the budget into a bound on the number of segments hosted + long defaultSegmentSizeBytes; + if (!segmentSizeBytes.isEmpty()) { + defaultSegmentSizeBytes = Math.max(1L, totalKnownBytes / segmentSizeBytes.size()); + } else if (tableSizeDetails != null && tableSizeDetails._reportedSizePerReplicaInBytes > 0 + && !initialAssignment.isEmpty()) { + defaultSegmentSizeBytes = + Math.max(1L, tableSizeDetails._reportedSizePerReplicaInBytes / initialAssignment.size()); + } else { + defaultSegmentSizeBytes = 1L; + } + Map initialServerBytes = new HashMap<>(); + for (Map.Entry> entry : initialAssignment.entrySet()) { + long sizeBytes = segmentSizeBytes.getOrDefault(entry.getKey(), defaultSegmentSizeBytes); + for (String instance : entry.getValue().keySet()) { + initialServerBytes.merge(instance, sizeBytes, Long::sum); + } + } + return new DiskUsageBudget(segmentSizeBytes, defaultSegmentSizeBytes, initialServerBytes, + new HashSet<>(initialAssignment.keySet())); + } + + long getSegmentSizeBytes(String segmentName) { + return _segmentSizeBytes.getOrDefault(segmentName, _defaultSegmentSizeBytes); + } + + /// Adds the segments that appeared since the anchor was taken - an uploaded segment, or a new consuming segment - + /// to the anchor of the servers hosting them. + /// + /// With this, this segment wouldn't be account for the budget if it doesn't need to be moved at all. + private void accountForNewSegments(Map> currentAssignment) { + for (Map.Entry> entry : currentAssignment.entrySet()) { + if (_accountedSegments.add(entry.getKey())) { + long sizeBytes = getSegmentSizeBytes(entry.getKey()); + for (String instance : entry.getValue().keySet()) { + _anchorServerBytes.merge(instance, sizeBytes, Long::sum); + } + } + } + } + + Map getServerToHostedBytes(Map> assignment) { + Map serverToHostedBytes = new HashMap<>(); + for (Map.Entry> entry : assignment.entrySet()) { + long sizeBytes = getSegmentSizeBytes(entry.getKey()); + for (String instance : entry.getValue().keySet()) { + serverToHostedBytes.merge(instance, sizeBytes, Long::sum); + } + } + return serverToHostedBytes; + } + + /// Totals the budgeted bytes of the segments sharing each pair of current and target instances, over whichever + /// part of the assignment it is given. Those segments are all assigned the same instances, so an instance added + /// for one of them takes on all of them. + Map, Set>, Long> getInstancePairToBytes( + Map> currentAssignment, Map> targetAssignment) { + Map, Set>, Long> instancePairToBytes = new HashMap<>(); + for (Map.Entry> entry : currentAssignment.entrySet()) { + Map targetInstanceStateMap = targetAssignment.get(entry.getKey()); + if (targetInstanceStateMap != null) { + instancePairToBytes.merge(Pair.of(entry.getValue().keySet(), targetInstanceStateMap.keySet()), + getSegmentSizeBytes(entry.getKey()), Long::sum); + } + } + return instancePairToBytes; + } + + /// Returns the bytes each server may still take on in this step, i.e. `ceiling - currently hosted`. + StepDiskBudget forStep(Map> currentAssignment, + Map> targetAssignment) { + accountForNewSegments(currentAssignment); + Map hostedBytes = getServerToHostedBytes(currentAssignment); + Map targetBytes = getServerToHostedBytes(targetAssignment); + Set servers = new HashSet<>(_anchorServerBytes.keySet()); + servers.addAll(hostedBytes.keySet()); + servers.addAll(targetBytes.keySet()); + Map remainingBytes = new HashMap<>(); + for (String server : servers) { + long ceilingBytes = + Math.max(_anchorServerBytes.getOrDefault(server, 0L), targetBytes.getOrDefault(server, 0L)); + remainingBytes.put(server, Math.max(0L, ceilingBytes - hostedBytes.getOrDefault(server, 0L))); + } + // Segments that share a current and target instance pair are all assigned the same instances, so a newly added + // instance takes on all of them. Total their size up so that the instances to add can be picked with that in + // mind rather than only on the number of segments to offload + return new StepDiskBudget(this, remainingBytes, + getInstancePairToBytes(currentAssignment, targetAssignment)); + } + } + + /// The [DiskUsageBudget] left for a single rebalance step, drawn down as segments are assigned. + @VisibleForTesting + static class StepDiskBudget { + private final DiskUsageBudget _budget; + private final Map _remainingBytes; + // this is to record for each strict replica group, how many bytes they contribute altogether since they'll be + // moved together + private final Map, Set>, Long> _instancePairToRequiredBytes; + + private StepDiskBudget(DiskUsageBudget budget, Map remainingBytes, + Map, Set>, Long> instancePairToRequiredBytes) { + _budget = budget; + _remainingBytes = remainingBytes; + _instancePairToRequiredBytes = instancePairToRequiredBytes; + } + + /// Moves the instances that cannot take on all the segments sharing this current and target instance pair to the + /// back of the list so that they will be picked as the next step first. + /// + /// The instances that do not fit are kept at the back rather than dropped, so that an assignment is still produced + /// when none of them fits and the budget stays the only thing that decides whether it is applied. + List> retainInstancesThatFit( + List> instancesInfo, Set currentInstances, + Set targetInstances) { + long requiredBytes = _instancePairToRequiredBytes.getOrDefault(Pair.of(currentInstances, targetInstances), 0L); + List> fits = new ArrayList<>(instancesInfo.size()); + for (Triple instanceInfo : instancesInfo) { + if (_remainingBytes.getOrDefault(instanceInfo.getLeft(), 0L) >= requiredBytes) { + fits.add(instanceInfo); + } + } + return fits; + } + + Map getRemainingBytes() { + return _remainingBytes; + } + + Map, Set>, Long> getInstancePairToBytes( + Map> currentAssignment, Map> targetAssignment) { + return _budget.getInstancePairToBytes(currentAssignment, targetAssignment); + } + + long getSegmentSizeBytes(String segmentName) { + return _budget.getSegmentSizeBytes(segmentName); + } + + /// Charges every server in `serversAdded` for hosting `segmentName`, and returns whether all of them had the + /// space for it. Nothing is charged when any one of them did not, so the segment can be left where it is. + boolean tryCharge(Set serversAdded, String segmentName) { + return tryCharge(serversAdded, getSegmentSizeBytes(segmentName)); + } + + /// Charges every server in `serversAdded` `sizeBytes`, and returns whether all of them had the space for it. + /// Nothing is charged when any one of them did not, so the segments can be left where they are. + boolean tryCharge(Set serversAdded, long sizeBytes) { + for (String server : serversAdded) { + if (_remainingBytes.getOrDefault(server, 0L) < sizeBytes) { + return false; + } + } + for (String server : serversAdded) { + _remainingBytes.merge(server, -sizeBytes, Long::sum); + } + return true; + } } private static Map> getNextStrictReplicaGroupAssignment( Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher, - DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) { + DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger, + @Nullable StepDiskBudget stepBudget) { Map> nextAssignment = new TreeMap<>(); Map numSegmentsToOffloadMap = getNumSegmentsToOffloadMap(currentAssignment, targetAssignment); Map, Set>, Set> assignmentMap = new HashMap<>(); @@ -1710,7 +1965,7 @@ private static Map> getNextStrictReplicaGroupAssignm // Directly update the nextAssignment with anyServerExhaustedBatchSize = false and return if batching is disabled updateNextAssignmentForPartitionIdStrictReplicaGroup(currentAssignment, targetAssignment, nextAssignment, false, minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, - availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor); + availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor, stepBudget); return nextAssignment; } @@ -1734,7 +1989,7 @@ private static Map> getNextStrictReplicaGroupAssignm Map firstEntryInstanceStateMap = firstEntry.getValue(); SingleSegmentAssignment firstAssignment = getNextSingleSegmentAssignment(firstEntryInstanceStateMap, targetAssignment.get(firstEntry.getKey()), - minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap); + minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, stepBudget); Set serversAdded = getServersAddedInSingleSegmentAssignment(firstEntryInstanceStateMap, firstAssignment._instanceStateMap); boolean anyServerExhaustedBatchSize = false; @@ -1755,7 +2010,7 @@ private static Map> getNextStrictReplicaGroupAssignm } updateNextAssignmentForPartitionIdStrictReplicaGroup(curAssignment, targetAssignment, nextAssignment, anyServerExhaustedBatchSize, minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, - availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor); + availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor, stepBudget); } } @@ -1770,11 +2025,17 @@ private static void updateNextAssignmentForPartitionIdStrictReplicaGroup( boolean lowDiskMode, Map numSegmentsToOffloadMap, Map, Set>, Set> assignmentMap, Map, Set> availableInstancesMap, Map serverToNumSegmentsAddedSoFar, - DataLossRiskAssessor dataLossRiskAssessor) { + DataLossRiskAssessor dataLossRiskAssessor, @Nullable StepDiskBudget stepBudget) { if (anyServerExhaustedBatchSize) { // Exhausted the batch size for at least 1 server, just copy over the remaining segments as is nextAssignment.putAll(currentAssignment); } else { + // Build this map so later on we can skip updating the entire replica group if the group's total segment size + // doesn't fit in the instances disk budget + Map, Set>, Long> groupToSizeBytes = stepBudget == null ? Map.of() + : stepBudget.getInstancePairToBytes(currentAssignment, targetAssignment); + Map, Set>, Boolean> groupToFitsDisk = new HashMap<>(); + // Process all the partitionIds even if segmentsAddedSoFar becomes larger than batchSizePerServer // Can only do bestEfforts w.r.t. StrictReplicaGroup since a whole partition must be moved together for // maintaining consistency @@ -1784,9 +2045,29 @@ private static void updateNextAssignmentForPartitionIdStrictReplicaGroup( Map targetInstanceStateMap = targetAssignment.get(segmentName); SingleSegmentAssignment assignment = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, minAvailableReplicas, - lowDiskMode, numSegmentsToOffloadMap, assignmentMap); + lowDiskMode, numSegmentsToOffloadMap, assignmentMap, stepBudget); Set assignedInstances = assignment._instanceStateMap.keySet(); Set availableInstances = assignment._availableInstances; + // Strict replica group routing requires every segment that shares the same current and target instances to + // move together, so the disk budget has to be charged for such a group as a whole. The disk budget constraint + // makes the entire group either all move or all stay + if (stepBudget != null) { + Pair, Set> group = + Pair.of(currentInstanceStateMap.keySet(), targetInstanceStateMap.keySet()); + Boolean groupFitsDisk = groupToFitsDisk.get(group); + if (groupFitsDisk == null) { + Set serversAdded = + getServersAddedInSingleSegmentAssignment(currentInstanceStateMap, assignment._instanceStateMap); + groupFitsDisk = + serversAdded.isEmpty() || stepBudget.tryCharge(serversAdded, groupToSizeBytes.getOrDefault(group, 0L)); + groupToFitsDisk.put(group, groupFitsDisk); + } + if (!groupFitsDisk) { + // A later step picks the group up once the drops have freed up the space, skip for now + nextAssignment.put(segmentName, currentInstanceStateMap); + continue; + } + } availableInstancesMap.compute(assignedInstances, (k, currentAvailableInstances) -> { if (currentAvailableInstances == null) { // First segment assigned to these instances, use the new assignment and update the available instances @@ -1998,7 +2279,7 @@ private static String generateDataLossRiskMessage(String segmentName, boolean is private static Map> getNextNonStrictReplicaGroupAssignment( Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer, - DataLossRiskAssessor dataLossRiskAssessor) { + DataLossRiskAssessor dataLossRiskAssessor, @Nullable StepDiskBudget stepBudget) { Map serverToNumSegmentsAddedSoFar = new HashMap<>(); Map> nextAssignment = new TreeMap<>(); Map numSegmentsToOffloadMap = getNumSegmentsToOffloadMap(currentAssignment, targetAssignment); @@ -2009,7 +2290,7 @@ private static Map> getNextNonStrictReplicaGroupAssi Map targetInstanceStateMap = targetAssignment.get(segmentName); Map nextInstanceStateMap = getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, minAvailableReplicas, - lowDiskMode, numSegmentsToOffloadMap, assignmentMap)._instanceStateMap; + lowDiskMode, numSegmentsToOffloadMap, assignmentMap, stepBudget)._instanceStateMap; Set serversAddedForSegment = getServersAddedInSingleSegmentAssignment(currentInstanceStateMap, nextInstanceStateMap); boolean anyServerExhaustedBatchSize = false; @@ -2021,8 +2302,14 @@ private static Map> getNextNonStrictReplicaGroupAssi } } } - if (anyServerExhaustedBatchSize) { - // Exhausted the batch size for at least 1 server, set to existing assignment + // Leave the segment where it is when one of the servers it would be added to cannot take on its size without + // going over the disk it started the rebalance with. A later step picks it up once the drops have freed up the + // space. Note that tryCharge must not run once the batch size is known to be exhausted, as the segment is not + // going to be moved in that case. + boolean anyServerOutOfDisk = !anyServerExhaustedBatchSize && stepBudget != null + && !serversAddedForSegment.isEmpty() && !stepBudget.tryCharge(serversAddedForSegment, segmentName); + if (anyServerExhaustedBatchSize || anyServerOutOfDisk) { + // Exhausted the batch size or the disk budget for at least 1 server, set to existing assignment nextAssignment.put(segmentName, currentInstanceStateMap); } else { // Add the next assignment and update the segments added so far counts @@ -2080,6 +2367,126 @@ private static void updateNumSegmentsToOffloadMap(Map numSegmen } } + /// Returns the minimum available replicas the rebalance has to keep up, derived from + /// `minReplicasToKeepUpForNoDowntime` and the replication of the segments to be moved. Shared by the rebalance and + /// by the pre-checks so that the two cannot drift apart. + /// + /// NOTE: + /// 1. The calculation is based on the number of replicas of the target assignment. In case of increasing the number + /// of replicas for the current assignment, the current instance state map might not have enough replicas to reach + /// the minimum available replicas requirement. In this scenario we don't want to fail the check, but keep all the + /// current instances as this is the best we can do, and can help the table get out of this state. + /// 2. Only check the segments to be moved because we don't need to maintain available replicas for segments not + /// being moved, including segments with all replicas OFFLINE (error segments during consumption). + /// + /// @return the minimum available replicas, or [#ILLEGAL_MIN_AVAILABLE_REPLICAS] when + /// `minReplicasToKeepUpForNoDowntime` is not less than the replication of the segments to be moved + @VisibleForTesting + static int getMinAvailableReplicas(Map> currentAssignment, + Map> targetAssignment, List segmentsToMove, + int minReplicasToKeepUpForNoDowntime, Logger tableRebalanceLogger) { + int numReplicas = getMinNumReplicas(targetAssignment, segmentsToMove); + int minAvailableReplicas; + if (minReplicasToKeepUpForNoDowntime >= 0) { + // For non-negative value, use it as min available replicas + if (minReplicasToKeepUpForNoDowntime >= numReplicas) { + return ILLEGAL_MIN_AVAILABLE_REPLICAS; + } + minAvailableReplicas = minReplicasToKeepUpForNoDowntime; + } else { + // For negative value, use it as max unavailable replicas + minAvailableReplicas = Math.max(numReplicas + minReplicasToKeepUpForNoDowntime, 0); + } + + int numCurrentAssignmentReplicas = getMinNumReplicas(currentAssignment, segmentsToMove); + if (minAvailableReplicas > numCurrentAssignmentReplicas) { + tableRebalanceLogger.warn("minAvailableReplicas: {} larger than existing number of replicas: {}, " + + "resetting minAvailableReplicas to {}", minAvailableReplicas, numCurrentAssignmentReplicas, + numCurrentAssignmentReplicas); + minAvailableReplicas = numCurrentAssignmentReplicas; + } + return minAvailableReplicas; + } + + private static int getMinNumReplicas(Map> assignment, List segments) { + int numReplicas = Integer.MAX_VALUE; + for (String segment : segments) { + numReplicas = Math.min(assignment.get(segment).size(), numReplicas); + } + return numReplicas; + } + + /// Replays a low disk mode rebalance to find the servers it cannot keep within the disk they start with. + /// + /// Low disk mode bounds the disk each server may take on to the larger of what it hosts when the rebalance starts + /// and what the target assignment places on it. When no progress at all is possible within those bounds - servers + /// circularly waiting for one another to drop first - the rebalance relaxes them for a step rather than stalling, + /// which is the only way a server can end up over its bound. This runs [#getNextAssignment] over the whole rebalance + /// up front to find out whether that happens, and for which servers, before a single segment is moved. + /// + /// The replay is exact for the assignment it is given, with one caveat: it does not read the segment partition ids, + /// so with `batchSizePerServer` enabled and strict replica group routing the segments are grouped more coarsely than + /// the rebalance groups them, which changes the pacing of the moves but not the bounds themselves. + /// + /// @return server to the most bytes it would be pushed over its bound by, empty when the rebalance can complete + /// within the bound of every server + public static Map getServersForcedOverDiskBudget(Map> currentAssignment, + Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, + int batchSizePerServer, @Nullable TableSizeReader.TableSubTypeSizeDetails tableSizeDetails, + Logger tableRebalanceLogger) { + DiskUsageBudget diskUsageBudget = DiskUsageBudget.create(currentAssignment, tableSizeDetails); + Object2IntOpenHashMap segmentPartitionIdMap = new Object2IntOpenHashMap<>(); + Map serverToBytesOverBudget = new TreeMap<>(); + Map> assignment = currentAssignment; + for (int step = 1; step <= MAX_DISK_BUDGET_REPLAY_STEPS; step++) { + if (assignment.equals(targetAssignment)) { + return serverToBytesOverBudget; + } + Map remainingBytes = diskUsageBudget.forStep(assignment, targetAssignment).getRemainingBytes(); + Map> nextAssignment; + try { + nextAssignment = + getNextAssignment(assignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, true, + batchSizePerServer, segmentPartitionIdMap, DEFAULT_PARTITION_ID_FETCHER, new NoOpRiskAssessor(), + tableRebalanceLogger, diskUsageBudget); + } catch (Exception e) { + tableRebalanceLogger.warn("Caught exception while replaying the rebalance to check the low disk mode disk " + + "usage, reporting what was found up to step {}", step, e); + return serverToBytesOverBudget; + } + if (nextAssignment.equals(assignment)) { + // The rebalance cannot progress at all, with or without the bounds. Nothing more to find + return serverToBytesOverBudget; + } + getServerToAddedBytes(assignment, nextAssignment, diskUsageBudget).forEach((server, addedBytes) -> { + long bytesOverBudget = addedBytes - remainingBytes.getOrDefault(server, 0L); + if (bytesOverBudget > 0) { + serverToBytesOverBudget.merge(server, bytesOverBudget, Math::max); + } + }); + assignment = nextAssignment; + } + tableRebalanceLogger.warn("Gave up replaying the rebalance to check the low disk mode disk usage after {} steps", + MAX_DISK_BUDGET_REPLAY_STEPS); + return serverToBytesOverBudget; + } + + /// Returns the bytes each server is assigned on top of what it already hosts. + private static Map getServerToAddedBytes(Map> currentAssignment, + Map> nextAssignment, DiskUsageBudget diskUsageBudget) { + Map serverToAddedBytes = new HashMap<>(); + for (Map.Entry> entry : nextAssignment.entrySet()) { + Map currentInstanceStateMap = currentAssignment.get(entry.getKey()); + long sizeBytes = diskUsageBudget.getSegmentSizeBytes(entry.getKey()); + for (String instance : entry.getValue().keySet()) { + if (currentInstanceStateMap == null || !currentInstanceStateMap.containsKey(instance)) { + serverToAddedBytes.merge(instance, sizeBytes, Long::sum); + } + } + } + return serverToAddedBytes; + } + /// Returns the next assignment for a segment based on the current instance state map and the target instance state /// map /// with regard to the minimum available replicas requirement. @@ -2088,7 +2495,8 @@ private static void updateNumSegmentsToOffloadMap(Map numSegmen @VisibleForTesting static SingleSegmentAssignment getNextSingleSegmentAssignment(Map currentInstanceStateMap, Map targetInstanceStateMap, int minAvailableReplicas, boolean lowDiskMode, - Map numSegmentsToOffloadMap, Map, Set>, Set> assignmentMap) { + Map numSegmentsToOffloadMap, Map, Set>, Set> assignmentMap, + @Nullable StepDiskBudget stepBudget) { Map nextInstanceStateMap = new TreeMap<>(); // Assign the segment the same way as other segments if the current and target instances are the same. We need this @@ -2151,6 +2559,14 @@ static SingleSegmentAssignment getNextSingleSegmentAssignment(Map> instancesInfo = getSortedInstancesOnNumSegmentsToOffload(targetInstanceStateMap, nextInstanceStateMap, numSegmentsToOffloadMap); + if (stepBudget != null) { + // Keep only the instances with enough disk budget. Adding one without it would have the whole group of + // segments rejected when it is charged, holding back the instances that could have taken them in this step + instancesInfo = stepBudget.retainInstancesThatFit(instancesInfo, currentInstances, targetInstances); + // Fewer instances than the target assignment has is fine: the segments gain the rest in a later step, once + // the instances that are out of space have dropped what they owe + numInstancesToAdd = Math.min(numInstancesToAdd, instancesInfo.size()); + } for (int i = 0; i < numInstancesToAdd; i++) { Triple instanceInfo = instancesInfo.get(i); nextInstanceStateMap.put(instanceInfo.getLeft(), instanceInfo.getMiddle()); 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 ecf5583965a2..842bf5a79819 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 @@ -176,6 +176,30 @@ private RebalancePreCheckerResult checkDiskUtilization(RebalanceConfig rebalance THRESHOLD); } + /// `lowDiskMode` is not a blanket guarantee: when the rebalance cannot progress at all within the disk the servers + /// start with, it goes over rather than stalling. The pre-check replays the rebalance to find that out instead of + /// assuming `lowDiskMode` always avoids the transient usage. No reachable rebalance has been found that trips it - + /// see `DefaultRebalancePreChecker#getServersForcedOverDiskBudget` - so the replay is stubbed out here to cover the + /// reporting. + @Test + public void testOverThresholdDuringRebalanceIsAnErrorWhenLowDiskModeCannotAvoidIt() { + setDiskUsage(100L, 320L); + DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker() { + @Override + protected Map getServersForcedOverDiskBudget(PreCheckContext preCheckContext) { + return Map.of(SERVER_1, 150L); + } + }; + RebalancePreCheckerResult result = preChecker.checkDiskUtilization( + getPreCheckContext(lowDiskMode(), CURRENT_ASSIGNMENT, TARGET_ASSIGNMENT), THRESHOLD); + assertEquals(result.getPreCheckStatus(), PreCheckStatus.ERROR); + assertEquals(result.getMessage(), + "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"); + } + private static RebalanceConfig lowDiskMode() { RebalanceConfig rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setLowDiskMode(true); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/LowDiskModeRebalanceSimulatorTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/LowDiskModeRebalanceSimulatorTest.java new file mode 100644 index 000000000000..83c7eb6528af --- /dev/null +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/LowDiskModeRebalanceSimulatorTest.java @@ -0,0 +1,730 @@ +/** + * 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 it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; +import org.apache.pinot.common.restlet.resources.RebalanceConfig; +import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils; +import org.apache.pinot.controller.util.TableSizeReader; +import org.slf4j.LoggerFactory; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Drives the real [TableRebalancer#getNextAssignment] step by step, the way [TableRebalancer] does, and checks the +/// three things low disk mode promises: +/// +/// 1. the rebalance always reaches the target assignment, i.e. deferring segment moves never wedges it; +/// 2. no server is pushed above `max(bytes it held that this rebalance did not place, bytes the target places on it)`, +/// except where the disk utilization pre-check says up front that it will be; +/// 3. under strict replica group routing, segments sharing a current and target instance pair always move together. +/// +/// Segment sizes are the ones [TableSizeReader] reports. Where a scenario supplies none, every segment weighs one +/// byte, which turns the budget into a bound on the number of segments hosted. +/// +/// [#main] runs randomized sweeps over a far wider space than the tests do, and reports how often the budget had to be +/// given up and how many steps it took. That is for comparing two versions of the assignment logic by hand, not for +/// CI, so it is deliberately not a test. +public class LowDiskModeRebalanceSimulatorTest { + private static final String ONLINE = "ONLINE"; + private static final int MAX_STEPS = 200; + private static final long MIB = 1024L * 1024; + private static final TableRebalancer.PartitionIdFetcher DUMMY_PARTITION_FETCHER = segmentName -> 0; + /// Reads the partition id back out of a segment name, so that batching groups segments the way it would in a + /// cluster rather than treating the whole table as one partition + private static final TableRebalancer.PartitionIdFetcher LLC_PARTITION_FETCHER = segmentName -> { + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); + return llcSegmentName == null ? 0 : llcSegmentName.getPartitionGroupId(); + }; + private static final TableRebalancer.DataLossRiskAssessor NO_DATA_LOSS_RISK = + new TableRebalancer.NoOpRiskAssessor(); + + // --------------------------------------------------------------------------------------------------------------- + // The three guarantees + // --------------------------------------------------------------------------------------------------------------- + + /// The budget defers segment moves, so the thing to rule out is that it defers them forever. + @Test + public void testRebalanceAlwaysReachesTheTargetAssignment() { + for (Scenario scenario : scenarios()) { + SimResult result = simulate(scenario); + assertTrue(result.reachedTarget(), scenario._name + ": " + result._outcome + result.report()); + } + } + + /// The bound, and the pre-check that reports where it cannot be held. + /// + /// When no progress at all is possible within the budget, the rebalance gives it up for a step rather than stalling, + /// so the bound is not absolute. What is absolute is that the pre-check replay names exactly the servers that go + /// over, before any segment moves: an operator told nothing is entitled to a rebalance that stays within the disk + /// every server started with. + @Test + public void testSequenceStaysWithinBudgetUnlessThePreCheckSaysOtherwise() { + for (Scenario scenario : scenarios()) { + SimResult result = simulate(scenario); + if (scenario._injectAtStep > 0) { + // The pre-check runs before the rebalance, so it cannot know about segments that appear while it runs. Those + // are credited to the anchor when first seen, which has to keep the rebalance inside the budget on its own + assertEquals(result._serversOverBudget, Set.of(), + scenario._name + ": went outside the budget while segments were being added" + result.report()); + continue; + } + Set reported = TableRebalancer.getServersForcedOverDiskBudget(scenario._currentAssignment, + scenario._targetAssignment, scenario._minAvailableReplicas, scenario._enableStrictReplicaGroup, + scenario._batchSizePerServer, toTableSizeDetails(scenario._segmentSizeBytes), + LoggerFactory.getLogger(getClass())).keySet(); + // Asserting what the pre-check reports, and not only that it agrees with the rebalance, is what makes this fail + // when a change starts giving the budget up on a shape that used to hold: the equality on its own passes a + // rebalance that goes over and says so. + if (scenario._withinBudget) { + assertEquals(reported, Set.of(), + scenario._name + ": the pre-check says the budget cannot be held for a shape that it should" + + result.report()); + } else { + assertTrue(!reported.isEmpty(), scenario._name + ": expected the pre-check to report a server, so that the " + + "case where it does is covered. If a change made this shape hold, clear its withinBudget flag" + + result.report()); + } + assertEquals(result._serversOverBudget, reported, + scenario._name + ": the servers that went over the budget are not the ones the pre-check named" + + result.report()); + } + } + + /// Strict replica group routing needs every segment of a partition on the same instances at all times, so the budget + /// has to charge a whole group of segments at once. Charging one segment at a time fits the first few segments of a + /// partition and not the rest, which splits the partition across replica groups. + @Test + public void testStrictReplicaGroupMovesGroupsTogether() { + for (Scenario scenario : scenarios()) { + if (scenario._enableStrictReplicaGroup) { + SimResult result = simulate(scenario); + assertEquals(result._groupSplits, Set.of(), + scenario._name + ": split a group of segments across instances" + result.report()); + } + } + } + + // --------------------------------------------------------------------------------------------------------------- + // Scenarios + // --------------------------------------------------------------------------------------------------------------- + + /// The shapes worth holding down: the overlapping server sets that per-segment sequencing could not bound, a pure + /// scale-out as a control, both routing modes, batching, uneven segment sizes, segments arriving mid-rebalance, and + /// the assignment a randomized search found hardest. + private static List scenarios() { + Random random = new Random(11); + List scenarios = new ArrayList<>(); + List threeOld = List.of("host1", "host2", "host3"); + List threeNew = List.of("host2", "host3", "host4"); + List sixOld = servers(0, 6); + List sixNew = servers(3, 6); + + for (boolean strict : List.of(false, true)) { + String suffix = strict ? " [strict]" : ""; + // The original shape: every server in the overlap both sheds and takes on segments + scenarios.add(new Scenario("3 -> 3 servers, 2 in overlap" + suffix, roundRobin(6, 2, threeOld, 0), + roundRobin(6, 2, threeNew, 0), 1, strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, Map.of(), 0)); + // Same, with uneven segment sizes, which is what a budget counting segments cannot see + Map> skewed = roundRobin(24, 2, sixOld, 0); + scenarios.add(new Scenario("6 -> 6 servers, 3 in overlap, uneven sizes" + suffix, skewed, + roundRobin(24, 2, sixNew, 0), 1, strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, + skewedSizes(skewed.keySet(), random), 0)); + // Replication 3 against a minimum of 1, so a segment has to gain two instances at once + Map> deep = roundRobin(18, 3, sixOld, 0); + scenarios.add(new Scenario("6 -> 6 servers, replication 3, uneven sizes" + suffix, deep, + roundRobin(18, 3, sixNew, 0), 1, strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, + skewedSizes(deep.keySet(), random), 0)); + // Batching, which defers whole partitions for reasons of its own + Map> batched = roundRobin(24, 2, sixOld, 0); + scenarios.add(new Scenario("6 -> 6 servers, batchSizePerServer 2" + suffix, batched, + roundRobin(24, 2, sixNew, 0), 1, strict, 2, skewedSizes(batched.keySet(), random), 0)); + // A control: no server in the old set gains anything, so the budget is never binding + scenarios.add(new Scenario("4 -> 8 servers, pure scale-out" + suffix, roundRobin(24, 2, servers(0, 4), 0), + roundRobin(24, 2, servers(0, 8), 0), 1, strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, Map.of(), 0)); + } + + scenarios.addAll(replicaGroupScenarios(random)); + scenarios.add(hardestKnownStrictScenario()); + scenarios.add(knownToNeedTheBudgetGivenUpScenario()); + scenarios.addAll(midRebalanceUploadScenarios()); + return scenarios; + } + + /// The one assignment known to need the budget given up, found by [#sweepRandomGroupStructures]. Included so that + /// the pre-check is checked against a rebalance that does go over, and not only against ones that do not. + /// + /// `host05` starts on 2041 MiB and the target places 1523 MiB on it, so its ceiling is what it started with. It is + /// pinned with eight segments it cannot drop without going below three available replicas, and a step arrives where + /// it is the only place left to put a segment and has nothing spare. The rebalance gives the budget up for that step + /// and `host05` ends up 31 MiB over, 2% above its ceiling. + private static Scenario knownToNeedTheBudgetGivenUpScenario() { + // Groups of segments sharing one current and one target instance set, then the size of every segment in MiB + String[][] groups = { + {"host00,host01,host02,host03", "host02,host04,host05,host06", "0,1,2,3,4,5,6,7,8,9"}, + {"host00,host02,host06", "host00,host03,host05,host06", "10,11,12"}, + {"host00,host01,host03,host06", "host01,host02,host03,host05", "13,14,15"}, + {"host02,host04,host05,host06", "host00,host01,host03,host04", "16,17,18,19,20,21,22,23"}, + {"host01,host02,host03,host04", "host00,host02,host05,host06", "24,25,26,27,28,29,30,31,32"} + }; + long[] sizesInMib = { + 31, 12, 17, 9, 16, 14, 431, 664, 23, 23, 23, 20, 12, 12, 28, 9, 29, 1162, 768, 14, 15, 31, 8, 14, 21, 9, 14, + 25, 15, 25, 18, 24, 28 + }; + Map> current = new TreeMap<>(); + Map> target = new TreeMap<>(); + Map sizes = new TreeMap<>(); + for (String[] group : groups) { + for (String index : group[2].split(",")) { + String segment = String.format("segment%03d", Integer.parseInt(index)); + current.put(segment, instanceStateMap(group[0].split(","))); + target.put(segment, instanceStateMap(group[1].split(","))); + sizes.put(segment, sizesInMib[Integer.parseInt(index)] * MIB); + } + } + Scenario scenario = new Scenario("assignment known to need the budget given up", current, target, 3, false, 1, + sizes, 0); + scenario._withinBudget = false; + return scenario; + } + + /// Replica group assignment, where the placement is structured rather than balanced: every segment of a partition + /// sits on one server per replica group, so a partition's segments always share a current and target instance pair. + /// That is the shape strict replica group routing is built for, and it moves partitions between the servers of a + /// replica group rather than spreading them over all servers the way a balanced assignment does. + private static List replicaGroupScenarios(Random random) { + List> twoBySmall = List.of(List.of("host00", "host01"), List.of("host02", "host03")); + List> twoByGrown = + List.of(List.of("host00", "host01", "host04"), List.of("host02", "host03", "host05")); + List> twoByRotated = List.of(List.of("host01", "host04"), List.of("host03", "host05")); + List> threeBySmall = + List.of(List.of("host00", "host01"), List.of("host02", "host03"), List.of("host04", "host05")); + List> threeByGrown = List.of(List.of("host00", "host01", "host06"), + List.of("host02", "host03", "host07"), List.of("host04", "host05", "host08")); + + List scenarios = new ArrayList<>(); + for (boolean strict : List.of(false, true)) { + String suffix = strict ? " [strict]" : ""; + scenarios.add(replicaGroupScenario("replica groups grow 2 -> 3 servers each" + suffix, twoBySmall, twoByGrown, + strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, random)); + scenarios.add(replicaGroupScenario("replica groups, one server replaced in each" + suffix, twoBySmall, + twoByRotated, strict, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, random)); + } + scenarios.add(replicaGroupScenario("replica groups shrink 3 -> 2 servers each [strict]", twoByGrown, twoBySmall, + true, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, random)); + scenarios.add(replicaGroupScenario("three replica groups grow 2 -> 3 servers each [strict]", threeBySmall, + threeByGrown, true, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, random)); + // Batching groups segments by partition, which only means anything once partition ids are real + scenarios.add(replicaGroupScenario("replica groups grow, batchSizePerServer 2 [strict]", twoBySmall, twoByGrown, + true, 2, random)); + return scenarios; + } + + private static Scenario replicaGroupScenario(String name, List> currentReplicaGroups, + List> targetReplicaGroups, boolean strict, int batchSizePerServer, Random random) { + Map> current = replicaGroupAssignment(9, 3, currentReplicaGroups); + Map> target = replicaGroupAssignment(9, 3, targetReplicaGroups); + Scenario scenario = new Scenario(name, current, target, 1, strict, batchSizePerServer, + skewedSizes(current.keySet(), random), 0); + scenario._partitionIdFetcher = LLC_PARTITION_FETCHER; + return scenario; + } + + /// Assigns `numPartitions` partitions of `segmentsPerPartition` segments over `replicaGroups`, one server per + /// replica group per partition, the shape `ReplicaGroupSegmentAssignmentStrategy` produces. Segments are named so + /// that their partition id can be read back, which is what batching groups them by. + private static Map> replicaGroupAssignment(int numPartitions, int segmentsPerPartition, + List> replicaGroups) { + Map> assignment = new TreeMap<>(); + for (int partition = 0; partition < numPartitions; partition++) { + List instances = new ArrayList<>(replicaGroups.size()); + for (List replicaGroup : replicaGroups) { + instances.add(replicaGroup.get(partition % replicaGroup.size())); + } + for (int sequence = 0; sequence < segmentsPerPartition; sequence++) { + assignment.put(String.format("myTable__%d__%d__20240101T000000Z", partition, sequence), + SegmentAssignmentUtils.getInstanceStateMap(instances, ONLINE)); + } + } + return assignment; + } + + /// The assignment a randomized search over group structures found hardest: `host02` has 318 MiB free under its own + /// ceiling while the group it has to take on is six segments totalling roughly 1.2 GiB, and strict replica group + /// routing needs all of them to move at once. Choosing the instances to add without regard for the budget left every + /// group blocked here, which gave the budget up and pushed `host02` 36% over. + private static Scenario hardestKnownStrictScenario() { + // Groups of segments sharing one current and one target instance set, and the size of each segment in MiB + String[][] groups = { + {"host02,host07", "host00,host01", "0,1,2"}, + {"host00,host01", "host02,host04", "3,4,5,6,7"}, + {"host04", "host00,host06", "8,9,10,11,12,13"}, + {"host04,host05", "host00,host02", "14,15,16,17,18,19"} + }; + long[] sizesInMib = {1099, 8, 1135, 351, 27, 23, 895, 12, 8, 9, 20, 14, 30, 14, 23, 1152, 26, 16, 17, 18}; + + Map> current = new TreeMap<>(); + Map> target = new TreeMap<>(); + Map sizes = new TreeMap<>(); + for (String[] group : groups) { + for (String index : group[2].split(",")) { + String segment = String.format("segment%03d", Integer.parseInt(index)); + current.put(segment, instanceStateMap(group[0].split(","))); + target.put(segment, instanceStateMap(group[1].split(","))); + sizes.put(segment, sizesInMib[Integer.parseInt(index)] * MIB); + } + } + return new Scenario("hardest known strict replica group assignment", current, target, 1, true, + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, sizes, 0); + } + + /// Segments can be uploaded, or start consuming, while a rebalance runs. They are absent from the anchor the budget + /// was built on, so they are credited to it when first seen: the ceiling of a server whose net change is a loss is + /// pinned to what it started with, and would otherwise have the headroom the rebalance needs eaten by them. + /// + /// Both placements are covered. Strict replica group assignment overrides a new segment onto its partition's + /// existing placement to keep the partition collocated, which mid-rebalance is the placement being moved away from, + /// so the target assignment naming somewhere else is the normal case there rather than the exception. + private static List midRebalanceUploadScenarios() { + List scenarios = new ArrayList<>(); + for (boolean targetAgrees : List.of(true, false)) { + Map> current = roundRobin(12, 2, List.of("host1", "host2", "host3"), 0); + Map> target = roundRobin(12, 2, List.of("host2", "host3", "host4"), 0); + Map sizes = new TreeMap<>(); + current.keySet().forEach(segment -> sizes.put(segment, 100 * MIB)); + + Scenario scenario = new Scenario( + "3 -> 3 servers, segments uploaded mid-rebalance, target " + (targetAgrees ? "agrees" : "names elsewhere"), + current, target, 1, false, RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, sizes, 3); + for (int i = 0; i < 4; i++) { + String segment = "uploaded" + i; + scenario._injectedCurrent.put(segment, instanceStateMap("host1", "host2")); + scenario._injectedTarget.put(segment, + targetAgrees ? instanceStateMap("host1", "host2") : instanceStateMap("host3", "host4")); + sizes.put(segment, 400 * MIB); + } + scenarios.add(scenario); + } + return scenarios; + } + + // --------------------------------------------------------------------------------------------------------------- + // Simulator + // --------------------------------------------------------------------------------------------------------------- + + /// Runs `scenario` to the target assignment, or until it stops making progress, tracking for every server the most + /// it ever hosts and whether the budget had to be given up to get there. + private static SimResult simulate(Scenario scenario) { + Map> current = deepCopy(scenario._currentAssignment); + Map> target = deepCopy(scenario._targetAssignment); + Map sizes = new TreeMap<>(scenario._segmentSizeBytes); + TableRebalancer.DiskUsageBudget budget = + TableRebalancer.DiskUsageBudget.create(current, toTableSizeDetails(scenario._segmentSizeBytes)); + Object2IntOpenHashMap segmentPartitionIdMap = new Object2IntOpenHashMap<>(); + + SimResult result = new SimResult(scenario); + result._anchor = hostedBytes(current, sizes); + result._peak = new TreeMap<>(result._anchor); + + for (int step = 1; step <= MAX_STEPS; step++) { + if (scenario._injectAtStep == step) { + current.putAll(scenario._injectedCurrent); + target.putAll(scenario._injectedTarget); + // A segment this rebalance did not place raises the anchor of whoever is holding it, so that it neither eats + // the headroom the rebalance needs nor lets the rebalance raise its own ceiling + scenario._injectedCurrent.forEach((segment, instanceStateMap) -> instanceStateMap.keySet() + .forEach(instance -> result._anchor.merge(instance, sizes.get(segment), Long::sum))); + } + + Map allowed = budget.forStep(current, target).getRemainingBytes(); + Map> next; + try { + next = TableRebalancer.getNextAssignment(current, target, scenario._minAvailableReplicas, + scenario._enableStrictReplicaGroup, true, scenario._batchSizePerServer, segmentPartitionIdMap, + scenario._partitionIdFetcher, NO_DATA_LOSS_RISK, budget); + } catch (Exception e) { + result._outcome = "threw " + e; + break; + } + if (next.equals(current)) { + result._outcome = "could not make progress"; + break; + } + + // A server assigned more bytes than it was allowed means the budget was given up for this step + bytesAdded(current, next, sizes).forEach((server, addedBytes) -> { + if (addedBytes > allowed.getOrDefault(server, 0L)) { + result._serversOverBudget.add(server); + } + }); + recordGroupSplits(scenario, current, target, next, step, result); + + current = next; + hostedBytes(current, sizes).forEach((server, bytes) -> result._peak.merge(server, bytes, Math::max)); + result._steps = step; + if (current.equals(target)) { + result._outcome = "reached the target assignment"; + break; + } + if (step == MAX_STEPS) { + result._outcome = "hit the step limit"; + } + } + result._target = hostedBytes(target, sizes); + return result; + } + + private static Map bytesAdded(Map> current, + Map> next, Map sizes) { + Map added = new TreeMap<>(); + for (Map.Entry> entry : next.entrySet()) { + Map currentInstanceStateMap = current.get(entry.getKey()); + for (String instance : entry.getValue().keySet()) { + if (currentInstanceStateMap == null || !currentInstanceStateMap.containsKey(instance)) { + added.merge(instance, sizes.getOrDefault(entry.getKey(), 1L), Long::sum); + } + } + } + return added; + } + + /// Under strict replica group routing every segment sharing a current and target instance pair has to be assigned + /// the same instances, or a partition ends up split across replica groups. + private static void recordGroupSplits(Scenario scenario, Map> current, + Map> target, Map> next, int step, SimResult result) { + if (!scenario._enableStrictReplicaGroup) { + return; + } + // Keyed the way the rebalance groups segments: by their current and target instances and by partition. Batching + // moves one partition at a time, so two partitions sharing a pair of instance sets are free to move in different + // steps - what must never happen is segments of the same partition being assigned different instances. + Map, Set> groupToNextInstances = new HashMap<>(); + for (String segment : current.keySet()) { + List group = List.of(current.get(segment).keySet(), target.get(segment).keySet(), + scenario._partitionIdFetcher.fetch(segment)); + Set nextInstances = next.get(segment).keySet(); + Set alreadyAssigned = groupToNextInstances.putIfAbsent(group, nextInstances); + if (alreadyAssigned != null && !alreadyAssigned.equals(nextInstances)) { + result._groupSplits.add(String.format("step %d: partition %s of %s -> %s was split between %s and %s", step, + group.get(2), group.get(0), group.get(1), alreadyAssigned, nextInstances)); + } + } + } + + // --------------------------------------------------------------------------------------------------------------- + // Randomized sweeps, for comparing two versions of the assignment logic by hand. Not tests. + // --------------------------------------------------------------------------------------------------------------- + + public static void main(String[] args) { + sweepServerSetShapes(); + sweepRandomGroupStructures(new Random(2027), 30_000); + } + + /// Every combination of old and new server set sizes, overlap, replication and minimum available replicas, with and + /// without strict replica group routing and batching, over unevenly sized segments. + private static void sweepServerSetShapes() { + Random random = new Random(7); + List results = new ArrayList<>(); + for (int numOldServers = 3; numOldServers <= 8; numOldServers++) { + for (int numNewServers = 3; numNewServers <= 8; numNewServers++) { + for (int shift = 1; shift < numOldServers; shift++) { + for (int replication = 2; replication <= Math.min(3, Math.min(numOldServers, numNewServers)); replication++) { + for (int minAvailableReplicas = 1; minAvailableReplicas < replication; minAvailableReplicas++) { + for (boolean strict : List.of(false, true)) { + for (int batchSizePerServer : List.of(RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, 2)) { + int numSegments = 12 * replication; + Map> current = + roundRobin(numSegments, replication, servers(0, numOldServers), 0); + results.add(simulate(new Scenario( + String.format("old=%d new=%d shift=%d replication=%d minAvail=%d strict=%s batch=%d", + numOldServers, numNewServers, shift, replication, minAvailableReplicas, strict, + batchSizePerServer), current, + roundRobin(numSegments, replication, servers(shift, numNewServers), 0), minAvailableReplicas, + strict, batchSizePerServer, skewedSizes(current.keySet(), random), 0))); + } + } + } + } + } + } + } + report("Server set shapes", results); + } + + /// Random current and target instance sets, group sizes, replica counts and size distributions. A rebalance the + /// budget cannot carry out is a rare structure, so this reaches shapes the fixed scenarios do not. + private static void sweepRandomGroupStructures(Random random, int numTrials) { + List results = new ArrayList<>(); + for (int trial = 0; trial < numTrials; trial++) { + int numServers = 3 + random.nextInt(8); + int replication = 2 + random.nextInt(3); + int minAvailableReplicas = 1 + random.nextInt(replication - 1); + List allServers = servers(0, numServers); + Map> current = new TreeMap<>(); + Map> target = new TreeMap<>(); + int segmentId = 0; + int numGroups = 2 + random.nextInt(7); + for (int group = 0; group < numGroups; group++) { + List currentInstances = pickServers(allServers, + minAvailableReplicas + random.nextInt(replication - minAvailableReplicas + 1), random); + List targetInstances = pickServers(allServers, replication, random); + int numSegmentsInGroup = 1 + random.nextInt(10); + for (int i = 0; i < numSegmentsInGroup; i++) { + String segment = String.format("segment%04d", segmentId++); + current.put(segment, SegmentAssignmentUtils.getInstanceStateMap(currentInstances, ONLINE)); + target.put(segment, SegmentAssignmentUtils.getInstanceStateMap(targetInstances, ONLINE)); + } + } + // Only run what the rebalance itself would accept + if (TableRebalancer.getMinAvailableReplicas(current, target, + SegmentAssignmentUtils.getSegmentsToMove(current, target), minAvailableReplicas, + LoggerFactory.getLogger(LowDiskModeRebalanceSimulatorTest.class)) != minAvailableReplicas) { + continue; + } + results.add(simulate(new Scenario("trial " + trial, current, target, minAvailableReplicas, random.nextBoolean(), + List.of(RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, 1, 2, 20).get(random.nextInt(4)), + skewedSizes(current.keySet(), random), 0))); + } + report("Random group structures", results); + } + + /// Prints how often the budget had to be given up, how far over it went and how many steps it took, which is what + /// there is to compare between two versions of the assignment logic. + private static void report(String label, List results) { + int gaveUpBudget = 0; + int notReached = 0; + int splits = 0; + long totalSteps = 0; + int maxSteps = 0; + double worstAmplification = 1; + SimResult worst = null; + for (SimResult result : results) { + if (!result._serversOverBudget.isEmpty()) { + gaveUpBudget++; + if (result.amplification() > worstAmplification) { + worstAmplification = result.amplification(); + worst = result; + } + } + notReached += result.reachedTarget() ? 0 : 1; + splits += result._groupSplits.size(); + totalSteps += result._steps; + maxSteps = Math.max(maxSteps, result._steps); + } + System.out.printf("%n%s, over %d scenarios%n", label, results.size()); + System.out.printf(" gave up the budget in : %d%n", gaveUpBudget); + System.out.printf(" worst amplification : %.2fx%n", worstAmplification); + System.out.printf(" split a group of segments: %d%n", splits); + System.out.printf(" did not reach the target : %d%n", notReached); + System.out.printf(" steps : %.1f mean, %d max%n", (double) totalSteps / results.size(), + maxSteps); + if (worst != null) { + System.out.println(worst.report()); + } + } + + // --------------------------------------------------------------------------------------------------------------- + // Model and helpers + // --------------------------------------------------------------------------------------------------------------- + + private static class Scenario { + final String _name; + final Map> _currentAssignment; + final Map> _targetAssignment; + final int _minAvailableReplicas; + final boolean _enableStrictReplicaGroup; + final int _batchSizePerServer; + final Map _segmentSizeBytes; + /// Step at which to add segments this rebalance did not place, 0 for none + final int _injectAtStep; + /// Whether the budget can be held for this shape. False for the one shape known to need it given up, which is + /// what exercises the pre-check reporting a server rather than reporting nothing. + boolean _withinBudget = true; + /// How the rebalance reads partition ids, which batching groups segments by. Replica group scenarios name their + /// segments so that the real fetcher resolves them, the rest have no meaningful partition. + TableRebalancer.PartitionIdFetcher _partitionIdFetcher = DUMMY_PARTITION_FETCHER; + final Map> _injectedCurrent = new TreeMap<>(); + final Map> _injectedTarget = new TreeMap<>(); + + Scenario(String name, Map> currentAssignment, + Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, + int batchSizePerServer, Map segmentSizeBytes, int injectAtStep) { + _name = name; + _currentAssignment = currentAssignment; + _targetAssignment = targetAssignment; + _minAvailableReplicas = minAvailableReplicas; + _enableStrictReplicaGroup = enableStrictReplicaGroup; + _batchSizePerServer = batchSizePerServer; + _segmentSizeBytes = segmentSizeBytes; + _injectAtStep = injectAtStep; + } + } + + private static class SimResult { + final Scenario _scenario; + final Set _serversOverBudget = new TreeSet<>(); + final Set _groupSplits = new TreeSet<>(); + Map _anchor = new TreeMap<>(); + Map _target = new TreeMap<>(); + Map _peak = new TreeMap<>(); + int _steps; + String _outcome = "did not finish"; + + SimResult(Scenario scenario) { + _scenario = scenario; + } + + boolean reachedTarget() { + return "reached the target assignment".equals(_outcome); + } + + private long bound(String server) { + return Math.max(_anchor.getOrDefault(server, 0L), _target.getOrDefault(server, 0L)); + } + + /// The most any server held, as a multiple of what it was allowed to hold. + double amplification() { + double worst = 1; + for (String server : _peak.keySet()) { + long bound = bound(server); + if (bound > 0) { + worst = Math.max(worst, (double) _peak.get(server) / bound); + } + } + return worst; + } + + /// Rendered only when an assertion fails, or for the worst scenario of a sweep. + String report() { + StringBuilder sb = new StringBuilder("\n ").append(_scenario._name).append(" — ").append(_outcome) + .append(" after ").append(_steps).append(" steps\n"); + sb.append(String.format(" %-8s %10s %10s %10s %10s%n", "server", "anchor", "target", "bound", "peak")); + for (String server : union(_anchor.keySet(), _peak.keySet())) { + long bound = bound(server); + long peak = _peak.getOrDefault(server, 0L); + sb.append(String.format(" %-8s %10s %10s %10s %10s%s%n", server, mib(_anchor.getOrDefault(server, 0L)), + mib(_target.getOrDefault(server, 0L)), mib(bound), mib(peak), peak > bound ? " OVER" : "")); + } + if (!_serversOverBudget.isEmpty()) { + sb.append(" budget given up for: ").append(_serversOverBudget).append('\n'); + } + _groupSplits.forEach(split -> sb.append(" ").append(split).append('\n')); + return sb.toString(); + } + } + + private static List servers(int startIndex, int count) { + List servers = new ArrayList<>(count); + for (int i = 0; i < count; i++) { + servers.add(String.format("host%02d", startIndex + i)); + } + return servers; + } + + private static List pickServers(List allServers, int count, Random random) { + List shuffled = new ArrayList<>(allServers); + for (int i = shuffled.size() - 1; i > 0; i--) { + shuffled.set(i, shuffled.set(random.nextInt(i + 1), shuffled.get(i))); + } + List picked = new ArrayList<>(shuffled.subList(0, Math.min(count, shuffled.size()))); + picked.sort(null); + return picked; + } + + /// Round-robins `numSegments` segments with `replication` replicas over `servers`, the shape + /// `BalanceNumSegmentAssignmentStrategy` produces. + private static Map> roundRobin(int numSegments, int replication, List servers, + int offset) { + Map> assignment = new TreeMap<>(); + int cursor = offset; + for (int i = 0; i < numSegments; i++) { + List instances = new ArrayList<>(replication); + for (int r = 0; r < replication; r++) { + instances.add(servers.get(cursor++ % servers.size())); + } + assignment.put(String.format("segment%03d", i), SegmentAssignmentUtils.getInstanceStateMap(instances, ONLINE)); + } + return assignment; + } + + /// Most segments small and a few an order of magnitude larger, which is what a table with mixed pushes and varying + /// retention looks like, and what a budget counting segments rather than bytes cannot see. + private static Map skewedSizes(Collection segments, Random random) { + Map segmentSizeBytes = new TreeMap<>(); + for (String segment : segments) { + segmentSizeBytes.put(segment, + random.nextInt(10) < 8 ? (8 + random.nextInt(24)) * MIB : (320 + random.nextInt(960)) * MIB); + } + return segmentSizeBytes; + } + + private static Map instanceStateMap(String... instances) { + return SegmentAssignmentUtils.getInstanceStateMap(List.of(instances), ONLINE); + } + + private static Map> deepCopy(Map> assignment) { + Map> copy = new TreeMap<>(); + assignment.forEach((segment, instanceStateMap) -> copy.put(segment, new TreeMap<>(instanceStateMap))); + return copy; + } + + private static Map hostedBytes(Map> assignment, Map sizes) { + Map bytes = new TreeMap<>(); + assignment.forEach((segment, instanceStateMap) -> instanceStateMap.keySet() + .forEach(instance -> bytes.merge(instance, sizes.getOrDefault(segment, 1L), Long::sum))); + return bytes; + } + + /// The sizes as [TableSizeReader] reports them, so the real extraction path is exercised. `null` where the scenario + /// supplies none, which makes every segment weigh one byte. + private static TableSizeReader.TableSubTypeSizeDetails toTableSizeDetails(Map segmentSizeBytes) { + if (segmentSizeBytes.isEmpty()) { + return null; + } + TableSizeReader.TableSubTypeSizeDetails tableSizeDetails = new TableSizeReader.TableSubTypeSizeDetails(); + segmentSizeBytes.forEach((segment, sizeBytes) -> { + TableSizeReader.SegmentSizeDetails segmentSizeDetails = new TableSizeReader.SegmentSizeDetails(); + segmentSizeDetails._maxReportedSizePerReplicaInBytes = sizeBytes; + tableSizeDetails._segments.put(segment, segmentSizeDetails); + }); + return tableSizeDetails; + } + + private static Set union(Set a, Set b) { + Set union = new TreeSet<>(a); + union.addAll(b); + return union; + } + + private static String mib(long bytes) { + return bytes >= MIB ? (bytes / MIB) + "M" : Long.toString(bytes); + } +} diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java index da98f32aba85..527bd7392e46 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java @@ -573,7 +573,7 @@ private TableRebalancer.SingleSegmentAssignment getNextSingleSegmentAssignment( } Map, Set>, Set> assignmentMap = new HashMap<>(); return TableRebalancer.getNextSingleSegmentAssignment(currentInstanceStateMap, targetInstanceStateMap, - minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap); + minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, null); } @Test @@ -1305,6 +1305,68 @@ public void testAssignmentWithLowDiskMode() { assertEquals(nextAssignment, targetAssignment); } + /// In low disk mode a server that is in both the current and the target assignment can be in the drop phase for + /// some segments while being in the add phase for others, which makes it hold both sets at once. Verify that the + /// adds to such a server are deferred until it has dropped everything the target assignment does not place on it. + @Test + public void testAssignmentWithLowDiskModeDefersAddsToServersWithPendingDrops() { + // 6 segments with replication 2 over host1/host2/host3, moving to host2/host3/host4. host2 and host3 are in both + // the current and the target assignment, so they both offload and onload segments. + Map> currentAssignment = new TreeMap<>(); + currentAssignment.put("segment1", SegmentAssignmentUtils.getInstanceStateMap(List.of("host1", "host2"), ONLINE)); + currentAssignment.put("segment2", SegmentAssignmentUtils.getInstanceStateMap(List.of("host3", "host1"), ONLINE)); + currentAssignment.put("segment3", SegmentAssignmentUtils.getInstanceStateMap(List.of("host2", "host3"), ONLINE)); + currentAssignment.put("segment4", SegmentAssignmentUtils.getInstanceStateMap(List.of("host1", "host2"), ONLINE)); + currentAssignment.put("segment5", SegmentAssignmentUtils.getInstanceStateMap(List.of("host3", "host1"), ONLINE)); + currentAssignment.put("segment6", SegmentAssignmentUtils.getInstanceStateMap(List.of("host2", "host3"), ONLINE)); + + Map> targetAssignment = new TreeMap<>(); + targetAssignment.put("segment1", SegmentAssignmentUtils.getInstanceStateMap(List.of("host2", "host3"), ONLINE)); + targetAssignment.put("segment2", SegmentAssignmentUtils.getInstanceStateMap(List.of("host4", "host2"), ONLINE)); + targetAssignment.put("segment3", SegmentAssignmentUtils.getInstanceStateMap(List.of("host3", "host4"), ONLINE)); + targetAssignment.put("segment4", SegmentAssignmentUtils.getInstanceStateMap(List.of("host2", "host3"), ONLINE)); + targetAssignment.put("segment5", SegmentAssignmentUtils.getInstanceStateMap(List.of("host4", "host2"), ONLINE)); + targetAssignment.put("segment6", SegmentAssignmentUtils.getInstanceStateMap(List.of("host3", "host4"), ONLINE)); + + // Every host hosts 4 segments to begin with, and host2 and host3 also host 4 in the target assignment + assertEquals(getNumSegmentsHosted(currentAssignment), Map.of("host1", 4, "host2", 4, "host3", 4)); + assertEquals(getNumSegmentsHosted(targetAssignment), Map.of("host2", 4, "host3", 4, "host4", 4)); + + // With no per-segment sizes available every segment counts as one byte. host1, host2 and host3 are all already + // at the disk they started the rebalance with, so none of them may take on a segment until they have dropped one, + // while host4 is empty and can take on its full target. + TableRebalancer.StepDiskBudget stepBudget = + TableRebalancer.DiskUsageBudget.create(currentAssignment, null).forStep(currentAssignment, targetAssignment); + assertEquals(stepBudget.getRemainingBytes(), Map.of("host1", 0L, "host2", 0L, "host3", 0L, "host4", 4L)); + + Map> nextAssignment = currentAssignment; + int maxNumSegmentsHostedByHost3 = 4; + for (int step = 1; step <= 10 && !nextAssignment.equals(targetAssignment); step++) { + nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 1, false, true, + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); + Map numSegmentsHosted = getNumSegmentsHosted(nextAssignment); + maxNumSegmentsHostedByHost3 = Math.max(maxNumSegmentsHostedByHost3, numSegmentsHosted.getOrDefault("host3", 0)); + if (step == 2) { + assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(List.of("host2"))); + assertEquals(nextAssignment.get("segment4").keySet(), new TreeSet<>(List.of("host2"))); + } + } + assertEquals(nextAssignment, targetAssignment); + // host3 starts and ends with 4 segments, so it must never host more than 4 at any point of the rebalance + assertEquals(maxNumSegmentsHostedByHost3, 4); + } + + private static Map getNumSegmentsHosted(Map> assignment) { + Map numSegmentsHosted = new TreeMap<>(); + for (Map instanceStateMap : assignment.values()) { + for (String instance : instanceStateMap.keySet()) { + numSegmentsHosted.merge(instance, 1, Integer::sum); + } + } + return numSegmentsHosted; + } + @Test public void testIsExternalViewConverged() { String offlineTableName = "testTable_OFFLINE";