From 986950de584f0fd146b056f09b2087f167c410b0 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Wed, 12 Feb 2025 11:19:52 -0800 Subject: [PATCH 1/7] Add a dry-run summary mode for TableRebalance which only returns a summary of the dry-run results --- .../resources/PinotTableRestletResource.java | 10 +- .../helix/core/rebalance/RebalanceConfig.java | 22 +- .../helix/core/rebalance/RebalanceResult.java | 11 +- .../rebalance/RebalanceSummaryResult.java | 127 ++++++++ .../helix/core/rebalance/TableRebalancer.java | 157 ++++++++-- .../tenant/DefaultTenantRebalancer.java | 5 +- .../TableRebalancerClusterStatelessTest.java | 280 +++++++++++++++++- 7 files changed, 573 insertions(+), 39 deletions(-) create mode 100644 pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java index 74361a62a82b..0942325677eb 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java @@ -604,6 +604,9 @@ public RebalanceResult rebalance( @ApiParam(value = "OFFLINE|REALTIME", required = true) @QueryParam("type") String tableTypeStr, @ApiParam(value = "Whether to rebalance table in dry-run mode") @DefaultValue("false") @QueryParam("dryRun") boolean dryRun, + @ApiParam(value = "Whether to return dry-run summary instead of full, dry-run must be enabled to use this") + @DefaultValue("false") @QueryParam("summary") + boolean summary, @ApiParam(value = "Whether to enable pre-checks for table, must be in dry-run mode to enable") @DefaultValue("false") @QueryParam("preChecks") boolean preChecks, @ApiParam(value = "Whether to reassign instances before reassigning segments") @DefaultValue("false") @@ -646,6 +649,7 @@ public RebalanceResult rebalance( String tableNameWithType = constructTableNameWithType(tableName, tableTypeStr); RebalanceConfig rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setDryRun(dryRun); + rebalanceConfig.setSummary(summary); rebalanceConfig.setPreChecks(preChecks); rebalanceConfig.setReassignInstances(reassignInstances); rebalanceConfig.setIncludeConsuming(includeConsuming); @@ -666,7 +670,7 @@ public RebalanceResult rebalance( String rebalanceJobId = TableRebalancer.createUniqueRebalanceJobIdentifier(); try { - if (dryRun || preChecks || downtime) { + if (dryRun || summary || preChecks || downtime) { // For dry-run, preChecks or rebalance with downtime, directly return the rebalance result as it should return // immediately return _pinotHelixResourceManager.rebalanceTable(tableNameWithType, rebalanceConfig, rebalanceJobId, false); @@ -687,7 +691,7 @@ public RebalanceResult rebalance( String errorMsg = String.format("Caught exception/error while rebalancing table: %s", tableNameWithType); LOGGER.error(errorMsg, t); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, null, null, null, - null); + null, null); } }); boolean isJobIdPersisted = waitForRebalanceToPersist( @@ -708,7 +712,7 @@ public RebalanceResult rebalance( return new RebalanceResult(dryRunResult.getJobId(), RebalanceResult.Status.IN_PROGRESS, "In progress, check controller logs for updates", dryRunResult.getInstanceAssignment(), dryRunResult.getTierInstanceAssignment(), dryRunResult.getSegmentAssignment(), - dryRunResult.getPreChecksResult()); + dryRunResult.getPreChecksResult(), dryRunResult.getRebalanceSummaryResult()); } else { // If dry-run failed or is no-op, return the dry-run result return dryRunResult; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java index e5bd39f0ec0e..06a9e171c3d7 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java @@ -34,6 +34,11 @@ public class RebalanceConfig { @ApiModelProperty(example = "false") private boolean _dryRun = false; + // Whether to return only dry-run summary instead of full dry-run output, can only be used in dry-run mode + @JsonProperty("summary") + @ApiModelProperty(example = "false") + private boolean _summary = false; + // Whether to perform pre-checks for rebalance. This only returns the status of each pre-check and does not fail // rebalance @JsonProperty("preChecks") @@ -124,6 +129,14 @@ public void setDryRun(boolean dryRun) { _dryRun = dryRun; } + public boolean isSummary() { + return _summary; + } + + public void setSummary(boolean summary) { + _summary = summary; + } + public boolean isPreChecks() { return _preChecks; } @@ -246,10 +259,10 @@ public void setRetryInitialDelayInMs(long retryInitialDelayInMs) { @Override public String toString() { - return "RebalanceConfig{" + "_dryRun=" + _dryRun + ", preChecks=" + _preChecks + ", _reassignInstances=" - + _reassignInstances + ", _includeConsuming=" + _includeConsuming + ", _bootstrap=" + _bootstrap - + ", _downtime=" + _downtime + ", _minAvailableReplicas=" + _minAvailableReplicas + ", _bestEfforts=" - + _bestEfforts + ", _externalViewCheckIntervalInMs=" + _externalViewCheckIntervalInMs + return "RebalanceConfig{" + "_dryRun=" + _dryRun + ", _summary=" + _summary + ", preChecks=" + _preChecks + + ", _reassignInstances=" + _reassignInstances + ", _includeConsuming=" + _includeConsuming + ", _bootstrap=" + + _bootstrap + ", _downtime=" + _downtime + ", _minAvailableReplicas=" + _minAvailableReplicas + + ", _bestEfforts=" + _bestEfforts + ", _externalViewCheckIntervalInMs=" + _externalViewCheckIntervalInMs + ", _externalViewStabilizationTimeoutInMs=" + _externalViewStabilizationTimeoutInMs + ", _updateTargetTier=" + _updateTargetTier + ", _heartbeatIntervalInMs=" + _heartbeatIntervalInMs + ", _heartbeatTimeoutInMs=" + _heartbeatTimeoutInMs + ", _maxAttempts=" + _maxAttempts + ", _retryInitialDelayInMs=" @@ -259,6 +272,7 @@ public String toString() { public static RebalanceConfig copy(RebalanceConfig cfg) { RebalanceConfig rc = new RebalanceConfig(); rc._dryRun = cfg._dryRun; + rc._summary = cfg._summary; rc._preChecks = cfg._preChecks; rc._reassignInstances = cfg._reassignInstances; rc._includeConsuming = cfg._includeConsuming; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceResult.java index 2b81d8d78b66..ac768a8c4638 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceResult.java @@ -42,6 +42,8 @@ public class RebalanceResult { private final String _description; @JsonInclude(JsonInclude.Include.NON_NULL) private final Map _preChecksResult; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RebalanceSummaryResult _rebalanceSummaryResult; @JsonCreator public RebalanceResult(@JsonProperty(value = "jobId", required = true) String jobId, @@ -50,7 +52,8 @@ public RebalanceResult(@JsonProperty(value = "jobId", required = true) String jo @JsonProperty("instanceAssignment") @Nullable Map instanceAssignment, @JsonProperty("tierInstanceAssignment") @Nullable Map tierInstanceAssignment, @JsonProperty("segmentAssignment") @Nullable Map> segmentAssignment, - @JsonProperty("preChecksResult") @Nullable Map preChecksResult) { + @JsonProperty("preChecksResult") @Nullable Map preChecksResult, + @JsonProperty("rebalanceSummaryResult") @Nullable RebalanceSummaryResult rebalanceSummaryResult) { _jobId = jobId; _status = status; _description = description; @@ -58,6 +61,7 @@ public RebalanceResult(@JsonProperty(value = "jobId", required = true) String jo _tierInstanceAssignment = tierInstanceAssignment; _segmentAssignment = segmentAssignment; _preChecksResult = preChecksResult; + _rebalanceSummaryResult = rebalanceSummaryResult; } @JsonProperty @@ -95,6 +99,11 @@ public Map getPreChecksResult() { return _preChecksResult; } + @JsonProperty + public RebalanceSummaryResult getRebalanceSummaryResult() { + return _rebalanceSummaryResult; + } + public enum Status { // FAILED if the job has ended with known exceptions; // ABORTED if the job is stopped by others but retry is still allowed; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java new file mode 100644 index 000000000000..f47d38b61f09 --- /dev/null +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -0,0 +1,127 @@ +/** + * 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 com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.Map; +import javax.annotation.Nullable; + + +/** + * Holds the summary data of the rebalance result + */ +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonInclude(JsonInclude.Include.NON_NULL) +public class RebalanceSummaryResult { + + public static class ServerSegmentChangeInfo { + public final int _totalNewSegments; + public final int _totalExistingSegments; + public final int _segmentsAdded; + public final int _segmentsDeleted; + public final int _segmentsUnchanged; + + @JsonCreator + public ServerSegmentChangeInfo(@JsonProperty("totalNewSegments") int totalNewSegments, + @JsonProperty("totalExistingSegments") int totalExistingSegments, + @JsonProperty("segmentsAdded") int segmentsAdded, @JsonProperty("segmentsDeleted") int segmentsDeleted, + @JsonProperty("segmentsUnchanged") int segmentsUnchanged) { + _totalNewSegments = totalNewSegments; + _totalExistingSegments = totalExistingSegments; + _segmentsAdded = segmentsAdded; + _segmentsDeleted = segmentsDeleted; + _segmentsUnchanged = segmentsUnchanged; + } + } + + public static class RebalanceChangeInfo { + public final int _existingValue; + public final int _newValue; + + @JsonCreator + public RebalanceChangeInfo(@JsonProperty("existingValue") int existingValue, + @JsonProperty("newValue") int newValue) { + _existingValue = existingValue; + _newValue = newValue; + } + } + + // TODO: Add stats about total data size to be moved and estimated calculations on how long the data move can take + // during rebalance. Estimations are fine based on total table size / total number of segments + private int _totalSegmentsToBeMoved; + @JsonInclude(JsonInclude.Include.NON_NULL) + private RebalanceChangeInfo _numServers; + @JsonInclude(JsonInclude.Include.NON_NULL) + private RebalanceChangeInfo _replicationFactor; + @JsonInclude(JsonInclude.Include.NON_NULL) + private RebalanceChangeInfo _numUniqueSegments; + @JsonInclude(JsonInclude.Include.NON_NULL) + private RebalanceChangeInfo _numTotalSegments; + @JsonInclude(JsonInclude.Include.NON_NULL) + private Map _serverSegmentChangeInfo; + + @JsonCreator + public RebalanceSummaryResult( + @JsonProperty("totalSegmentsToBeMoved") int totalSegmentsToBeMoved, + @JsonProperty("numServers") @Nullable RebalanceChangeInfo numServers, + @JsonProperty("replicationFactor") @Nullable RebalanceChangeInfo replicationFactor, + @JsonProperty("numUniqueSegments") @Nullable RebalanceChangeInfo numUniqueSegments, + @JsonProperty("numTotalSegments") @Nullable RebalanceChangeInfo numTotalSegments, + @JsonProperty("serverSegmentChangeInfo") @Nullable Map serverSegmentChangeInfo) { + _totalSegmentsToBeMoved = totalSegmentsToBeMoved; + _numServers = numServers; + _replicationFactor = replicationFactor; + _numUniqueSegments = numUniqueSegments; + _numTotalSegments = numTotalSegments; + _serverSegmentChangeInfo = serverSegmentChangeInfo; + } + + @JsonProperty + public int getTotalSegmentsToBeMoved() { + return _totalSegmentsToBeMoved; + } + + @JsonProperty + public RebalanceChangeInfo getNumServers() { + return _numServers; + } + + @JsonProperty + public RebalanceChangeInfo getReplicationFactor() { + return _replicationFactor; + } + + @JsonProperty + public RebalanceChangeInfo getNumUniqueSegments() { + return _numUniqueSegments; + } + + @JsonProperty + public RebalanceChangeInfo getNumTotalSegments() { + return _numTotalSegments; + } + + @JsonProperty + public Map getServerSegmentChangeInfo() { + return _serverSegmentChangeInfo; + } +} 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 5dc7c4d8eae2..abacfb759b7c 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 @@ -175,6 +175,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb rebalanceJobId = createUniqueRebalanceJobIdentifier(); } boolean dryRun = rebalanceConfig.isDryRun(); + boolean summary = rebalanceConfig.isSummary(); boolean preChecks = rebalanceConfig.isPreChecks(); boolean reassignInstances = rebalanceConfig.isReassignInstances(); boolean includeConsuming = rebalanceConfig.isIncludeConsuming(); @@ -189,14 +190,19 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb && RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE.equalsIgnoreCase( tableConfig.getRoutingConfig().getInstanceSelectorType()); LOGGER.info( - "Start rebalancing table: {} with dryRun: {}, preChecks: {}, reassignInstances: {}, includeConsuming: {}, " - + "bootstrap: {}, downtime: {}, minReplicasToKeepUpForNoDowntime: {}, enableStrictReplicaGroup: {}, " - + "lowDiskMode: {}, bestEfforts: {}, externalViewCheckIntervalInMs: {}, " + "Start rebalancing table: {} with dryRun: {}, summary: {}, preChecks: {}, reassignInstances: {}, " + + "includeConsuming: {}, bootstrap: {}, downtime: {}, minReplicasToKeepUpForNoDowntime: {}, " + + "enableStrictReplicaGroup: {}, lowDiskMode: {}, bestEfforts: {}, externalViewCheckIntervalInMs: {}, " + "externalViewStabilizationTimeoutInMs: {}", - tableNameWithType, dryRun, preChecks, reassignInstances, includeConsuming, bootstrap, downtime, + tableNameWithType, dryRun, summary, preChecks, reassignInstances, includeConsuming, bootstrap, downtime, minReplicasToKeepUpForNoDowntime, enableStrictReplicaGroup, lowDiskMode, bestEfforts, externalViewCheckIntervalInMs, externalViewStabilizationTimeoutInMs); + if (summary && !dryRun) { + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, + "Must enable dry-run mode to use summary mode", null, null, null, null, null); + } + // Perform pre-checks if enabled Map preChecksResult = null; if (preChecks) { @@ -205,7 +211,8 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb String errorMsg = String.format("Pre-checks can only be enabled in dry-run mode, not triggering rebalance for " + "table: %s with rebalanceJobId: %s", tableNameWithType, rebalanceJobId); LOGGER.error(errorMsg); - return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, null, null, null, null); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, null, null, null, null, + null); } if (_rebalancePreChecker != null) { preChecksResult = _rebalancePreChecker.check(rebalanceJobId, tableNameWithType, tableConfig); @@ -222,21 +229,21 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb "For rebalanceId: %s, caught exception while fetching IdealState for table: %s, aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, - "Caught exception while fetching IdealState: " + e, null, null, null, preChecksResult); + "Caught exception while fetching IdealState: " + e, null, null, null, preChecksResult, null); } if (currentIdealState == null) { onReturnFailure( String.format("For rebalanceId: %s, cannot find the IdealState for table: %s, aborting the rebalance", rebalanceJobId, tableNameWithType), null); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Cannot find the IdealState for table", - null, null, null, preChecksResult); + null, null, null, preChecksResult, null); } if (!currentIdealState.isEnabled() && !downtime) { onReturnFailure(String.format( "For rebalanceId: %s, cannot rebalance disabled table: %s without downtime, aborting the rebalance", rebalanceJobId, tableNameWithType), null); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, - "Cannot rebalance disabled table without downtime", null, null, null, preChecksResult); + "Cannot rebalance disabled table without downtime", null, null, null, preChecksResult, null); } LOGGER.info("For rebalanceId: {}, processing instance partitions for table: {}", rebalanceJobId, tableNameWithType); @@ -254,7 +261,8 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb "For rebalanceId: %s, caught exception while fetching/calculating instance partitions for table: %s, " + "aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, - "Caught exception while fetching/calculating instance partitions: " + e, null, null, null, preChecksResult); + "Caught exception while fetching/calculating instance partitions: " + e, null, null, null, preChecksResult, + null); } // Calculate instance partitions for tiers if configured @@ -273,7 +281,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb + "aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while fetching/calculating tier instance partitions: " + e, null, null, null, - preChecksResult); + preChecksResult, null); } LOGGER.info("For rebalanceId: {}, calculating the target assignment for table: {}", rebalanceJobId, @@ -291,7 +299,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb + "rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while calculating target assignment: " + e, instancePartitionsMap, - tierToInstancePartitionsMap, null, preChecksResult); + tierToInstancePartitionsMap, null, preChecksResult, null); } boolean segmentAssignmentUnchanged = currentAssignment.equals(targetAssignment); @@ -299,6 +307,11 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb + "segmentAssignmentUnchanged: {} for table: {}", rebalanceJobId, instancePartitionsUnchanged, tierInstancePartitionsUnchanged, segmentAssignmentUnchanged, tableNameWithType); + RebalanceSummaryResult summaryResult = null; + if (summary) { + summaryResult = calculateDryRunSummary(currentAssignment, targetAssignment, tableNameWithType, rebalanceJobId); + } + if (segmentAssignmentUnchanged) { LOGGER.info("Table: {} is already balanced", tableNameWithType); if (instancePartitionsUnchanged && tierInstancePartitionsUnchanged) { @@ -306,28 +319,38 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb String.format("For rebalanceId: %s, instance unchanged and table: %s is already balanced", rebalanceJobId, tableNameWithType)); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.NO_OP, "Table is already balanced", - instancePartitionsMap, tierToInstancePartitionsMap, targetAssignment, preChecksResult); + summary ? null : instancePartitionsMap, summary ? null : tierToInstancePartitionsMap, + summary ? null : targetAssignment, preChecksResult, summaryResult); } else { if (dryRun) { return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, - "Instance reassigned in dry-run mode, table is already balanced", instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + "Instance reassigned in dry-run mode, table is already balanced", + summary ? null : instancePartitionsMap, summary ? null : tierToInstancePartitionsMap, + summary ? null : targetAssignment, preChecksResult, summaryResult); } else { _tableRebalanceObserver.onSuccess( String.format("For rebalanceId: %s, instance reassigned but table: %s is already balanced", rebalanceJobId, tableNameWithType)); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, - "Instance reassigned, table is already balanced", instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult); + "Instance reassigned, table is already balanced", summary ? null : instancePartitionsMap, + summary ? null : tierToInstancePartitionsMap, summary ? null : targetAssignment, preChecksResult, + summaryResult); } } } + if (summary) { + LOGGER.info("For rebalanceId: {}, rebalancing table: {} in dry-run summary mode, returning the summary only", + rebalanceJobId, tableNameWithType); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, "Dry-run summary mode", null, + null, null, preChecksResult, summaryResult); + } + if (dryRun) { LOGGER.info("For rebalanceId: {}, rebalancing table: {} in dry-run mode, returning the target assignment", rebalanceJobId, tableNameWithType); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, "Dry-run mode", instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } if (downtime) { @@ -352,14 +375,14 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, "Success with downtime (replaced IdealState with the target segment assignment, ExternalView might not " + "reach the target segment assignment yet)", instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult); + targetAssignment, preChecksResult, summaryResult); } catch (Exception e) { onReturnFailure(String.format( "For rebalanceId: %s, caught exception while updating IdealState for table: %s, aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while updating IdealState: " + e, instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult); + targetAssignment, preChecksResult, summaryResult); } } @@ -391,7 +414,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb minReplicasToKeepUpForNoDowntime, tableNameWithType, numReplicas), null); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Illegal min available replicas config", instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult); + targetAssignment, preChecksResult, summaryResult); } minAvailableReplicas = minReplicasToKeepUpForNoDowntime; } else { @@ -432,12 +455,12 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb if (_tableRebalanceObserver.isStopped()) { return new RebalanceResult(rebalanceJobId, _tableRebalanceObserver.getStopStatus(), "Caught exception while waiting for ExternalView to converge: " + e, instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } _tableRebalanceObserver.onError(errorMsg); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while waiting for ExternalView to converge: " + e, instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } // Re-calculate the target assignment if IdealState changed while waiting for ExternalView to converge @@ -486,7 +509,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb + "aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while re-calculating the target assignment: " + e, instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } } else { LOGGER.info("For rebalanceId:{}, no state change found for segments to be moved, " @@ -510,7 +533,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.DONE, "Success with minAvailableReplicas: " + minAvailableReplicas + " (both IdealState and ExternalView should reach the target segment assignment)", - instancePartitionsMap, tierToInstancePartitionsMap, targetAssignment, preChecksResult); + instancePartitionsMap, tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } // Record change of current ideal state and the new target @@ -519,7 +542,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb if (_tableRebalanceObserver.isStopped()) { return new RebalanceResult(rebalanceJobId, _tableRebalanceObserver.getStopStatus(), "Rebalance has stopped already before updating the IdealState", instancePartitionsMap, - tierToInstancePartitionsMap, targetAssignment, preChecksResult); + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } Map> nextAssignment = getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, @@ -550,7 +573,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb + "aborting the rebalance", rebalanceJobId, tableNameWithType), e); return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, "Caught exception while updating IdealState: " + e, instancePartitionsMap, tierToInstancePartitionsMap, - targetAssignment, preChecksResult); + targetAssignment, preChecksResult, summaryResult); } segmentsToMonitor = new HashSet<>(segmentsToMove); @@ -559,6 +582,88 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb } } + private RebalanceSummaryResult calculateDryRunSummary(Map> currentAssignment, + Map> targetAssignment, String tableNameWithType, String rebalanceJobId) { + LOGGER.info("Calculating rebalance summary for table: {} with rebalanceJobId: {}", tableNameWithType, + rebalanceJobId); + int existingReplicationFactor = 0; + int newReplicationFactor = 0; + Map> existingServersToSegmentMap = new HashMap<>(); + Map> newServersToSegmentMap = new HashMap<>(); + + for (Map.Entry> entrySet : currentAssignment.entrySet()) { + existingReplicationFactor = entrySet.getValue().size(); + for (Map.Entry segmentEntrySet : entrySet.getValue().entrySet()) { + existingServersToSegmentMap.putIfAbsent(segmentEntrySet.getKey(), new HashSet<>()); + existingServersToSegmentMap.get(segmentEntrySet.getKey()).add(entrySet.getKey()); + } + } + + for (Map.Entry> entrySet : targetAssignment.entrySet()) { + newReplicationFactor = entrySet.getValue().size(); + for (Map.Entry segmentEntrySet : entrySet.getValue().entrySet()) { + newServersToSegmentMap.putIfAbsent(segmentEntrySet.getKey(), new HashSet<>()); + newServersToSegmentMap.get(segmentEntrySet.getKey()).add(entrySet.getKey()); + } + } + RebalanceSummaryResult.RebalanceChangeInfo replicationFactor + = new RebalanceSummaryResult.RebalanceChangeInfo(existingReplicationFactor, newReplicationFactor); + + int existingNumServers = existingServersToSegmentMap.keySet().size(); + int newNumServers = newServersToSegmentMap.keySet().size(); + RebalanceSummaryResult.RebalanceChangeInfo numServers + = new RebalanceSummaryResult.RebalanceChangeInfo(existingNumServers, newNumServers); + + Map serverSegmentChangeInfoMap = new HashMap<>(); + int segmentsNotMoved = 0; + for (Map.Entry> entry : newServersToSegmentMap.entrySet()) { + String server = entry.getKey(); + Set segmentMap = entry.getValue(); + int totalNewSegments = segmentMap.size(); + + Set newSegmentList = new HashSet<>(segmentMap); + Set existingSegmentList = new HashSet<>(); + int segmentsUnchanged = 0; + int totalExistingSegments = 0; + if (existingServersToSegmentMap.containsKey(server)) { + totalExistingSegments = existingServersToSegmentMap.get(server).size(); + existingSegmentList.addAll(existingServersToSegmentMap.get(server)); + Set intersection = new HashSet<>(existingServersToSegmentMap.get(server)); + intersection.retainAll(newSegmentList); + segmentsUnchanged = intersection.size(); + segmentsNotMoved += segmentsUnchanged; + } + newSegmentList.removeAll(existingSegmentList); + int segmentsAdded = newSegmentList.size(); + int segmentsDeleted = existingSegmentList.size() - segmentsUnchanged; + + serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(totalNewSegments, + totalExistingSegments, segmentsAdded, segmentsDeleted, segmentsUnchanged)); + } + + for (Map.Entry> entry : existingServersToSegmentMap.entrySet()) { + String server = entry.getKey(); + if (!serverSegmentChangeInfoMap.containsKey(server)) { + serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(0, + entry.getValue().size(), 0, entry.getValue().size(), 0)); + } + } + + RebalanceSummaryResult.RebalanceChangeInfo uniqueNumSegments + = new RebalanceSummaryResult.RebalanceChangeInfo(currentAssignment.size(), targetAssignment.size()); + + int existingNumberSegmentsTotal = existingReplicationFactor * currentAssignment.size(); + int newNumberSegmentsTotal = newReplicationFactor * targetAssignment.size(); + RebalanceSummaryResult.RebalanceChangeInfo totalNumSegments + = new RebalanceSummaryResult.RebalanceChangeInfo(existingNumberSegmentsTotal, newNumberSegmentsTotal); + + int totalSegmentsToBeMoved = newNumberSegmentsTotal - segmentsNotMoved; + LOGGER.info("Calculated rebalance summary for table: {} with rebalanceJobId: {}", tableNameWithType, + rebalanceJobId); + return new RebalanceSummaryResult(totalSegmentsToBeMoved, numServers, replicationFactor, uniqueNumSegments, + totalNumSegments, serverSegmentChangeInfoMap); + } + private void onReturnFailure(String errorMsg, Exception e) { if (e != null) { LOGGER.warn(errorMsg, e); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java index e11547652941..5a45bf836512 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java @@ -60,7 +60,7 @@ public TenantRebalanceResult rebalance(TenantRebalanceConfig config) { false)); } catch (TableNotFoundException exception) { rebalanceResult.put(table, new RebalanceResult(null, RebalanceResult.Status.FAILED, exception.getMessage(), - null, null, null, null)); + null, null, null, null, null)); } }); if (config.isDryRun()) { @@ -71,7 +71,8 @@ public TenantRebalanceResult rebalance(TenantRebalanceConfig config) { if (result.getStatus() == RebalanceResult.Status.DONE) { rebalanceResult.put(table, new RebalanceResult(result.getJobId(), RebalanceResult.Status.IN_PROGRESS, "In progress, check controller task status for the", result.getInstanceAssignment(), - result.getTierInstanceAssignment(), result.getSegmentAssignment(), result.getPreChecksResult())); + result.getTierInstanceAssignment(), result.getSegmentAssignment(), result.getRebalanceSummaryResult(), + result.getPreChecksResult())); } } } 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 de83ec6eee7b..3e284b9f87df 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 @@ -100,6 +100,15 @@ public void testRebalance() // Rebalance should fail without creating the table RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + + // Rebalance with dry-run summary should fail without creating the table + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + assertNull(rebalanceResult.getRebalanceSummaryResult()); // Create the table addDummySchema(RAW_TABLE_NAME); @@ -114,6 +123,24 @@ public void testRebalance() Map> oldSegmentAssignment = _helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields(); + // Rebalance with dry-run summary should return NO_OP status + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + + // Dry-run mode should not change the IdealState + assertEquals(_helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields(), + oldSegmentAssignment); + // Rebalance should return NO_OP status rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); @@ -137,8 +164,50 @@ public void testRebalance() addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX + (numServers + i), true); } + // Rebalance in dry-run summary mode with added servers + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 14); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + + Map serverSegmentChangeInfoMap = + rebalanceResult.getRebalanceSummaryResult().getServerSegmentChangeInfo(); + assertNotNull(serverSegmentChangeInfoMap); + for (int i = 0; i < numServers; i++) { + // Original servers should be losing some segments + String newServer = SERVER_INSTANCE_ID_PREFIX + i; + RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); + assertTrue(serverSegmentChange._segmentsDeleted > 0); + assertTrue(serverSegmentChange._segmentsUnchanged > 0); + assertTrue(serverSegmentChange._totalExistingSegments > 0); + assertTrue(serverSegmentChange._totalNewSegments > 0); + assertEquals(serverSegmentChange._segmentsAdded, 0); + } + for (int i = 0; i < numServersToAdd; i++) { + // New servers should only get new segments + String newServer = SERVER_INSTANCE_ID_PREFIX + (numServers + i); + RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); + assertTrue(serverSegmentChange._segmentsAdded > 0); + assertEquals(serverSegmentChange._totalExistingSegments, 0); + assertEquals(serverSegmentChange._totalNewSegments, serverSegmentChange._segmentsAdded); + assertEquals(serverSegmentChange._segmentsDeleted, 0); + assertEquals(serverSegmentChange._segmentsUnchanged, 0); + } + + // Dry-run mode should not change the IdealState + assertEquals(_helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields(), + oldSegmentAssignment); + // Rebalance in dry-run mode - RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setDryRun(true); rebalanceConfig.setPreChecks(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); @@ -186,6 +255,21 @@ public void testRebalance() assertEquals(_helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields(), oldSegmentAssignment); + // Rebalance dry-run summary with 3 min available replicas should not be impacted since actual rebalance does not + // occur + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceConfig.setMinAvailableReplicas(3); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + // Rebalance with 3 min available replicas should fail as the table only have 3 replicas rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setMinAvailableReplicas(3); @@ -218,6 +302,38 @@ public void testRebalance() new InstanceAssignmentConfig(tagPoolConfig, null, replicaGroupPartitionConfig, null, false))); _helixResourceManager.updateTableConfig(tableConfig); + // Try dry-run summary mode + // No need to reassign instances because instances should be automatically assigned when updating the table config + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 11); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + + serverSegmentChangeInfoMap = rebalanceResult.getRebalanceSummaryResult().getServerSegmentChangeInfo(); + assertNotNull(serverSegmentChangeInfoMap); + for (int i = 0; i < numServers + numServersToAdd; i++) { + String newServer = SERVER_INSTANCE_ID_PREFIX + i; + RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); + assertEquals(serverSegmentChange._totalNewSegments, 5); + // Ensure not all segments moved + assertTrue(serverSegmentChange._segmentsUnchanged > 0); + // Ensure all segments has something assigned prior to rebalance + assertTrue(serverSegmentChange._totalExistingSegments > 0); + } + + // Dry-run mode should not change the IdealState + assertEquals(_helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields(), + newSegmentAssignment); + + // Try actual rebalance // No need to reassign instances because instances should be automatically assigned when updating the table config rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); @@ -264,6 +380,21 @@ public void testRebalance() tableConfig.setInstanceAssignmentConfigMap(null); _helixResourceManager.updateTableConfig(tableConfig); + // Try dry-run summary mode without reassignment to ensure that existing instance partitions are used + // no movement should occur + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + // Without instances reassignment, the rebalance should return status NO_OP, and the existing instance partitions // should be used rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -271,6 +402,22 @@ public void testRebalance() assertEquals(rebalanceResult.getInstanceAssignment(), instanceAssignment); assertEquals(rebalanceResult.getSegmentAssignment(), newSegmentAssignment); + // Try dry-run summary mode with reassignment + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceConfig.setReassignInstances(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + // No move expected since already balanced + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + // With instances reassignment, the rebalance should return status DONE, the existing instance partitions should be // removed, and the default instance partitions should be used rebalanceConfig = new RebalanceConfig(); @@ -300,6 +447,21 @@ public void testRebalance() TagNameUtils.getOfflineTagForTenant(null)); } + // Try dry-run summary mode + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceConfig.setDowntime(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + // Rebalance with downtime should succeed rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setDowntime(true); @@ -327,6 +489,16 @@ public void testRebalance() } } + // Try summary mode without dry-run set + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + _helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME); } @@ -367,7 +539,23 @@ public void testRebalanceWithTiers() _helixResourceManager.getTableIdealState(OFFLINE_TIERED_TABLE_NAME).getRecord().getMapFields(); TableRebalancer tableRebalancer = new TableRebalancer(_helixManager); - RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); + + // Try dry-run summary mode + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + + // Run actual table rebalance + rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); // Segment assignment should not change assertEquals(rebalanceResult.getSegmentAssignment(), oldSegmentAssignment); @@ -382,6 +570,20 @@ public void testRebalanceWithTiers() } _helixResourceManager.createServerTenant(new Tenant(TenantRole.SERVER, TIER_B_NAME, 3, 3, 0)); + // Try dry-run summary mode, should be no-op + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); @@ -400,6 +602,20 @@ public void testRebalanceWithTiers() TierFactory.PINOT_SERVER_STORAGE_TYPE, NO_TIER_NAME + "_OFFLINE", null, null))); _helixResourceManager.updateTableConfig(tableConfig); + // Try dry-run summary mode, some moves should occur + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 9); + // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); @@ -458,7 +674,22 @@ public void testRebalanceWithTiersAndInstanceAssignments() _helixResourceManager.getTableIdealState(OFFLINE_TIERED_TABLE_NAME).getRecord().getMapFields(); TableRebalancer tableRebalancer = new TableRebalancer(_helixManager); - RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); + + // Try dry-run summary mode + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + + rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); // Segment assignment should not change assertEquals(rebalanceResult.getSegmentAssignment(), oldSegmentAssignment); @@ -469,6 +700,21 @@ public void testRebalanceWithTiersAndInstanceAssignments() "replicaAssignment" + TIER_A_NAME + "_" + SERVER_INSTANCE_ID_PREFIX + i, false); } _helixResourceManager.createServerTenant(new Tenant(TenantRole.SERVER, "replicaAssignment" + TIER_A_NAME, 6, 6, 0)); + + // Try dry-run summary mode + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); @@ -481,6 +727,20 @@ public void testRebalanceWithTiersAndInstanceAssignments() TierFactory.PINOT_SERVER_STORAGE_TYPE, "replicaAssignment" + TIER_A_NAME + "_OFFLINE", null, null))); _helixResourceManager.updateTableConfig(tableConfig); + // Try dry-run summary mode + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 30); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); @@ -504,6 +764,20 @@ public void testRebalanceWithTiersAndInstanceAssignments() new InstanceAssignmentConfig(tagPoolConfig, null, replicaGroupPartitionConfig, null, false))); _helixResourceManager.updateTableConfig(tableConfig); + // Try dry-run summary mode + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNull(rebalanceResult.getInstanceAssignment()); + assertNull(rebalanceResult.getTierInstanceAssignment()); + assertNull(rebalanceResult.getSegmentAssignment()); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 13); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); assertTrue(rebalanceResult.getTierInstanceAssignment().containsKey(TIER_A_NAME)); From 6085faa44316768856872ea73e744c57330daa0f Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Thu, 13 Feb 2025 13:46:51 -0800 Subject: [PATCH 2/7] Add table size option --- .../controller/BaseControllerStarter.java | 1 + .../helix/core/PinotHelixResourceManager.java | 18 ++- .../rebalance/RebalanceSummaryResult.java | 44 +++++- .../helix/core/rebalance/TableRebalancer.java | 61 ++++++- .../tenant/DefaultTenantRebalancer.java | 4 +- .../tenant/TenantRebalanceResult.java | 2 +- .../TableRebalancerClusterStatelessTest.java | 13 +- .../tests/OfflineClusterIntegrationTest.java | 149 +++++++++++++++++- 8 files changed, 274 insertions(+), 18 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java b/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java index 5761cb27e5ec..9a4d080c7c22 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java @@ -512,6 +512,7 @@ private void setUpPinotController() { _tableSizeReader = new TableSizeReader(_executorService, _connectionManager, _controllerMetrics, _helixResourceManager, _leadControllerManager); + _helixResourceManager.registerTableSizeReader(_tableSizeReader); _storageQuotaChecker = new StorageQuotaChecker(_tableSizeReader, _controllerMetrics, _leadControllerManager, _helixResourceManager, _config); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java index 748f9d3c286e..b30a117bec2f 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java @@ -159,6 +159,7 @@ import org.apache.pinot.controller.helix.core.rebalance.TableRebalancer; import org.apache.pinot.controller.helix.core.rebalance.ZkBasedTableRebalanceObserver; import org.apache.pinot.controller.helix.starter.HelixConfig; +import org.apache.pinot.controller.util.TableSizeReader; import org.apache.pinot.segment.spi.SegmentMetadata; import org.apache.pinot.spi.config.DatabaseConfig; import org.apache.pinot.spi.config.instance.Instance; @@ -241,6 +242,7 @@ private enum LineageUpdateType { private TableCache _tableCache; private final LineageManager _lineageManager; private final RebalancePreChecker _rebalancePreChecker; + private TableSizeReader _tableSizeReader; public PinotHelixResourceManager(String zkURL, String helixClusterName, @Nullable String dataDir, boolean isSingleTenantCluster, boolean enableBatchMessageMode, int deletedSegmentsRetentionInDays, @@ -449,6 +451,15 @@ public RebalancePreChecker getRebalancePreChecker() { } /** + * Get the table size reader. + * + * @return table size reader + */ + public TableSizeReader getTableSizeReader() { + return _tableSizeReader; + } + +/** * Instance related APIs */ @@ -1943,6 +1954,10 @@ public void registerPinotLLCRealtimeSegmentManager(PinotLLCRealtimeSegmentManage _pinotLLCRealtimeSegmentManager = pinotLLCRealtimeSegmentManager; } + public void registerTableSizeReader(TableSizeReader tableSizeReader) { + _tableSizeReader = tableSizeReader; + } + private void assignInstances(TableConfig tableConfig, boolean override) { String tableNameWithType = tableConfig.getTableName(); String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); @@ -3612,7 +3627,8 @@ public RebalanceResult rebalanceTable(String tableNameWithType, TableConfig tabl tierToSegmentsMap = updateTargetTier(rebalanceJobId, tableNameWithType, tableConfig); } TableRebalancer tableRebalancer = - new TableRebalancer(_helixZkManager, zkBasedTableRebalanceObserver, _controllerMetrics, _rebalancePreChecker); + new TableRebalancer(_helixZkManager, zkBasedTableRebalanceObserver, _controllerMetrics, _rebalancePreChecker, + _tableSizeReader); return tableRebalancer.rebalance(tableConfig, rebalanceConfig, rebalanceJobId, tierToSegmentsMap); } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java index f47d38b61f09..da77ff562522 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -67,27 +67,39 @@ public RebalanceChangeInfo(@JsonProperty("existingValue") int existingValue, // TODO: Add stats about total data size to be moved and estimated calculations on how long the data move can take // during rebalance. Estimations are fine based on total table size / total number of segments - private int _totalSegmentsToBeMoved; + private final int _totalSegmentsToBeMoved; + private final int _numServersGettingNewSegments; + private final long _estimatedAverageSegmentSizeInBytes; + private final long _totalEstimatedDataToBeMovedInBytes; + private final double _totalEstimatedTimeToMoveDataInSecs; @JsonInclude(JsonInclude.Include.NON_NULL) - private RebalanceChangeInfo _numServers; + private final RebalanceChangeInfo _numServers; @JsonInclude(JsonInclude.Include.NON_NULL) - private RebalanceChangeInfo _replicationFactor; + private final RebalanceChangeInfo _replicationFactor; @JsonInclude(JsonInclude.Include.NON_NULL) - private RebalanceChangeInfo _numUniqueSegments; + private final RebalanceChangeInfo _numUniqueSegments; @JsonInclude(JsonInclude.Include.NON_NULL) - private RebalanceChangeInfo _numTotalSegments; + private final RebalanceChangeInfo _numTotalSegments; @JsonInclude(JsonInclude.Include.NON_NULL) - private Map _serverSegmentChangeInfo; + private final Map _serverSegmentChangeInfo; @JsonCreator public RebalanceSummaryResult( @JsonProperty("totalSegmentsToBeMoved") int totalSegmentsToBeMoved, + @JsonProperty("numServersGettingNewSegments") int numServersGettingNewSegments, + @JsonProperty("estimatedAverageSegmentSizeInBytes") long estimatedAverageSegmentSizeInBytes, + @JsonProperty("totalEstimatedDataToBeMovedInBytes") long totalEstimatedDataToBeMovedInBytes, + @JsonProperty("totalEstimatedTimeToMoveDataInSecs") double totalEstimatedTimeToMoveDataInSecs, @JsonProperty("numServers") @Nullable RebalanceChangeInfo numServers, @JsonProperty("replicationFactor") @Nullable RebalanceChangeInfo replicationFactor, @JsonProperty("numUniqueSegments") @Nullable RebalanceChangeInfo numUniqueSegments, @JsonProperty("numTotalSegments") @Nullable RebalanceChangeInfo numTotalSegments, @JsonProperty("serverSegmentChangeInfo") @Nullable Map serverSegmentChangeInfo) { _totalSegmentsToBeMoved = totalSegmentsToBeMoved; + _numServersGettingNewSegments = numServersGettingNewSegments; + _estimatedAverageSegmentSizeInBytes = estimatedAverageSegmentSizeInBytes; + _totalEstimatedDataToBeMovedInBytes = totalEstimatedDataToBeMovedInBytes; + _totalEstimatedTimeToMoveDataInSecs = totalEstimatedTimeToMoveDataInSecs; _numServers = numServers; _replicationFactor = replicationFactor; _numUniqueSegments = numUniqueSegments; @@ -100,6 +112,26 @@ public int getTotalSegmentsToBeMoved() { return _totalSegmentsToBeMoved; } + @JsonProperty + public int getNumServersGettingNewSegments() { + return _numServersGettingNewSegments; + } + + @JsonProperty + public long getEstimatedAverageSegmentSizeInBytes() { + return _estimatedAverageSegmentSizeInBytes; + } + + @JsonProperty + public long getTotalEstimatedDataToBeMovedInBytes() { + return _totalEstimatedDataToBeMovedInBytes; + } + + @JsonProperty + public double getTotalEstimatedTimeToMoveDataInSecs() { + return _totalEstimatedTimeToMoveDataInSecs; + } + @JsonProperty public RebalanceChangeInfo getNumServers() { return _numServers; 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 abacfb759b7c..29819ae650a5 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 @@ -48,6 +48,7 @@ import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils; import org.apache.pinot.common.assignment.InstancePartitions; import org.apache.pinot.common.assignment.InstancePartitionsUtils; +import org.apache.pinot.common.exception.InvalidConfigException; import org.apache.pinot.common.metrics.ControllerMetrics; import org.apache.pinot.common.metrics.ControllerTimer; import org.apache.pinot.common.tier.PinotServerTierStorage; @@ -60,6 +61,7 @@ import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentFactory; import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils; import org.apache.pinot.controller.helix.core.assignment.segment.StrictRealtimeSegmentAssignment; +import org.apache.pinot.controller.util.TableSizeReader; import org.apache.pinot.spi.config.table.RoutingConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -120,9 +122,11 @@ public class TableRebalancer { private final TableRebalanceObserver _tableRebalanceObserver; private final ControllerMetrics _controllerMetrics; private final RebalancePreChecker _rebalancePreChecker; + private final TableSizeReader _tableSizeReader; public TableRebalancer(HelixManager helixManager, @Nullable TableRebalanceObserver tableRebalanceObserver, - @Nullable ControllerMetrics controllerMetrics, @Nullable RebalancePreChecker rebalancePreChecker) { + @Nullable ControllerMetrics controllerMetrics, @Nullable RebalancePreChecker rebalancePreChecker, + @Nullable TableSizeReader tableSizeReader) { _helixManager = helixManager; if (tableRebalanceObserver != null) { _tableRebalanceObserver = tableRebalanceObserver; @@ -132,10 +136,11 @@ public TableRebalancer(HelixManager helixManager, @Nullable TableRebalanceObserv _helixDataAccessor = helixManager.getHelixDataAccessor(); _controllerMetrics = controllerMetrics; _rebalancePreChecker = rebalancePreChecker; + _tableSizeReader = tableSizeReader; } public TableRebalancer(HelixManager helixManager) { - this(helixManager, null, null, null); + this(helixManager, null, null, null, null); } public static String createUniqueRebalanceJobIdentifier() { @@ -582,10 +587,28 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb } } + private long calculateTableSizePerReplicaInBytes(String tableNameWithType) { + long tableSizePerReplicaInBytes = -1; + if (_tableSizeReader == null) { + LOGGER.warn("tableSizeReader is null, cannot calculate table size for table {}!", tableNameWithType); + return tableSizePerReplicaInBytes; + } + try { + TableSizeReader.TableSubTypeSizeDetails tableSizeDetails = + _tableSizeReader.getTableSubtypeSize(tableNameWithType, 30_000); + tableSizePerReplicaInBytes = tableSizeDetails._reportedSizePerReplicaInBytes; + } catch (InvalidConfigException e) { + String errMsg = String.format("Caught exception while trying to fetch table size details for table: %s", + tableNameWithType); + LOGGER.error(errMsg, e); + } + return tableSizePerReplicaInBytes; + } + private RebalanceSummaryResult calculateDryRunSummary(Map> currentAssignment, Map> targetAssignment, String tableNameWithType, String rebalanceJobId) { - LOGGER.info("Calculating rebalance summary for table: {} with rebalanceJobId: {}", tableNameWithType, - rebalanceJobId); + LOGGER.info("Calculating rebalance summary for table: {} with rebalanceJobId: {}", + tableNameWithType, rebalanceJobId); int existingReplicationFactor = 0; int newReplicationFactor = 0; Map> existingServersToSegmentMap = new HashMap<>(); @@ -616,6 +639,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map serverSegmentChangeInfoMap = new HashMap<>(); int segmentsNotMoved = 0; + int numServersGettingDataAdded = 0; for (Map.Entry> entry : newServersToSegmentMap.entrySet()) { String server = entry.getKey(); Set segmentMap = entry.getValue(); @@ -636,6 +660,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0 ? 1 : 0; serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(totalNewSegments, totalExistingSegments, segmentsAdded, segmentsDeleted, segmentsUnchanged)); @@ -658,10 +683,34 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0 && numServersGettingDataAdded > 0) { + // Do some processing to figure out what the time to download might be + // Assume that data is evenly distributed across all servers + long totalDataPerServerInBytes = totalEstimatedDataToBeMovedInBytes / numServersGettingDataAdded; + // TODO: Pick a good threshold to calculate estimated time to rebalance. For now assume 100 MB/s data download + // + process rate + estimatedTimeToRebalanceInSec = ((double) totalDataPerServerInBytes) / (100.0D * 1024.0D * 1024.0D); + } + return estimatedTimeToRebalanceInSec; } private void onReturnFailure(String errorMsg, Exception e) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java index 5a45bf836512..661fff70d918 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java @@ -71,8 +71,8 @@ public TenantRebalanceResult rebalance(TenantRebalanceConfig config) { if (result.getStatus() == RebalanceResult.Status.DONE) { rebalanceResult.put(table, new RebalanceResult(result.getJobId(), RebalanceResult.Status.IN_PROGRESS, "In progress, check controller task status for the", result.getInstanceAssignment(), - result.getTierInstanceAssignment(), result.getSegmentAssignment(), result.getRebalanceSummaryResult(), - result.getPreChecksResult())); + result.getTierInstanceAssignment(), result.getSegmentAssignment(), result.getPreChecksResult(), + result.getRebalanceSummaryResult())); } } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalanceResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalanceResult.java index c96c25b2509f..ed0caf0f2933 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalanceResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalanceResult.java @@ -36,7 +36,7 @@ public TenantRebalanceResult(String jobId, Map rebalanc _rebalanceTableResults = new HashMap<>(); rebalanceTableResults.forEach((table, result) -> { _rebalanceTableResults.put(table, new RebalanceResult(result.getJobId(), result.getStatus(), - result.getDescription(), null, null, null, null)); + result.getDescription(), null, null, null, null, null)); }); } } 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 3e284b9f87df..d3687015342c 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 @@ -93,7 +93,8 @@ public void testRebalance() DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker(); preChecker.init(_helixResourceManager, Executors.newFixedThreadPool(10)); - TableRebalancer tableRebalancer = new TableRebalancer(_helixManager, null, null, preChecker); + TableRebalancer tableRebalancer = new TableRebalancer(_helixManager, null, null, preChecker, + _helixResourceManager.getTableSizeReader()); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(); @@ -177,6 +178,7 @@ public void testRebalance() assertNull(rebalanceResult.getSegmentAssignment()); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 3); Map serverSegmentChangeInfoMap = rebalanceResult.getRebalanceSummaryResult().getServerSegmentChangeInfo(); @@ -394,6 +396,7 @@ public void testRebalance() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); // Without instances reassignment, the rebalance should return status NO_OP, and the existing instance partitions // should be used @@ -417,6 +420,7 @@ public void testRebalance() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); // With instances reassignment, the rebalance should return status DONE, the existing instance partitions should be // removed, and the default instance partitions should be used @@ -461,6 +465,7 @@ public void testRebalance() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 3); // Rebalance with downtime should succeed rebalanceConfig = new RebalanceConfig(); @@ -553,6 +558,7 @@ public void testRebalanceWithTiers() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); // Run actual table rebalance rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -583,6 +589,7 @@ public void testRebalanceWithTiers() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -615,6 +622,7 @@ public void testRebalanceWithTiers() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 9); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 6); // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -688,6 +696,7 @@ public void testRebalanceWithTiersAndInstanceAssignments() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); @@ -714,6 +723,7 @@ public void testRebalanceWithTiersAndInstanceAssignments() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -740,6 +750,7 @@ public void testRebalanceWithTiersAndInstanceAssignments() assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 30); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 6); // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index 6750a7340d7a..d48dacf966ba 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -24,6 +24,8 @@ import java.io.File; import java.io.IOException; import java.net.URI; +import java.net.URISyntaxException; +import java.net.URL; import java.sql.ResultSet; import java.sql.Statement; import java.sql.Timestamp; @@ -64,10 +66,12 @@ import org.apache.pinot.common.utils.ServiceStatus; import org.apache.pinot.common.utils.SimpleHttpResponse; import org.apache.pinot.common.utils.http.HttpClient; +import org.apache.pinot.controller.api.resources.ServerRebalanceJobStatusResponse; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.helix.core.rebalance.DefaultRebalancePreChecker; import org.apache.pinot.controller.helix.core.rebalance.RebalanceConfig; import org.apache.pinot.controller.helix.core.rebalance.RebalanceResult; +import org.apache.pinot.controller.helix.core.rebalance.RebalanceSummaryResult; import org.apache.pinot.controller.helix.core.rebalance.TableRebalancer; import org.apache.pinot.core.operator.query.NonScanBasedAggregationOperator; import org.apache.pinot.segment.spi.index.ForwardIndexConfig; @@ -288,7 +292,8 @@ public void setUp() _resourceManager = _controllerStarter.getHelixResourceManager(); DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker(); preChecker.init(_helixResourceManager, Executors.newFixedThreadPool(10)); - _tableRebalancer = new TableRebalancer(_resourceManager.getHelixZkManager(), null, null, preChecker); + _tableRebalancer = new TableRebalancer(_resourceManager.getHelixZkManager(), null, null, preChecker, + _resourceManager.getTableSizeReader()); } private void reloadAllSegments(String testQuery, boolean forceDownload, long numTotalDocs) @@ -4018,4 +4023,146 @@ public void testResponseWithClientRequestId(boolean useMultiStageQueryEngine) assertEquals(result.get("clientRequestId").asText(), clientRequestId); } + + @Test + public void testRebalanceDryRunSummary() + throws Exception { + // setup the rebalance config + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDryRun(true); + + TableConfig tableConfig = getOfflineTableConfig(); + + // Ensure summary status is null if not enabled + RebalanceResult rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + + // Enable summary, nothing is set + rebalanceConfig.setSummary(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.NO_OP, false, getNumServers(), getNumServers(), + tableConfig.getReplication()); + + // Add a new server (to force change in instance assignment) and enable reassignInstances + BaseServerStarter serverStarter1 = startOneServer(NUM_SERVERS); + rebalanceConfig.setReassignInstances(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.DONE, true, getNumServers(), + getNumServers() + 1, tableConfig.getReplication()); + + // Disable dry-run, summary still enabled + rebalanceConfig.setDryRun(false); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + + // Disable summary along with dry-run to do a real rebalance + rebalanceConfig.setSummary(false); + rebalanceConfig.setDowntime(true); + rebalanceConfig.setReassignInstances(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + + waitForRebalanceToComplete(rebalanceResult, 600_000L); + + // Untag the added server + _resourceManager.updateInstanceTags(serverStarter1.getInstanceId(), "", false); + + // Re-enable dry-run + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.DONE, true, getNumServers() + 1, + getNumServers(), tableConfig.getReplication()); + + // Disable dry-run and summary to do a real rebalance + rebalanceConfig.setDryRun(false); + rebalanceConfig.setSummary(false); + rebalanceConfig.setDowntime(true); + rebalanceConfig.setReassignInstances(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertNull(rebalanceResult.getRebalanceSummaryResult()); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + + waitForRebalanceToComplete(rebalanceResult, 600_000L); + + // Stop the server + serverStarter1.stop(); + TestUtils.waitForCondition(aVoid -> _resourceManager.dropInstance(serverStarter1.getInstanceId()).isSuccessful(), + 60_000L, "Failed to drop added server"); + + // Try dry-run with summary again + rebalanceConfig.setDryRun(true); + rebalanceConfig.setSummary(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.NO_OP, false, getNumServers(), getNumServers(), + tableConfig.getReplication()); + } + + private void checkRebalanceDryRunSummary(RebalanceResult rebalanceResult, RebalanceResult.Status expectedStatus, + boolean isSegmentsToBeMoved, int existingNumServers, int newNumServers, int replicationFactor) { + assertEquals(rebalanceResult.getStatus(), expectedStatus); + RebalanceSummaryResult summaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(summaryResult); + assertEquals(summaryResult.getReplicationFactor()._existingValue, replicationFactor, + "Existing replication factor doesn't match expected"); + assertEquals(summaryResult.getReplicationFactor()._existingValue, summaryResult.getReplicationFactor()._newValue, + "Existing and new replication factor doesn't match"); + assertEquals(summaryResult.getNumServers()._existingValue, existingNumServers, + "Existing number of servers don't match"); + assertEquals(summaryResult.getNumServers()._newValue, newNumServers, "New number of servers don't match"); + if (_tableSize > 0) { + assertTrue(summaryResult.getEstimatedAverageSegmentSizeInBytes() > 0L, + "Avg segment size expected to be > 0 but found to be 0"); + } + if (existingNumServers != newNumServers) { + assertTrue(summaryResult.getNumServersGettingNewSegments() > 0, "Expected number of servers should be > 0"); + } else { + assertEquals(summaryResult.getNumServersGettingNewSegments(), 0, + "Expected number of servers getting new segments should be 0"); + } + + if (isSegmentsToBeMoved) { + assertTrue(summaryResult.getTotalSegmentsToBeMoved() > 0, "Segments to be moved should be > 0"); + assertEquals(summaryResult.getTotalEstimatedDataToBeMovedInBytes(), + summaryResult.getTotalSegmentsToBeMoved() * summaryResult.getEstimatedAverageSegmentSizeInBytes(), + "Estimated data to be moved in bytes doesn't match"); + assertTrue(summaryResult.getTotalEstimatedTimeToMoveDataInSecs() > 0.0D, + "Estimated time to move segments should be more than 0.0 seconds"); + } else { + assertEquals(summaryResult.getTotalSegmentsToBeMoved(), 0, "Segments to be moved should be 0"); + assertEquals(summaryResult.getTotalEstimatedDataToBeMovedInBytes(), 0L, + "Estimated data to be moved in bytes should be 0"); + assertEquals(summaryResult.getTotalEstimatedTimeToMoveDataInSecs(), 0.0D, + "Estimated time to move segments should be 0.0D"); + } + } + + protected void waitForRebalanceToComplete(RebalanceResult rebalanceResult, long timeoutMs) + throws Exception { + String jobId = rebalanceResult.getJobId(); + if (rebalanceResult.getStatus() != RebalanceResult.Status.IN_PROGRESS) { + return; + } + + TestUtils.waitForCondition(aVoid -> { + try { + String requestUrl = getControllerRequestURLBuilder().forTableRebalanceStatus(jobId); + try { + SimpleHttpResponse httpResponse = + HttpClient.wrapAndThrowHttpException(getHttpClient().sendGetRequest(new URL(requestUrl).toURI(), null)); + + ServerRebalanceJobStatusResponse serverRebalanceJobStatusResponse = + JsonUtils.stringToObject(httpResponse.getResponse(), ServerRebalanceJobStatusResponse.class); + RebalanceResult.Status status = serverRebalanceJobStatusResponse.getTableRebalanceProgressStats().getStatus(); + return status != RebalanceResult.Status.IN_PROGRESS; + } catch (HttpErrorStatusException | URISyntaxException e) { + throw new IOException(e); + } + } catch (Exception e) { + return null; + } + }, 1000L, timeoutMs, "Failed to load all segments after rebalance"); + } } From 35ba42c8e14991dc626a3567379f5a7ea25d0610 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Sun, 16 Feb 2025 16:29:37 -0800 Subject: [PATCH 3/7] Address review comments to improve summary --- .../rebalance/RebalanceSummaryResult.java | 300 +++++++++++++----- .../helix/core/rebalance/TableRebalancer.java | 34 +- .../TableRebalancerClusterStatelessTest.java | 205 +++++++----- .../tests/OfflineClusterIntegrationTest.java | 82 ++++- 4 files changed, 430 insertions(+), 191 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java index da77ff562522..e9a80242e1d7 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.List; import java.util.Map; import javax.annotation.Nullable; @@ -34,126 +35,249 @@ public class RebalanceSummaryResult { public static class ServerSegmentChangeInfo { - public final int _totalNewSegments; - public final int _totalExistingSegments; - public final int _segmentsAdded; - public final int _segmentsDeleted; - public final int _segmentsUnchanged; + private final ServerStatus _serverStatus; + private final int _totalSegmentsAfterRebalance; + private final int _totalSegmentsBeforeRebalance; + private final int _segmentsAdded; + private final int _segmentsDeleted; + private final int _segmentsUnchanged; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final List _tagList; + /** + * Constructor for ServerSegmentChangeInfo + * @param serverStatus server status, whether it was added, removed, or unchanged as part of this rebalance + * @param totalSegmentsAfterRebalance expected total segments on this server after rebalance + * @param totalSegmentsBeforeRebalance current number of segments on this server before rebalance + * @param segmentsAdded number of segments expected to be added as part of this rebalance + * @param segmentsDeleted number of segments expected to be deleted as part of this rebalance + * @param segmentsUnchanged number of segments that aren't moving from this server as part of this rebalance + * @param tagList server tag list + */ @JsonCreator - public ServerSegmentChangeInfo(@JsonProperty("totalNewSegments") int totalNewSegments, - @JsonProperty("totalExistingSegments") int totalExistingSegments, + public ServerSegmentChangeInfo(@JsonProperty("serverStatus") ServerStatus serverStatus, + @JsonProperty("totalSegmentsAfterRebalance") int totalSegmentsAfterRebalance, + @JsonProperty("totalSegmentsBeforeRebalance") int totalSegmentsBeforeRebalance, @JsonProperty("segmentsAdded") int segmentsAdded, @JsonProperty("segmentsDeleted") int segmentsDeleted, - @JsonProperty("segmentsUnchanged") int segmentsUnchanged) { - _totalNewSegments = totalNewSegments; - _totalExistingSegments = totalExistingSegments; + @JsonProperty("segmentsUnchanged") int segmentsUnchanged, + @JsonProperty("tagList") @Nullable List tagList) { + _serverStatus = serverStatus; + _totalSegmentsAfterRebalance = totalSegmentsAfterRebalance; + _totalSegmentsBeforeRebalance = totalSegmentsBeforeRebalance; _segmentsAdded = segmentsAdded; _segmentsDeleted = segmentsDeleted; _segmentsUnchanged = segmentsUnchanged; + _tagList = tagList; + } + + @JsonProperty + public ServerStatus getServerStatus() { + return _serverStatus; + } + + @JsonProperty + public int getTotalSegmentsAfterRebalance() { + return _totalSegmentsAfterRebalance; + } + + @JsonProperty + public int getTotalSegmentsBeforeRebalance() { + return _totalSegmentsBeforeRebalance; + } + + @JsonProperty + public int getSegmentsAdded() { + return _segmentsAdded; + } + + @JsonProperty + public int getSegmentsDeleted() { + return _segmentsDeleted; + } + + @JsonProperty + public int getSegmentsUnchanged() { + return _segmentsUnchanged; + } + + @JsonProperty + public List getTagList() { + return _tagList; } } public static class RebalanceChangeInfo { - public final int _existingValue; - public final int _newValue; + private final int _valueBeforeRebalance; + private final int _expectedValueAfterRebalance; + /** + * Constructor for RebalanceChangeInfo + * @param valueBeforeRebalance current value before rebalance + * @param expectedValueAfterRebalance expected value after rebalance + */ @JsonCreator - public RebalanceChangeInfo(@JsonProperty("existingValue") int existingValue, - @JsonProperty("newValue") int newValue) { - _existingValue = existingValue; - _newValue = newValue; + public RebalanceChangeInfo(@JsonProperty("valueBeforeRebalance") int valueBeforeRebalance, + @JsonProperty("expectedValueAfterRebalance") int expectedValueAfterRebalance) { + _valueBeforeRebalance = valueBeforeRebalance; + _expectedValueAfterRebalance = expectedValueAfterRebalance; } - } - // TODO: Add stats about total data size to be moved and estimated calculations on how long the data move can take - // during rebalance. Estimations are fine based on total table size / total number of segments - private final int _totalSegmentsToBeMoved; - private final int _numServersGettingNewSegments; - private final long _estimatedAverageSegmentSizeInBytes; - private final long _totalEstimatedDataToBeMovedInBytes; - private final double _totalEstimatedTimeToMoveDataInSecs; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final RebalanceChangeInfo _numServers; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final RebalanceChangeInfo _replicationFactor; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final RebalanceChangeInfo _numUniqueSegments; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final RebalanceChangeInfo _numTotalSegments; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final Map _serverSegmentChangeInfo; + @JsonProperty + public int getValueBeforeRebalance() { + return _valueBeforeRebalance; + } - @JsonCreator - public RebalanceSummaryResult( - @JsonProperty("totalSegmentsToBeMoved") int totalSegmentsToBeMoved, - @JsonProperty("numServersGettingNewSegments") int numServersGettingNewSegments, - @JsonProperty("estimatedAverageSegmentSizeInBytes") long estimatedAverageSegmentSizeInBytes, - @JsonProperty("totalEstimatedDataToBeMovedInBytes") long totalEstimatedDataToBeMovedInBytes, - @JsonProperty("totalEstimatedTimeToMoveDataInSecs") double totalEstimatedTimeToMoveDataInSecs, - @JsonProperty("numServers") @Nullable RebalanceChangeInfo numServers, - @JsonProperty("replicationFactor") @Nullable RebalanceChangeInfo replicationFactor, - @JsonProperty("numUniqueSegments") @Nullable RebalanceChangeInfo numUniqueSegments, - @JsonProperty("numTotalSegments") @Nullable RebalanceChangeInfo numTotalSegments, - @JsonProperty("serverSegmentChangeInfo") @Nullable Map serverSegmentChangeInfo) { - _totalSegmentsToBeMoved = totalSegmentsToBeMoved; - _numServersGettingNewSegments = numServersGettingNewSegments; - _estimatedAverageSegmentSizeInBytes = estimatedAverageSegmentSizeInBytes; - _totalEstimatedDataToBeMovedInBytes = totalEstimatedDataToBeMovedInBytes; - _totalEstimatedTimeToMoveDataInSecs = totalEstimatedTimeToMoveDataInSecs; - _numServers = numServers; - _replicationFactor = replicationFactor; - _numUniqueSegments = numUniqueSegments; - _numTotalSegments = numTotalSegments; - _serverSegmentChangeInfo = serverSegmentChangeInfo; + @JsonProperty + public int getExpectedValueAfterRebalance() { + return _expectedValueAfterRebalance; + } } - @JsonProperty - public int getTotalSegmentsToBeMoved() { - return _totalSegmentsToBeMoved; - } + public static class ServerInfo { + private final int _numServersGettingNewSegments; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RebalanceChangeInfo _numServers; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final Map _serverSegmentChangeInfo; - @JsonProperty - public int getNumServersGettingNewSegments() { - return _numServersGettingNewSegments; - } + /** + * Constructor for ServerInfo + * @param numServersGettingNewSegments total number of servers receiving new segments as part of this rebalance + * @param numServers number of servers before and after this rebalance + * @param serverSegmentChangeInfo per server statistics for this rebalance + */ + @JsonCreator + public ServerInfo(@JsonProperty("numServersGettingNewSegments") int numServersGettingNewSegments, + @JsonProperty("numServers") @Nullable RebalanceChangeInfo numServers, + @JsonProperty("serverSegmentChangeInfo") + @Nullable Map serverSegmentChangeInfo) { + _numServersGettingNewSegments = numServersGettingNewSegments; + _numServers = numServers; + _serverSegmentChangeInfo = serverSegmentChangeInfo; + } - @JsonProperty - public long getEstimatedAverageSegmentSizeInBytes() { - return _estimatedAverageSegmentSizeInBytes; - } + @JsonProperty + public int getNumServersGettingNewSegments() { + return _numServersGettingNewSegments; + } - @JsonProperty - public long getTotalEstimatedDataToBeMovedInBytes() { - return _totalEstimatedDataToBeMovedInBytes; - } + @JsonProperty + public RebalanceChangeInfo getNumServers() { + return _numServers; + } - @JsonProperty - public double getTotalEstimatedTimeToMoveDataInSecs() { - return _totalEstimatedTimeToMoveDataInSecs; + @JsonProperty + public Map getServerSegmentChangeInfo() { + return _serverSegmentChangeInfo; + } } - @JsonProperty - public RebalanceChangeInfo getNumServers() { - return _numServers; + public static class SegmentInfo { + private final int _totalSegmentsToBeMoved; + private final long _estimatedAverageSegmentSizeInBytes; + private final long _totalEstimatedDataToBeMovedInBytes; + private final double _totalEstimatedTimeToMoveDataInSecs; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RebalanceChangeInfo _replicationFactor; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RebalanceChangeInfo _numSegmentsInSingleReplica; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final RebalanceChangeInfo _numSegmentsAcrossAllReplicas; + + /** + * Constructor for SegmentInfo + * @param totalSegmentsToBeMoved total number of segments to be moved as part of this rebalance + * @param estimatedAverageSegmentSizeInBytes estimated average size of segments in bytes + * @param totalEstimatedDataToBeMovedInBytes total estimated amount of data to be moved as part of this rebalance + * @param totalEstimatedTimeToMoveDataInSecs total estimated time to move the segments as part of this rebalance + * @param replicationFactor replication factor before and after this rebalance + * @param numSegmentsInSingleReplica number of segments in single replica before and after this rebalance + * @param numSegmentsAcrossAllReplicas total number of segments across all replicas before and after this rebalance + */ + @JsonCreator + public SegmentInfo(@JsonProperty("totalSegmentsToBeMoved") int totalSegmentsToBeMoved, + @JsonProperty("estimatedAverageSegmentSizeInBytes") long estimatedAverageSegmentSizeInBytes, + @JsonProperty("totalEstimatedDataToBeMovedInBytes") long totalEstimatedDataToBeMovedInBytes, + @JsonProperty("totalEstimatedTimeToMoveDataInSecs") double totalEstimatedTimeToMoveDataInSecs, + @JsonProperty("replicationFactor") @Nullable RebalanceChangeInfo replicationFactor, + @JsonProperty("numSegmentsInSingleReplica") @Nullable RebalanceChangeInfo numSegmentsInSingleReplica, + @JsonProperty("numSegmentsAcrossAllReplicas") @Nullable RebalanceChangeInfo numSegmentsAcrossAllReplicas) { + _totalSegmentsToBeMoved = totalSegmentsToBeMoved; + _estimatedAverageSegmentSizeInBytes = estimatedAverageSegmentSizeInBytes; + _totalEstimatedDataToBeMovedInBytes = totalEstimatedDataToBeMovedInBytes; + _totalEstimatedTimeToMoveDataInSecs = totalEstimatedTimeToMoveDataInSecs; + _replicationFactor = replicationFactor; + _numSegmentsInSingleReplica = numSegmentsInSingleReplica; + _numSegmentsAcrossAllReplicas = numSegmentsAcrossAllReplicas; + } + + @JsonProperty + public int getTotalSegmentsToBeMoved() { + return _totalSegmentsToBeMoved; + } + + @JsonProperty + public long getEstimatedAverageSegmentSizeInBytes() { + return _estimatedAverageSegmentSizeInBytes; + } + + @JsonProperty + public long getTotalEstimatedDataToBeMovedInBytes() { + return _totalEstimatedDataToBeMovedInBytes; + } + + @JsonProperty + public double getTotalEstimatedTimeToMoveDataInSecs() { + return _totalEstimatedTimeToMoveDataInSecs; + } + + @JsonProperty + public RebalanceChangeInfo getReplicationFactor() { + return _replicationFactor; + } + + @JsonProperty + public RebalanceChangeInfo getNumSegmentsInSingleReplica() { + return _numSegmentsInSingleReplica; + } + + @JsonProperty + public RebalanceChangeInfo getNumSegmentsAcrossAllReplicas() { + return _numSegmentsAcrossAllReplicas; + } } - @JsonProperty - public RebalanceChangeInfo getReplicationFactor() { - return _replicationFactor; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final ServerInfo _serverInfo; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final SegmentInfo _segmentInfo; + + /** + * Constructor for RebalanceSummaryResult + * @param serverInfo server related summary information + * @param segmentInfo segment related summary information + */ + @JsonCreator + public RebalanceSummaryResult(@JsonProperty("serverInfo") @Nullable ServerInfo serverInfo, + @JsonProperty("segmentInfo") @Nullable SegmentInfo segmentInfo) { + _serverInfo = serverInfo; + _segmentInfo = segmentInfo; } @JsonProperty - public RebalanceChangeInfo getNumUniqueSegments() { - return _numUniqueSegments; + public ServerInfo getServerInfo() { + return _serverInfo; } @JsonProperty - public RebalanceChangeInfo getNumTotalSegments() { - return _numTotalSegments; + public SegmentInfo getSegmentInfo() { + return _segmentInfo; } - @JsonProperty - public Map getServerSegmentChangeInfo() { - return _serverSegmentChangeInfo; + public enum ServerStatus { + // ADDED if the server is newly added as part of rebalance; + // REMOVED if the server is removed as part of rebalance; + // UNCHANGED if the server status is unchanged as part of rebalance; + ADDED, REMOVED, UNCHANGED } } 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 29819ae650a5..c785cdc031ed 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 @@ -43,6 +43,7 @@ import org.apache.helix.PropertyKey; import org.apache.helix.model.ExternalView; import org.apache.helix.model.IdealState; +import org.apache.helix.model.InstanceConfig; import org.apache.helix.zookeeper.datamodel.ZNRecord; import org.apache.helix.zookeeper.zkclient.exception.ZkBadVersionException; import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils; @@ -637,6 +638,13 @@ private RebalanceSummaryResult calculateDryRunSummary(Map instanceConfigs = _helixDataAccessor.getChildValues( + _helixDataAccessor.keyBuilder().instanceConfigs(), true); + Map> instanceToTagsMap = new HashMap<>(); + for (InstanceConfig instanceConfig : instanceConfigs) { + instanceToTagsMap.put(instanceConfig.getInstanceName(), instanceConfig.getTags()); + } + Map serverSegmentChangeInfoMap = new HashMap<>(); int segmentsNotMoved = 0; int numServersGettingDataAdded = 0; @@ -649,6 +657,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map existingSegmentList = new HashSet<>(); int segmentsUnchanged = 0; int totalExistingSegments = 0; + RebalanceSummaryResult.ServerStatus serverStatus = RebalanceSummaryResult.ServerStatus.ADDED; if (existingServersToSegmentMap.containsKey(server)) { totalExistingSegments = existingServersToSegmentMap.get(server).size(); existingSegmentList.addAll(existingServersToSegmentMap.get(server)); @@ -656,30 +665,33 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0 ? 1 : 0; - serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(totalNewSegments, - totalExistingSegments, segmentsAdded, segmentsDeleted, segmentsUnchanged)); + serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(serverStatus, + totalNewSegments, totalExistingSegments, segmentsAdded, segmentsDeleted, segmentsUnchanged, + instanceToTagsMap.getOrDefault(server, null))); } for (Map.Entry> entry : existingServersToSegmentMap.entrySet()) { String server = entry.getKey(); if (!serverSegmentChangeInfoMap.containsKey(server)) { - serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(0, - entry.getValue().size(), 0, entry.getValue().size(), 0)); + serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo( + RebalanceSummaryResult.ServerStatus.REMOVED, 0, entry.getValue().size(), 0, entry.getValue().size(), 0, + instanceToTagsMap.getOrDefault(server, null))); } } - RebalanceSummaryResult.RebalanceChangeInfo uniqueNumSegments + RebalanceSummaryResult.RebalanceChangeInfo numSegmentsInSingleReplica = new RebalanceSummaryResult.RebalanceChangeInfo(currentAssignment.size(), targetAssignment.size()); int existingNumberSegmentsTotal = existingReplicationFactor * currentAssignment.size(); int newNumberSegmentsTotal = newReplicationFactor * targetAssignment.size(); - RebalanceSummaryResult.RebalanceChangeInfo totalNumSegments + RebalanceSummaryResult.RebalanceChangeInfo numSegmentsAcrossAllReplicas = new RebalanceSummaryResult.RebalanceChangeInfo(existingNumberSegmentsTotal, newNumberSegmentsTotal); int totalSegmentsToBeMoved = newNumberSegmentsTotal - segmentsNotMoved; @@ -692,11 +704,15 @@ private RebalanceSummaryResult calculateDryRunSummary(Map serverSegmentChangeInfoMap = - rebalanceResult.getRebalanceSummaryResult().getServerSegmentChangeInfo(); + rebalanceSummaryResult.getServerInfo().getServerSegmentChangeInfo(); assertNotNull(serverSegmentChangeInfoMap); for (int i = 0; i < numServers; i++) { // Original servers should be losing some segments String newServer = SERVER_INSTANCE_ID_PREFIX + i; RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); - assertTrue(serverSegmentChange._segmentsDeleted > 0); - assertTrue(serverSegmentChange._segmentsUnchanged > 0); - assertTrue(serverSegmentChange._totalExistingSegments > 0); - assertTrue(serverSegmentChange._totalNewSegments > 0); - assertEquals(serverSegmentChange._segmentsAdded, 0); + assertTrue(serverSegmentChange.getSegmentsDeleted() > 0); + assertTrue(serverSegmentChange.getSegmentsUnchanged() > 0); + assertTrue(serverSegmentChange.getTotalSegmentsBeforeRebalance() > 0); + assertTrue(serverSegmentChange.getTotalSegmentsAfterRebalance() > 0); + assertEquals(serverSegmentChange.getSegmentsAdded(), 0); } for (int i = 0; i < numServersToAdd; i++) { // New servers should only get new segments String newServer = SERVER_INSTANCE_ID_PREFIX + (numServers + i); RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); - assertTrue(serverSegmentChange._segmentsAdded > 0); - assertEquals(serverSegmentChange._totalExistingSegments, 0); - assertEquals(serverSegmentChange._totalNewSegments, serverSegmentChange._segmentsAdded); - assertEquals(serverSegmentChange._segmentsDeleted, 0); - assertEquals(serverSegmentChange._segmentsUnchanged, 0); + assertTrue(serverSegmentChange.getSegmentsAdded() > 0); + assertEquals(serverSegmentChange.getTotalSegmentsBeforeRebalance(), 0); + assertEquals(serverSegmentChange.getTotalSegmentsAfterRebalance(), serverSegmentChange.getSegmentsAdded()); + assertEquals(serverSegmentChange.getSegmentsDeleted(), 0); + assertEquals(serverSegmentChange.getSegmentsUnchanged(), 0); } // Dry-run mode should not change the IdealState @@ -262,15 +270,20 @@ public void testRebalance() rebalanceConfig = new RebalanceConfig(); rebalanceConfig.setDryRun(true); rebalanceConfig.setSummary(true); + rebalanceConfig.setPreChecks(true); rebalanceConfig.setMinAvailableReplicas(3); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); + assertNotNull(rebalanceResult.getPreChecksResult()); // Rebalance with 3 min available replicas should fail as the table only have 3 replicas rebalanceConfig = new RebalanceConfig(); @@ -311,24 +324,27 @@ public void testRebalance() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 11); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 11); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); - serverSegmentChangeInfoMap = rebalanceResult.getRebalanceSummaryResult().getServerSegmentChangeInfo(); + serverSegmentChangeInfoMap = rebalanceSummaryResult.getServerInfo().getServerSegmentChangeInfo(); assertNotNull(serverSegmentChangeInfoMap); for (int i = 0; i < numServers + numServersToAdd; i++) { String newServer = SERVER_INSTANCE_ID_PREFIX + i; RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChange = serverSegmentChangeInfoMap.get(newServer); - assertEquals(serverSegmentChange._totalNewSegments, 5); + assertEquals(serverSegmentChange.getTotalSegmentsAfterRebalance(), 5); // Ensure not all segments moved - assertTrue(serverSegmentChange._segmentsUnchanged > 0); + assertTrue(serverSegmentChange.getSegmentsUnchanged() > 0); // Ensure all segments has something assigned prior to rebalance - assertTrue(serverSegmentChange._totalExistingSegments > 0); + assertTrue(serverSegmentChange.getTotalSegmentsBeforeRebalance() > 0); } // Dry-run mode should not change the IdealState @@ -389,14 +405,17 @@ public void testRebalance() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); // Without instances reassignment, the rebalance should return status NO_OP, and the existing instance partitions // should be used @@ -412,15 +431,18 @@ public void testRebalance() rebalanceConfig.setReassignInstances(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); // No move expected since already balanced - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); // With instances reassignment, the rebalance should return status DONE, the existing instance partitions should be // removed, and the default instance partitions should be used @@ -458,14 +480,17 @@ public void testRebalance() rebalanceConfig.setDowntime(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 3); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 15); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 3); // Rebalance with downtime should succeed rebalanceConfig = new RebalanceConfig(); @@ -505,6 +530,7 @@ public void testRebalance() assertNull(rebalanceResult.getSegmentAssignment()); _helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME); + executorService.shutdown(); } /** @@ -551,14 +577,17 @@ public void testRebalanceWithTiers() rebalanceConfig.setSummary(true); RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + RebalanceSummaryResult rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); // Run actual table rebalance rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -582,14 +611,17 @@ public void testRebalanceWithTiers() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -615,14 +647,17 @@ public void testRebalanceWithTiers() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 15); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 9); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 6); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 15); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 9); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 6); // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -689,14 +724,17 @@ public void testRebalanceWithTiersAndInstanceAssignments() rebalanceConfig.setSummary(true); RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + RebalanceSummaryResult rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); @@ -716,14 +754,17 @@ public void testRebalanceWithTiersAndInstanceAssignments() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 0); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 0); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 0); // rebalance is NOOP and no change in assignment caused by new instances rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -743,14 +784,17 @@ public void testRebalanceWithTiersAndInstanceAssignments() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 30); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 3); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServersGettingNewSegments(), 6); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 30); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 3); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServersGettingNewSegments(), 6); // rebalance should change assignment rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); @@ -781,13 +825,16 @@ public void testRebalanceWithTiersAndInstanceAssignments() rebalanceConfig.setSummary(true); rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); - assertNotNull(rebalanceResult.getRebalanceSummaryResult()); + rebalanceSummaryResult = rebalanceResult.getRebalanceSummaryResult(); + assertNotNull(rebalanceSummaryResult); + assertNotNull(rebalanceSummaryResult.getServerInfo()); + assertNotNull(rebalanceSummaryResult.getSegmentInfo()); assertNull(rebalanceResult.getInstanceAssignment()); assertNull(rebalanceResult.getTierInstanceAssignment()); assertNull(rebalanceResult.getSegmentAssignment()); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getTotalSegmentsToBeMoved(), 13); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._existingValue, 6); - assertEquals(rebalanceResult.getRebalanceSummaryResult().getNumServers()._newValue, 6); + assertEquals(rebalanceSummaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 13); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), 6); + assertEquals(rebalanceSummaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), 6); rebalanceResult = tableRebalancer.rebalance(tableConfig, new RebalanceConfig(), null); assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index d48dacf966ba..037aef489e8a 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -42,6 +42,7 @@ import java.util.Properties; import java.util.Set; import java.util.UUID; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -180,6 +181,7 @@ public class OfflineClusterIntegrationTest extends BaseClusterIntegrationTestSet private PinotHelixResourceManager _resourceManager; private TableRebalancer _tableRebalancer; + private ExecutorService _executorService; protected int getNumBrokers() { return NUM_BROKERS; @@ -291,7 +293,8 @@ public void setUp() _resourceManager = _controllerStarter.getHelixResourceManager(); DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker(); - preChecker.init(_helixResourceManager, Executors.newFixedThreadPool(10)); + _executorService = Executors.newFixedThreadPool(10); + preChecker.init(_helixResourceManager, _executorService); _tableRebalancer = new TableRebalancer(_resourceManager.getHelixZkManager(), null, null, preChecker, _resourceManager.getTableSizeReader()); } @@ -3067,6 +3070,7 @@ public void tearDown() stopController(); stopZk(); FileUtils.deleteDirectory(_tempDir); + _executorService.shutdown(); } private void testInstanceDecommission() @@ -4098,6 +4102,18 @@ public void testRebalanceDryRunSummary() rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.NO_OP, false, getNumServers(), getNumServers(), tableConfig.getReplication()); + + // Try dry-run with summary and pre-checks + rebalanceConfig.setPreChecks(true); + rebalanceResult = _tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + checkRebalanceDryRunSummary(rebalanceResult, RebalanceResult.Status.NO_OP, false, getNumServers(), getNumServers(), + tableConfig.getReplication()); + assertNotNull(rebalanceResult.getPreChecksResult()); + assertTrue(rebalanceResult.getPreChecksResult().containsKey(DefaultRebalancePreChecker.NEEDS_RELOAD_STATUS)); + assertTrue(rebalanceResult.getPreChecksResult().containsKey(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT)); + assertEquals(rebalanceResult.getPreChecksResult().get(DefaultRebalancePreChecker.NEEDS_RELOAD_STATUS), "false"); + assertEquals(rebalanceResult.getPreChecksResult().get(DefaultRebalancePreChecker.IS_MINIMIZE_DATA_MOVEMENT), + "false"); } private void checkRebalanceDryRunSummary(RebalanceResult rebalanceResult, RebalanceResult.Status expectedStatus, @@ -4105,38 +4121,74 @@ private void checkRebalanceDryRunSummary(RebalanceResult rebalanceResult, Rebala assertEquals(rebalanceResult.getStatus(), expectedStatus); RebalanceSummaryResult summaryResult = rebalanceResult.getRebalanceSummaryResult(); assertNotNull(summaryResult); - assertEquals(summaryResult.getReplicationFactor()._existingValue, replicationFactor, + assertNotNull(summaryResult.getServerInfo()); + assertNotNull(summaryResult.getSegmentInfo()); + assertEquals(summaryResult.getSegmentInfo().getReplicationFactor().getValueBeforeRebalance(), replicationFactor, "Existing replication factor doesn't match expected"); - assertEquals(summaryResult.getReplicationFactor()._existingValue, summaryResult.getReplicationFactor()._newValue, + assertEquals(summaryResult.getSegmentInfo().getReplicationFactor().getValueBeforeRebalance(), + summaryResult.getSegmentInfo().getReplicationFactor().getExpectedValueAfterRebalance(), "Existing and new replication factor doesn't match"); - assertEquals(summaryResult.getNumServers()._existingValue, existingNumServers, + assertEquals(summaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), existingNumServers, "Existing number of servers don't match"); - assertEquals(summaryResult.getNumServers()._newValue, newNumServers, "New number of servers don't match"); + assertEquals(summaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), newNumServers, + "New number of servers don't match"); if (_tableSize > 0) { - assertTrue(summaryResult.getEstimatedAverageSegmentSizeInBytes() > 0L, + assertTrue(summaryResult.getSegmentInfo().getEstimatedAverageSegmentSizeInBytes() > 0L, "Avg segment size expected to be > 0 but found to be 0"); } if (existingNumServers != newNumServers) { - assertTrue(summaryResult.getNumServersGettingNewSegments() > 0, "Expected number of servers should be > 0"); + assertTrue(summaryResult.getServerInfo().getNumServersGettingNewSegments() > 0, + "Expected number of servers should be > 0"); } else { - assertEquals(summaryResult.getNumServersGettingNewSegments(), 0, + assertEquals(summaryResult.getServerInfo().getNumServersGettingNewSegments(), 0, "Expected number of servers getting new segments should be 0"); } if (isSegmentsToBeMoved) { - assertTrue(summaryResult.getTotalSegmentsToBeMoved() > 0, "Segments to be moved should be > 0"); - assertEquals(summaryResult.getTotalEstimatedDataToBeMovedInBytes(), - summaryResult.getTotalSegmentsToBeMoved() * summaryResult.getEstimatedAverageSegmentSizeInBytes(), + assertTrue(summaryResult.getSegmentInfo().getTotalSegmentsToBeMoved() > 0, + "Segments to be moved should be > 0"); + assertEquals(summaryResult.getSegmentInfo().getTotalEstimatedDataToBeMovedInBytes(), + summaryResult.getSegmentInfo().getTotalSegmentsToBeMoved() + * summaryResult.getSegmentInfo().getEstimatedAverageSegmentSizeInBytes(), "Estimated data to be moved in bytes doesn't match"); - assertTrue(summaryResult.getTotalEstimatedTimeToMoveDataInSecs() > 0.0D, + assertTrue(summaryResult.getSegmentInfo().getTotalEstimatedTimeToMoveDataInSecs() > 0.0D, "Estimated time to move segments should be more than 0.0 seconds"); } else { - assertEquals(summaryResult.getTotalSegmentsToBeMoved(), 0, "Segments to be moved should be 0"); - assertEquals(summaryResult.getTotalEstimatedDataToBeMovedInBytes(), 0L, + assertEquals(summaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0, "Segments to be moved should be 0"); + assertEquals(summaryResult.getSegmentInfo().getTotalEstimatedDataToBeMovedInBytes(), 0L, "Estimated data to be moved in bytes should be 0"); - assertEquals(summaryResult.getTotalEstimatedTimeToMoveDataInSecs(), 0.0D, + assertEquals(summaryResult.getSegmentInfo().getTotalEstimatedTimeToMoveDataInSecs(), 0.0D, "Estimated time to move segments should be 0.0D"); } + + // Validate server status stats with numServers information + Map serverSegmentChangeInfoMap = + summaryResult.getServerInfo().getServerSegmentChangeInfo(); + int numServersAdded = 0; + int numServersRemoved = 0; + int numServersUnchanged = 0; + for (RebalanceSummaryResult.ServerSegmentChangeInfo serverSegmentChangeInfo : serverSegmentChangeInfoMap.values()) { + switch (serverSegmentChangeInfo.getServerStatus()) { + case UNCHANGED: + numServersUnchanged++; + break; + case ADDED: + numServersAdded++; + break; + case REMOVED: + numServersRemoved++; + break; + default: + Assert.fail(String.format("Unknown server status encountered: %s", + serverSegmentChangeInfo.getServerStatus())); + break; + } + } + + Assert.assertEquals(summaryResult.getServerInfo().getNumServers().getValueBeforeRebalance(), + numServersRemoved + numServersUnchanged); + Assert.assertEquals(summaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), + numServersAdded + numServersUnchanged); } protected void waitForRebalanceToComplete(RebalanceResult rebalanceResult, long timeoutMs) From beb5c9b408ba3378cb68026280d086fef7b382b0 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Tue, 18 Feb 2025 08:38:12 -0800 Subject: [PATCH 4/7] Trigger Build From 3ce1cced600eb7b3f1da4c9068a0427605854991 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Tue, 18 Feb 2025 15:54:23 -0800 Subject: [PATCH 5/7] Address review comments --- .../rebalance/DefaultRebalancePreChecker.java | 3 +- .../rebalance/RebalanceSummaryResult.java | 19 +++--- .../helix/core/rebalance/TableRebalancer.java | 62 ++++++++----------- .../tests/OfflineClusterIntegrationTest.java | 10 +-- 4 files changed, 41 insertions(+), 53 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 945a98e48e17..b6a89208eeef 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 @@ -53,8 +53,7 @@ public void init(PinotHelixResourceManager pinotHelixResourceManager, @Nullable } @Override - public Map check(String rebalanceJobId, String tableNameWithType, - TableConfig tableConfig) { + public Map check(String rebalanceJobId, String tableNameWithType, TableConfig tableConfig) { LOGGER.info("Start pre-checks for table: {} with rebalanceJobId: {}", tableNameWithType, rebalanceJobId); Map preCheckResult = new HashMap<>(); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java index e9a80242e1d7..3733472e1625 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -173,10 +173,11 @@ public Map getServerSegmentChangeInfo() { } public static class SegmentInfo { + // TODO: Add a metric to estimate the total time it will take to rebalance private final int _totalSegmentsToBeMoved; + private final int _maxSegmentsAddedToASingleServer; private final long _estimatedAverageSegmentSizeInBytes; private final long _totalEstimatedDataToBeMovedInBytes; - private final double _totalEstimatedTimeToMoveDataInSecs; @JsonInclude(JsonInclude.Include.NON_NULL) private final RebalanceChangeInfo _replicationFactor; @JsonInclude(JsonInclude.Include.NON_NULL) @@ -187,25 +188,25 @@ public static class SegmentInfo { /** * Constructor for SegmentInfo * @param totalSegmentsToBeMoved total number of segments to be moved as part of this rebalance + * @param maxSegmentsAddedToASingleServer maximum segments added to a single server as part of this rebalance * @param estimatedAverageSegmentSizeInBytes estimated average size of segments in bytes * @param totalEstimatedDataToBeMovedInBytes total estimated amount of data to be moved as part of this rebalance - * @param totalEstimatedTimeToMoveDataInSecs total estimated time to move the segments as part of this rebalance * @param replicationFactor replication factor before and after this rebalance * @param numSegmentsInSingleReplica number of segments in single replica before and after this rebalance * @param numSegmentsAcrossAllReplicas total number of segments across all replicas before and after this rebalance */ @JsonCreator public SegmentInfo(@JsonProperty("totalSegmentsToBeMoved") int totalSegmentsToBeMoved, + @JsonProperty("maxSegmentsAddedToASingleServer") int maxSegmentsAddedToASingleServer, @JsonProperty("estimatedAverageSegmentSizeInBytes") long estimatedAverageSegmentSizeInBytes, @JsonProperty("totalEstimatedDataToBeMovedInBytes") long totalEstimatedDataToBeMovedInBytes, - @JsonProperty("totalEstimatedTimeToMoveDataInSecs") double totalEstimatedTimeToMoveDataInSecs, @JsonProperty("replicationFactor") @Nullable RebalanceChangeInfo replicationFactor, @JsonProperty("numSegmentsInSingleReplica") @Nullable RebalanceChangeInfo numSegmentsInSingleReplica, @JsonProperty("numSegmentsAcrossAllReplicas") @Nullable RebalanceChangeInfo numSegmentsAcrossAllReplicas) { _totalSegmentsToBeMoved = totalSegmentsToBeMoved; + _maxSegmentsAddedToASingleServer = maxSegmentsAddedToASingleServer; _estimatedAverageSegmentSizeInBytes = estimatedAverageSegmentSizeInBytes; _totalEstimatedDataToBeMovedInBytes = totalEstimatedDataToBeMovedInBytes; - _totalEstimatedTimeToMoveDataInSecs = totalEstimatedTimeToMoveDataInSecs; _replicationFactor = replicationFactor; _numSegmentsInSingleReplica = numSegmentsInSingleReplica; _numSegmentsAcrossAllReplicas = numSegmentsAcrossAllReplicas; @@ -216,6 +217,11 @@ public int getTotalSegmentsToBeMoved() { return _totalSegmentsToBeMoved; } + @JsonProperty + public int getMaxSegmentsAddedToASingleServer() { + return _maxSegmentsAddedToASingleServer; + } + @JsonProperty public long getEstimatedAverageSegmentSizeInBytes() { return _estimatedAverageSegmentSizeInBytes; @@ -226,11 +232,6 @@ public long getTotalEstimatedDataToBeMovedInBytes() { return _totalEstimatedDataToBeMovedInBytes; } - @JsonProperty - public double getTotalEstimatedTimeToMoveDataInSecs() { - return _totalEstimatedTimeToMoveDataInSecs; - } - @JsonProperty public RebalanceChangeInfo getReplicationFactor() { return _replicationFactor; 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 c785cdc031ed..005d078c93ab 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 @@ -594,15 +594,16 @@ private long calculateTableSizePerReplicaInBytes(String tableNameWithType) { LOGGER.warn("tableSizeReader is null, cannot calculate table size for table {}!", tableNameWithType); return tableSizePerReplicaInBytes; } + LOGGER.info("Fetching the table size for rebalance summary for table: {}", tableNameWithType); try { + // TODO: Consider making the timeoutMs for fetching table size via table rebalancer configurable TableSizeReader.TableSubTypeSizeDetails tableSizeDetails = _tableSizeReader.getTableSubtypeSize(tableNameWithType, 30_000); tableSizePerReplicaInBytes = tableSizeDetails._reportedSizePerReplicaInBytes; } catch (InvalidConfigException e) { - String errMsg = String.format("Caught exception while trying to fetch table size details for table: %s", - tableNameWithType); - LOGGER.error(errMsg, e); + LOGGER.error("Caught exception while trying to fetch table size details for table: {}", tableNameWithType, e); } + LOGGER.info("Fetched the table size for rebalance summary for table: {}", tableNameWithType); return tableSizePerReplicaInBytes; } @@ -617,24 +618,22 @@ private RebalanceSummaryResult calculateDryRunSummary(Map> entrySet : currentAssignment.entrySet()) { existingReplicationFactor = entrySet.getValue().size(); - for (Map.Entry segmentEntrySet : entrySet.getValue().entrySet()) { - existingServersToSegmentMap.putIfAbsent(segmentEntrySet.getKey(), new HashSet<>()); - existingServersToSegmentMap.get(segmentEntrySet.getKey()).add(entrySet.getKey()); + for (String segmentKey : entrySet.getValue().keySet()) { + existingServersToSegmentMap.computeIfAbsent(segmentKey, k -> new HashSet<>()).add(entrySet.getKey()); } } for (Map.Entry> entrySet : targetAssignment.entrySet()) { newReplicationFactor = entrySet.getValue().size(); - for (Map.Entry segmentEntrySet : entrySet.getValue().entrySet()) { - newServersToSegmentMap.putIfAbsent(segmentEntrySet.getKey(), new HashSet<>()); - newServersToSegmentMap.get(segmentEntrySet.getKey()).add(entrySet.getKey()); + for (String segmentKey : entrySet.getValue().keySet()) { + newServersToSegmentMap.computeIfAbsent(segmentKey, k -> new HashSet<>()).add(entrySet.getKey()); } } RebalanceSummaryResult.RebalanceChangeInfo replicationFactor = new RebalanceSummaryResult.RebalanceChangeInfo(existingReplicationFactor, newReplicationFactor); - int existingNumServers = existingServersToSegmentMap.keySet().size(); - int newNumServers = newServersToSegmentMap.keySet().size(); + int existingNumServers = existingServersToSegmentMap.size(); + int newNumServers = newServersToSegmentMap.size(); RebalanceSummaryResult.RebalanceChangeInfo numServers = new RebalanceSummaryResult.RebalanceChangeInfo(existingNumServers, newNumServers); @@ -648,28 +647,31 @@ private RebalanceSummaryResult calculateDryRunSummary(Map serverSegmentChangeInfoMap = new HashMap<>(); int segmentsNotMoved = 0; int numServersGettingDataAdded = 0; + int maxSegmentsAddedToServer = 0; for (Map.Entry> entry : newServersToSegmentMap.entrySet()) { String server = entry.getKey(); - Set segmentMap = entry.getValue(); - int totalNewSegments = segmentMap.size(); + Set segmentSet = entry.getValue(); + int totalNewSegments = segmentSet.size(); - Set newSegmentList = new HashSet<>(segmentMap); - Set existingSegmentList = new HashSet<>(); + Set newSegmentList = new HashSet<>(segmentSet); + Set existingSegmentSet = new HashSet<>(); int segmentsUnchanged = 0; int totalExistingSegments = 0; RebalanceSummaryResult.ServerStatus serverStatus = RebalanceSummaryResult.ServerStatus.ADDED; if (existingServersToSegmentMap.containsKey(server)) { - totalExistingSegments = existingServersToSegmentMap.get(server).size(); - existingSegmentList.addAll(existingServersToSegmentMap.get(server)); - Set intersection = new HashSet<>(existingServersToSegmentMap.get(server)); + Set segmentSetForServer = existingServersToSegmentMap.get(server); + totalExistingSegments = segmentSetForServer.size(); + existingSegmentSet.addAll(segmentSetForServer); + Set intersection = new HashSet<>(segmentSetForServer); intersection.retainAll(newSegmentList); segmentsUnchanged = intersection.size(); segmentsNotMoved += segmentsUnchanged; serverStatus = RebalanceSummaryResult.ServerStatus.UNCHANGED; } - newSegmentList.removeAll(existingSegmentList); + newSegmentList.removeAll(existingSegmentSet); int segmentsAdded = newSegmentList.size(); - int segmentsDeleted = existingSegmentList.size() - segmentsUnchanged; + maxSegmentsAddedToServer = Math.max(maxSegmentsAddedToServer, segmentsAdded); + int segmentsDeleted = existingSegmentSet.size() - segmentsUnchanged; numServersGettingDataAdded += segmentsAdded > 0 ? 1 : 0; serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(serverStatus, @@ -701,13 +703,13 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0 && numServersGettingDataAdded > 0) { - // Do some processing to figure out what the time to download might be - // Assume that data is evenly distributed across all servers - long totalDataPerServerInBytes = totalEstimatedDataToBeMovedInBytes / numServersGettingDataAdded; - // TODO: Pick a good threshold to calculate estimated time to rebalance. For now assume 100 MB/s data download - // + process rate - estimatedTimeToRebalanceInSec = ((double) totalDataPerServerInBytes) / (100.0D * 1024.0D * 1024.0D); - } - return estimatedTimeToRebalanceInSec; - } - private void onReturnFailure(String errorMsg, Exception e) { if (e != null) { LOGGER.warn(errorMsg, e); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index 037aef489e8a..d25ff1747947 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -296,7 +296,7 @@ public void setUp() _executorService = Executors.newFixedThreadPool(10); preChecker.init(_helixResourceManager, _executorService); _tableRebalancer = new TableRebalancer(_resourceManager.getHelixZkManager(), null, null, preChecker, - _resourceManager.getTableSizeReader()); + _resourceManager.getTableSizeReader()); } private void reloadAllSegments(String testQuery, boolean forceDownload, long numTotalDocs) @@ -4151,14 +4151,14 @@ private void checkRebalanceDryRunSummary(RebalanceResult rebalanceResult, Rebala summaryResult.getSegmentInfo().getTotalSegmentsToBeMoved() * summaryResult.getSegmentInfo().getEstimatedAverageSegmentSizeInBytes(), "Estimated data to be moved in bytes doesn't match"); - assertTrue(summaryResult.getSegmentInfo().getTotalEstimatedTimeToMoveDataInSecs() > 0.0D, - "Estimated time to move segments should be more than 0.0 seconds"); + assertTrue(summaryResult.getSegmentInfo().getMaxSegmentsAddedToASingleServer() > 0, + "Estimated max number of segments to move in a single server should be > 0"); } else { assertEquals(summaryResult.getSegmentInfo().getTotalSegmentsToBeMoved(), 0, "Segments to be moved should be 0"); assertEquals(summaryResult.getSegmentInfo().getTotalEstimatedDataToBeMovedInBytes(), 0L, "Estimated data to be moved in bytes should be 0"); - assertEquals(summaryResult.getSegmentInfo().getTotalEstimatedTimeToMoveDataInSecs(), 0.0D, - "Estimated time to move segments should be 0.0D"); + assertEquals(summaryResult.getSegmentInfo().getMaxSegmentsAddedToASingleServer(), 0, + "Estimated max number of segments to move in a single server should be 0"); } // Validate server status stats with numServers information From 20060fadee0ece9d5f718411642c1670bc78a2d8 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Tue, 18 Feb 2025 16:27:29 -0800 Subject: [PATCH 6/7] Address review - add server level lists of added, removed, unchanged, and servers getting new segments --- .../rebalance/RebalanceSummaryResult.java | 41 +++++++++++++++++++ .../helix/core/rebalance/TableRebalancer.java | 19 +++++++-- .../tests/OfflineClusterIntegrationTest.java | 9 +++- 3 files changed, 63 insertions(+), 6 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java index 3733472e1625..81a1418993ac 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -24,6 +24,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import java.util.List; import java.util.Map; +import java.util.Set; import javax.annotation.Nullable; @@ -138,21 +139,41 @@ public static class ServerInfo { @JsonInclude(JsonInclude.Include.NON_NULL) private final RebalanceChangeInfo _numServers; @JsonInclude(JsonInclude.Include.NON_NULL) + private final Set _serversAdded; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final Set _serversRemoved; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final Set _serversUnchanged; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final Set _serversGettingNewSegments; + @JsonInclude(JsonInclude.Include.NON_NULL) private final Map _serverSegmentChangeInfo; /** * Constructor for ServerInfo * @param numServersGettingNewSegments total number of servers receiving new segments as part of this rebalance * @param numServers number of servers before and after this rebalance + * @param serversAdded set of servers getting added as part of this rebalance + * @param serversRemoved set of servers getting removed as part of this rebalance + * @param serversUnchanged set of servers existing both before and as part of this rebalance + * @param serversGettingNewSegments set of servers getting segments added * @param serverSegmentChangeInfo per server statistics for this rebalance */ @JsonCreator public ServerInfo(@JsonProperty("numServersGettingNewSegments") int numServersGettingNewSegments, @JsonProperty("numServers") @Nullable RebalanceChangeInfo numServers, + @JsonProperty("serversAdded") @Nullable Set serversAdded, + @JsonProperty("serversRemoved") @Nullable Set serversRemoved, + @JsonProperty("serversUnchanged") @Nullable Set serversUnchanged, + @JsonProperty("serversGettingNewSegments") @Nullable Set serversGettingNewSegments, @JsonProperty("serverSegmentChangeInfo") @Nullable Map serverSegmentChangeInfo) { _numServersGettingNewSegments = numServersGettingNewSegments; _numServers = numServers; + _serversAdded = serversAdded; + _serversRemoved = serversRemoved; + _serversUnchanged = serversUnchanged; + _serversGettingNewSegments = serversGettingNewSegments; _serverSegmentChangeInfo = serverSegmentChangeInfo; } @@ -166,6 +187,26 @@ public RebalanceChangeInfo getNumServers() { return _numServers; } + @JsonProperty + public Set getServersAdded() { + return _serversAdded; + } + + @JsonProperty + public Set getServersRemoved() { + return _serversRemoved; + } + + @JsonProperty + public Set getServersUnchanged() { + return _serversUnchanged; + } + + @JsonProperty + public Set getServersGettingNewSegments() { + return _serversGettingNewSegments; + } + @JsonProperty public Map getServerSegmentChangeInfo() { return _serverSegmentChangeInfo; 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 005d078c93ab..e23a6153f120 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 @@ -644,9 +644,12 @@ private RebalanceSummaryResult calculateDryRunSummary(Map serversAdded = new HashSet<>(); + Set serversRemoved = new HashSet<>(); + Set serversUnchanged = new HashSet<>(); + Set serversGettingNewSegments = new HashSet<>(); Map serverSegmentChangeInfoMap = new HashMap<>(); int segmentsNotMoved = 0; - int numServersGettingDataAdded = 0; int maxSegmentsAddedToServer = 0; for (Map.Entry> entry : newServersToSegmentMap.entrySet()) { String server = entry.getKey(); @@ -667,12 +670,17 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0) { + serversGettingNewSegments.add(server); + } maxSegmentsAddedToServer = Math.max(maxSegmentsAddedToServer, segmentsAdded); int segmentsDeleted = existingSegmentSet.size() - segmentsUnchanged; - numServersGettingDataAdded += segmentsAdded > 0 ? 1 : 0; serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo(serverStatus, totalNewSegments, totalExistingSegments, segmentsAdded, segmentsDeleted, segmentsUnchanged, @@ -682,6 +690,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map> entry : existingServersToSegmentMap.entrySet()) { String server = entry.getKey(); if (!serverSegmentChangeInfoMap.containsKey(server)) { + serversRemoved.add(server); serverSegmentChangeInfoMap.put(server, new RebalanceSummaryResult.ServerSegmentChangeInfo( RebalanceSummaryResult.ServerStatus.REMOVED, 0, entry.getValue().size(), 0, entry.getValue().size(), 0, instanceToTagsMap.getOrDefault(server, null))); @@ -704,8 +713,10 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0L, "Avg segment size expected to be > 0 but found to be 0"); } + assertEquals(summaryResult.getServerInfo().getNumServersGettingNewSegments(), + summaryResult.getServerInfo().getServersGettingNewSegments().size()); if (existingNumServers != newNumServers) { assertTrue(summaryResult.getServerInfo().getNumServersGettingNewSegments() > 0, "Expected number of servers should be > 0"); @@ -4189,10 +4191,13 @@ private void checkRebalanceDryRunSummary(RebalanceResult rebalanceResult, Rebala numServersRemoved + numServersUnchanged); Assert.assertEquals(summaryResult.getServerInfo().getNumServers().getExpectedValueAfterRebalance(), numServersAdded + numServersUnchanged); + + assertEquals(numServersAdded, summaryResult.getServerInfo().getServersAdded().size()); + assertEquals(numServersRemoved, summaryResult.getServerInfo().getServersRemoved().size()); + assertEquals(numServersUnchanged, summaryResult.getServerInfo().getServersUnchanged().size()); } - protected void waitForRebalanceToComplete(RebalanceResult rebalanceResult, long timeoutMs) - throws Exception { + protected void waitForRebalanceToComplete(RebalanceResult rebalanceResult, long timeoutMs) { String jobId = rebalanceResult.getJobId(); if (rebalanceResult.getStatus() != RebalanceResult.Status.IN_PROGRESS) { return; From 6a04c661f6115bb9b4a81bedd1652a6eba3e0d3c Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Wed, 19 Feb 2025 11:11:16 -0800 Subject: [PATCH 7/7] Address review comments --- .../rebalance/RebalanceSummaryResult.java | 54 +++++++++---------- .../helix/core/rebalance/TableRebalancer.java | 13 +++-- 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java index 81a1418993ac..3169c9b3676c 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceSummaryResult.java @@ -35,6 +35,33 @@ @JsonInclude(JsonInclude.Include.NON_NULL) public class RebalanceSummaryResult { + @JsonInclude(JsonInclude.Include.NON_NULL) + private final ServerInfo _serverInfo; + @JsonInclude(JsonInclude.Include.NON_NULL) + private final SegmentInfo _segmentInfo; + + /** + * Constructor for RebalanceSummaryResult + * @param serverInfo server related summary information + * @param segmentInfo segment related summary information + */ + @JsonCreator + public RebalanceSummaryResult(@JsonProperty("serverInfo") @Nullable ServerInfo serverInfo, + @JsonProperty("segmentInfo") @Nullable SegmentInfo segmentInfo) { + _serverInfo = serverInfo; + _segmentInfo = segmentInfo; + } + + @JsonProperty + public ServerInfo getServerInfo() { + return _serverInfo; + } + + @JsonProperty + public SegmentInfo getSegmentInfo() { + return _segmentInfo; + } + public static class ServerSegmentChangeInfo { private final ServerStatus _serverStatus; private final int _totalSegmentsAfterRebalance; @@ -289,33 +316,6 @@ public RebalanceChangeInfo getNumSegmentsAcrossAllReplicas() { } } - @JsonInclude(JsonInclude.Include.NON_NULL) - private final ServerInfo _serverInfo; - @JsonInclude(JsonInclude.Include.NON_NULL) - private final SegmentInfo _segmentInfo; - - /** - * Constructor for RebalanceSummaryResult - * @param serverInfo server related summary information - * @param segmentInfo segment related summary information - */ - @JsonCreator - public RebalanceSummaryResult(@JsonProperty("serverInfo") @Nullable ServerInfo serverInfo, - @JsonProperty("segmentInfo") @Nullable SegmentInfo segmentInfo) { - _serverInfo = serverInfo; - _segmentInfo = segmentInfo; - } - - @JsonProperty - public ServerInfo getServerInfo() { - return _serverInfo; - } - - @JsonProperty - public SegmentInfo getSegmentInfo() { - return _segmentInfo; - } - public enum ServerStatus { // ADDED if the server is newly added as part of rebalance; // REMOVED if the server is removed as part of rebalance; 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 e23a6153f120..57e43385572e 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 @@ -315,6 +315,8 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb RebalanceSummaryResult summaryResult = null; if (summary) { + // Calculate summary here itself so that even if the table is already balanced, the caller can verify whether that + // is expected or not based on the summary results summaryResult = calculateDryRunSummary(currentAssignment, targetAssignment, tableNameWithType, rebalanceJobId); } @@ -603,7 +605,8 @@ private long calculateTableSizePerReplicaInBytes(String tableNameWithType) { } catch (InvalidConfigException e) { LOGGER.error("Caught exception while trying to fetch table size details for table: {}", tableNameWithType, e); } - LOGGER.info("Fetched the table size for rebalance summary for table: {}", tableNameWithType); + LOGGER.info("Fetched the table size (per replica size: {}) for rebalance summary for table: {}", + tableSizePerReplicaInBytes, tableNameWithType); return tableSizePerReplicaInBytes; } @@ -656,7 +659,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map segmentSet = entry.getValue(); int totalNewSegments = segmentSet.size(); - Set newSegmentList = new HashSet<>(segmentSet); + Set newSegmentSet = new HashSet<>(segmentSet); Set existingSegmentSet = new HashSet<>(); int segmentsUnchanged = 0; int totalExistingSegments = 0; @@ -666,7 +669,7 @@ private RebalanceSummaryResult calculateDryRunSummary(Map intersection = new HashSet<>(segmentSetForServer); - intersection.retainAll(newSegmentList); + intersection.retainAll(newSegmentSet); segmentsUnchanged = intersection.size(); segmentsNotMoved += segmentsUnchanged; serverStatus = RebalanceSummaryResult.ServerStatus.UNCHANGED; @@ -674,8 +677,8 @@ private RebalanceSummaryResult calculateDryRunSummary(Map 0) { serversGettingNewSegments.add(server); }