You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xb...@apache.org on 2023/10/18 21:33:12 UTC
[pinot] branch master updated: Add support to retry failed table rebalance (#11740)
This is an automated email from the ASF dual-hosted git repository.
xbli pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 0cf091a7ef Add support to retry failed table rebalance (#11740)
0cf091a7ef is described below
commit 0cf091a7ef822470fd2c0195a926b35e8cddcf72
Author: Xiaobing <61...@users.noreply.github.com>
AuthorDate: Wed Oct 18 14:33:07 2023 -0700
Add support to retry failed table rebalance (#11740)
* add support to retry failed table rebalance
* add exponential retry backoff and more metrics
* add CANCELLED job status and cancelRebalance restful API for users to cancel rebalance
---
.../configs/controller.yml | 6 +
.../etc/jmx_prometheus_javaagent/configs/pinot.yml | 6 +
.../pinot/common/metrics/ControllerMeter.java | 3 +
.../pinot/common/metrics/ControllerTimer.java | 1 +
.../pinot/controller/BaseControllerStarter.java | 5 +
.../apache/pinot/controller/ControllerConf.java | 22 +-
.../api/resources/PinotTableRestletResource.java | 126 +++++--
.../ServerRebalanceJobStatusResponse.java | 11 +
.../helix/core/PinotHelixResourceManager.java | 147 +++++---
.../core/rebalance/NoOpTableRebalanceObserver.java | 10 +
.../helix/core/rebalance/RebalanceChecker.java | 331 ++++++++++++++++++
.../helix/core/rebalance/RebalanceConfig.java | 66 ++++
.../core/rebalance/RebalanceJobConstants.java | 2 +
.../helix/core/rebalance/RebalanceResult.java | 6 +-
.../core/rebalance/TableRebalanceContext.java | 94 +++++
.../core/rebalance/TableRebalanceObserver.java | 4 +
.../rebalance/TableRebalanceProgressStats.java | 6 +-
.../helix/core/rebalance/TableRebalancer.java | 35 ++
.../rebalance/ZkBasedTableRebalanceObserver.java | 123 +++++--
.../tenant/ZkBasedTenantRebalanceObserver.java | 4 +-
...ControllerPeriodicTaskStarterStatelessTest.java | 2 +-
.../helix/core/rebalance/RebalanceCheckerTest.java | 388 +++++++++++++++++++++
.../TestZkBasedTableRebalanceObserver.java | 6 +-
.../java/org/apache/pinot/core/auth/Actions.java | 1 +
...PartialUpsertTableRebalanceIntegrationTest.java | 5 +-
25 files changed, 1284 insertions(+), 126 deletions(-)
diff --git a/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/controller.yml b/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/controller.yml
index cfd4f9af31..c4071887ed 100644
--- a/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/controller.yml
+++ b/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/controller.yml
@@ -106,6 +106,12 @@ rules:
labels:
table: "$1"
taskType: "$2"
+- pattern: "\"org.apache.pinot.common.metrics\"<type=\"ControllerMetrics\", name=\"pinot.controller.([^\\.]*?)\\.(\\w+).tableRebalanceExecutionTimeMs\"><>(\\w+)"
+ name: "pinot_controller_tableRebalanceExecutionTimeMs_$3"
+ cache: true
+ labels:
+ table: "$1"
+ result: "$2"
- pattern: "\"org.apache.pinot.common.metrics\"<type=\"ControllerMetrics\", name=\"pinot.controller.taskStatus.([^\\.]*?)\\.(\\w+)\"><>(\\w+)"
name: "pinot_controller_taskStatus_$3"
cache: true
diff --git a/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/pinot.yml b/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/pinot.yml
index 6660ee42d7..1c435be787 100644
--- a/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/pinot.yml
+++ b/docker/images/pinot/etc/jmx_prometheus_javaagent/configs/pinot.yml
@@ -106,6 +106,12 @@ rules:
labels:
table: "$1"
taskType: "$2"
+- pattern: "\"org.apache.pinot.common.metrics\"<type=\"ControllerMetrics\", name=\"pinot.controller.([^\\.]*?)\\.(\\w+).tableRebalanceExecutionTimeMs\"><>(\\w+)"
+ name: "pinot_controller_tableRebalanceExecutionTimeMs_$3"
+ cache: true
+ labels:
+ table: "$1"
+ result: "$2"
- pattern: "\"org.apache.pinot.common.metrics\"<type=\"ControllerMetrics\", name=\"pinot.controller.taskStatus.(\\w+)\\.(\\w+)\"><>(\\w+)"
name: "pinot_controller_taskStatus_$3"
cache: true
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
index ac933acf6a..2e84f8b510 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerMeter.java
@@ -61,6 +61,9 @@ public enum ControllerMeter implements AbstractMetrics.Meter {
LLC_SEGMENTS_DEEP_STORE_UPLOAD_RETRY_ERROR("LLCSegmentDeepStoreUploadRetryError", false),
SEGMENT_MISSING_DEEP_STORE_LINK("RealtimeSegmentMissingDeepStoreLink", false),
DELETED_TMP_SEGMENT_COUNT("DeletedTmpSegmentCount", false),
+ TABLE_REBALANCE_FAILURE_DETECTED("TableRebalanceFailureDetected", false),
+ TABLE_REBALANCE_RETRY("TableRebalanceRetry", false),
+ TABLE_REBALANCE_RETRY_TOO_MANY_TIMES("TableRebalanceRetryTooManyTimes", false),
NUMBER_ADHOC_TASKS_SUBMITTED("adhocTasks", false);
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerTimer.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerTimer.java
index 6627382db6..64608d0037 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerTimer.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ControllerTimer.java
@@ -26,6 +26,7 @@ import org.apache.pinot.common.Utils;
*
*/
public enum ControllerTimer implements AbstractMetrics.Timer {
+ TABLE_REBALANCE_EXECUTION_TIME_MS("tableRebalanceExecutionTimeMs", false),
CRON_SCHEDULER_JOB_EXECUTION_TIME_MS("cronSchedulerJobExecutionTimeMs", false);
private final String _timerName;
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 185b1d506e..2093aab0b4 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
@@ -93,6 +93,7 @@ import org.apache.pinot.controller.helix.core.minion.PinotTaskManager;
import org.apache.pinot.controller.helix.core.minion.TaskMetricsEmitter;
import org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
import org.apache.pinot.controller.helix.core.realtime.SegmentCompletionManager;
+import org.apache.pinot.controller.helix.core.rebalance.RebalanceChecker;
import org.apache.pinot.controller.helix.core.rebalance.tenant.DefaultTenantRebalancer;
import org.apache.pinot.controller.helix.core.rebalance.tenant.TenantRebalancer;
import org.apache.pinot.controller.helix.core.relocation.SegmentRelocator;
@@ -174,6 +175,7 @@ public abstract class BaseControllerStarter implements ServiceStartable {
protected SegmentRelocator _segmentRelocator;
protected RetentionManager _retentionManager;
protected SegmentStatusChecker _segmentStatusChecker;
+ protected RebalanceChecker _rebalanceChecker;
protected RealtimeConsumerMonitor _realtimeConsumerMonitor;
protected PinotTaskManager _taskManager;
protected TaskManagerStatusCache<TaskGeneratorMostRecentRunInfo> _taskManagerStatusCache;
@@ -818,6 +820,9 @@ public abstract class BaseControllerStarter implements ServiceStartable {
new SegmentStatusChecker(_helixResourceManager, _leadControllerManager, _config, _controllerMetrics,
_executorService);
periodicTasks.add(_segmentStatusChecker);
+ _rebalanceChecker = new RebalanceChecker(_helixResourceManager, _leadControllerManager, _config, _controllerMetrics,
+ _executorService);
+ periodicTasks.add(_rebalanceChecker);
_realtimeConsumerMonitor =
new RealtimeConsumerMonitor(_config, _helixResourceManager, _leadControllerManager, _controllerMetrics,
_executorService);
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java
index cb6ee99c4f..b8e2247b08 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java
@@ -178,9 +178,9 @@ public class ControllerConf extends PinotConfiguration {
public static final String DEPRECATED_SEGMENT_RELOCATOR_FREQUENCY_IN_SECONDS =
"controller.segment.relocator.frequencyInSeconds";
public static final String SEGMENT_RELOCATOR_FREQUENCY_PERIOD = "controller.segment.relocator.frequencyPeriod";
+ public static final String REBALANCE_CHECKER_FREQUENCY_PERIOD = "controller.rebalance.checker.frequencyPeriod";
// Because segment level validation is expensive and requires heavy ZK access, we run segment level validation
- // with a
- // separate interval
+ // with a separate interval
// Deprecated as of 0.8.0
@Deprecated
public static final String DEPRECATED_SEGMENT_LEVEL_VALIDATION_INTERVAL_IN_SECONDS =
@@ -199,6 +199,8 @@ public class ControllerConf extends PinotConfiguration {
// RealtimeSegmentRelocator has been rebranded as SegmentRelocator
public static final String DEPRECATED_REALTIME_SEGMENT_RELOCATION_INITIAL_DELAY_IN_SECONDS =
"controller.realtimeSegmentRelocation.initialDelayInSeconds";
+ public static final String REBALANCE_CHECKER_INITIAL_DELAY_IN_SECONDS =
+ "controller.rebalanceChecker.initialDelayInSeconds";
public static final String SEGMENT_RELOCATOR_INITIAL_DELAY_IN_SECONDS =
"controller.segmentRelocator.initialDelayInSeconds";
public static final String SEGMENT_RELOCATOR_ENABLE_LOCAL_TIER_MIGRATION =
@@ -237,7 +239,8 @@ public class ControllerConf extends PinotConfiguration {
private static final int DEFAULT_OFFLINE_SEGMENT_INTERVAL_CHECKER_FREQUENCY_IN_SECONDS = 24 * 60 * 60; // 24 Hours.
private static final int DEFAULT_REALTIME_SEGMENT_VALIDATION_FREQUENCY_IN_SECONDS = 60 * 60; // 1 Hour.
private static final int DEFAULT_BROKER_RESOURCE_VALIDATION_FREQUENCY_IN_SECONDS = 60 * 60; // 1 Hour.
- private static final int DEFAULT_STATUS_CONTROLLER_FREQUENCY_IN_SECONDS = 5 * 60; // 5 minutes
+ private static final int DEFAULT_STATUS_CHECKER_FREQUENCY_IN_SECONDS = 5 * 60; // 5 minutes
+ private static final int DEFAULT_REBALANCE_CHECKER_FREQUENCY_IN_SECONDS = 5 * 60; // 5 minutes
private static final int DEFAULT_TASK_METRICS_EMITTER_FREQUENCY_IN_SECONDS = 5 * 60; // 5 minutes
private static final int DEFAULT_STATUS_CONTROLLER_WAIT_FOR_PUSH_TIME_IN_SECONDS = 10 * 60; // 10 minutes
private static final int DEFAULT_TASK_MANAGER_FREQUENCY_IN_SECONDS = -1; // Disabled
@@ -604,7 +607,7 @@ public class ControllerConf extends PinotConfiguration {
return Optional.ofNullable(getProperty(ControllerPeriodicTasksConf.STATUS_CHECKER_FREQUENCY_PERIOD))
.map(period -> (int) convertPeriodToSeconds(period)).orElseGet(
() -> getProperty(ControllerPeriodicTasksConf.DEPRECATED_STATUS_CHECKER_FREQUENCY_IN_SECONDS,
- ControllerPeriodicTasksConf.DEFAULT_STATUS_CONTROLLER_FREQUENCY_IN_SECONDS));
+ ControllerPeriodicTasksConf.DEFAULT_STATUS_CHECKER_FREQUENCY_IN_SECONDS));
}
public void setStatusCheckerFrequencyInSeconds(int statusCheckerFrequencyInSeconds) {
@@ -612,6 +615,17 @@ public class ControllerConf extends PinotConfiguration {
Integer.toString(statusCheckerFrequencyInSeconds));
}
+ public int getRebalanceCheckerFrequencyInSeconds() {
+ return Optional.ofNullable(getProperty(ControllerPeriodicTasksConf.REBALANCE_CHECKER_FREQUENCY_PERIOD))
+ .map(period -> (int) convertPeriodToSeconds(period)).orElse(
+ ControllerPeriodicTasksConf.DEFAULT_REBALANCE_CHECKER_FREQUENCY_IN_SECONDS);
+ }
+
+ public long getRebalanceCheckerInitialDelayInSeconds() {
+ return getProperty(ControllerPeriodicTasksConf.REBALANCE_CHECKER_INITIAL_DELAY_IN_SECONDS,
+ ControllerPeriodicTasksConf.getRandomInitialDelayInSeconds());
+ }
+
public int getRealtimeConsumerMonitorRunFrequency() {
return Optional.ofNullable(getProperty(ControllerPeriodicTasksConf.RT_CONSUMER_MONITOR_FREQUENCY_PERIOD))
.map(period -> (int) convertPeriodToSeconds(period)).orElse(
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 58abb6c366..f0216da013 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
@@ -95,6 +95,7 @@ import org.apache.pinot.controller.helix.core.minion.PinotHelixTaskResourceManag
import org.apache.pinot.controller.helix.core.rebalance.RebalanceConfig;
import org.apache.pinot.controller.helix.core.rebalance.RebalanceJobConstants;
import org.apache.pinot.controller.helix.core.rebalance.RebalanceResult;
+import org.apache.pinot.controller.helix.core.rebalance.TableRebalanceContext;
import org.apache.pinot.controller.helix.core.rebalance.TableRebalanceProgressStats;
import org.apache.pinot.controller.helix.core.rebalance.TableRebalancer;
import org.apache.pinot.controller.recommender.RecommenderDriver;
@@ -615,11 +616,16 @@ public class PinotTableRestletResource {
@QueryParam("externalViewCheckIntervalInMs") long externalViewCheckIntervalInMs,
@ApiParam(value = "How long to wait till external view converges with ideal states") @DefaultValue("3600000")
@QueryParam("externalViewStabilizationTimeoutInMs") long externalViewStabilizationTimeoutInMs,
+ @ApiParam(value = "How often to make a status update (i.e. heartbeat)") @DefaultValue("300000")
+ @QueryParam("heartbeatIntervalInMs") long heartbeatIntervalInMs,
+ @ApiParam(value = "How long to wait for next status update (i.e. heartbeat) before the job is considered failed")
+ @DefaultValue("3600000") @QueryParam("heartbeatTimeoutInMs") long heartbeatTimeoutInMs,
+ @ApiParam(value = "Max number of attempts to rebalance") @DefaultValue("3") @QueryParam("maxAttempts")
+ int maxAttempts, @ApiParam(value = "Initial delay to exponentially backoff retry") @DefaultValue("300000")
+ @QueryParam("retryInitialDelayInMs") long retryInitialDelayInMs,
@ApiParam(value = "Whether to update segment target tier as part of the rebalance") @DefaultValue("false")
@QueryParam("updateTargetTier") boolean updateTargetTier) {
-
String tableNameWithType = constructTableNameWithType(tableName, tableTypeStr);
-
RebalanceConfig rebalanceConfig = new RebalanceConfig();
rebalanceConfig.setDryRun(dryRun);
rebalanceConfig.setReassignInstances(reassignInstances);
@@ -630,6 +636,12 @@ public class PinotTableRestletResource {
rebalanceConfig.setBestEfforts(bestEfforts);
rebalanceConfig.setExternalViewCheckIntervalInMs(externalViewCheckIntervalInMs);
rebalanceConfig.setExternalViewStabilizationTimeoutInMs(externalViewStabilizationTimeoutInMs);
+ heartbeatIntervalInMs = Math.max(externalViewCheckIntervalInMs, heartbeatIntervalInMs);
+ rebalanceConfig.setHeartbeatIntervalInMs(heartbeatIntervalInMs);
+ heartbeatTimeoutInMs = Math.max(heartbeatTimeoutInMs, 3 * heartbeatIntervalInMs);
+ rebalanceConfig.setHeartbeatTimeoutInMs(heartbeatTimeoutInMs);
+ rebalanceConfig.setMaxAttempts(maxAttempts);
+ rebalanceConfig.setRetryInitialDelayInMs(retryInitialDelayInMs);
rebalanceConfig.setUpdateTargetTier(updateTargetTier);
String rebalanceJobId = TableRebalancer.createUniqueRebalanceJobIdentifier();
@@ -666,6 +678,79 @@ public class PinotTableRestletResource {
}
}
+ @DELETE
+ @Produces(MediaType.APPLICATION_JSON)
+ @Authenticate(AccessType.UPDATE)
+ @Path("/tables/{tableName}/rebalance")
+ @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.CANCEL_REBALANCE)
+ @ApiOperation(value = "Cancel all rebalance jobs for the given table, and noop if no rebalance is running", notes =
+ "Cancel all rebalance jobs for the given table, and noop if no rebalance is running")
+ public List<String> cancelRebalance(
+ @ApiParam(value = "Name of the table to rebalance", required = true) @PathParam("tableName") String tableName,
+ @ApiParam(value = "OFFLINE|REALTIME", required = true) @QueryParam("type") String tableTypeStr) {
+ String tableNameWithType = constructTableNameWithType(tableName, tableTypeStr);
+ List<String> cancelledJobIds = new ArrayList<>();
+ boolean updated =
+ _pinotHelixResourceManager.updateJobsForTable(tableNameWithType, ControllerJobType.TABLE_REBALANCE,
+ jobMetadata -> {
+ String jobId = jobMetadata.get(CommonConstants.ControllerJob.JOB_ID);
+ try {
+ String jobStatsInStr = jobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS);
+ TableRebalanceProgressStats jobStats =
+ JsonUtils.stringToObject(jobStatsInStr, TableRebalanceProgressStats.class);
+ if (jobStats.getStatus() != RebalanceResult.Status.IN_PROGRESS) {
+ return;
+ }
+ cancelledJobIds.add(jobId);
+ LOGGER.info("Cancel rebalance job: {} for table: {}", jobId, tableNameWithType);
+ jobStats.setStatus(RebalanceResult.Status.CANCELLED);
+ jobMetadata.put(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS,
+ JsonUtils.objectToString(jobStats));
+ } catch (Exception e) {
+ LOGGER.error("Failed to cancel rebalance job: {} for table: {}", jobId, tableNameWithType, e);
+ }
+ });
+ LOGGER.info("Tried to cancel existing jobs at best effort and done: {}", updated);
+ return cancelledJobIds;
+ }
+
+ @GET
+ @Produces(MediaType.APPLICATION_JSON)
+ @Authenticate(AccessType.UPDATE)
+ @Path("/rebalanceStatus/{jobId}")
+ @Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.GET_REBALANCE_STATUS)
+ @ApiOperation(value = "Gets detailed stats of a rebalance operation",
+ notes = "Gets detailed stats of a rebalance operation")
+ public ServerRebalanceJobStatusResponse rebalanceStatus(
+ @ApiParam(value = "Rebalance Job Id", required = true) @PathParam("jobId") String jobId)
+ throws JsonProcessingException {
+ Map<String, String> controllerJobZKMetadata =
+ _pinotHelixResourceManager.getControllerJobZKMetadata(jobId, ControllerJobType.TABLE_REBALANCE);
+
+ if (controllerJobZKMetadata == null) {
+ throw new ControllerApplicationException(LOGGER, "Failed to find controller job id: " + jobId,
+ Response.Status.NOT_FOUND);
+ }
+ ServerRebalanceJobStatusResponse serverRebalanceJobStatusResponse = new ServerRebalanceJobStatusResponse();
+ TableRebalanceProgressStats tableRebalanceProgressStats = JsonUtils.stringToObject(
+ controllerJobZKMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS),
+ TableRebalanceProgressStats.class);
+ serverRebalanceJobStatusResponse.setTableRebalanceProgressStats(tableRebalanceProgressStats);
+
+ long timeSinceStartInSecs = 0L;
+ if (RebalanceResult.Status.DONE != tableRebalanceProgressStats.getStatus()) {
+ timeSinceStartInSecs = (System.currentTimeMillis() - tableRebalanceProgressStats.getStartTimeMs()) / 1000;
+ }
+ serverRebalanceJobStatusResponse.setTimeElapsedSinceStartInSeconds(timeSinceStartInSecs);
+
+ String jobCtxInStr = controllerJobZKMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_CONTEXT);
+ if (StringUtils.isNotEmpty(jobCtxInStr)) {
+ TableRebalanceContext jobCtx = JsonUtils.stringToObject(jobCtxInStr, TableRebalanceContext.class);
+ serverRebalanceJobStatusResponse.setTableRebalanceContext(jobCtx);
+ }
+ return serverRebalanceJobStatusResponse;
+ }
+
@GET
@Produces(MediaType.APPLICATION_JSON)
@Path("/tables/{tableName}/state")
@@ -724,37 +809,6 @@ public class PinotTableRestletResource {
}
}
- @GET
- @Produces(MediaType.APPLICATION_JSON)
- @Authenticate(AccessType.UPDATE)
- @Path("/rebalanceStatus/{jobId}")
- @Authorize(targetType = TargetType.CLUSTER, action = Actions.Cluster.GET_REBALANCE_STATUS)
- @ApiOperation(value = "Gets detailed stats of a rebalance operation",
- notes = "Gets detailed stats of a rebalance operation")
- public ServerRebalanceJobStatusResponse rebalanceStatus(
- @ApiParam(value = "Rebalance Job Id", required = true) @PathParam("jobId") String jobId)
- throws JsonProcessingException {
- Map<String, String> controllerJobZKMetadata =
- _pinotHelixResourceManager.getControllerJobZKMetadata(jobId, ControllerJobType.TABLE_REBALANCE);
-
- if (controllerJobZKMetadata == null) {
- throw new ControllerApplicationException(LOGGER, "Failed to find controller job id: " + jobId,
- Response.Status.NOT_FOUND);
- }
- TableRebalanceProgressStats tableRebalanceProgressStats = JsonUtils.stringToObject(
- controllerJobZKMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS),
- TableRebalanceProgressStats.class);
- long timeSinceStartInSecs = 0L;
- if (!RebalanceResult.Status.DONE.toString().equals(tableRebalanceProgressStats.getStatus())) {
- timeSinceStartInSecs = (System.currentTimeMillis() - tableRebalanceProgressStats.getStartTimeMs()) / 1000;
- }
-
- ServerRebalanceJobStatusResponse serverRebalanceJobStatusResponse = new ServerRebalanceJobStatusResponse();
- serverRebalanceJobStatusResponse.setTableRebalanceProgressStats(tableRebalanceProgressStats);
- serverRebalanceJobStatusResponse.setTimeElapsedSinceStartInSeconds(timeSinceStartInSecs);
- return serverRebalanceJobStatusResponse;
- }
-
@GET
@Path("/tables/{tableName}/stats")
@Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.GET_METADATA)
@@ -987,10 +1041,10 @@ public class PinotTableRestletResource {
}
Map<String, Map<String, String>> result = new HashMap<>();
for (String tableNameWithType : tableNamesWithType) {
- result.putAll(_pinotHelixResourceManager.getAllJobsForTable(tableNameWithType,
- jobTypesToFilter == null ? validJobTypes : jobTypesToFilter));
+ result.putAll(_pinotHelixResourceManager.getAllJobs(jobTypesToFilter == null ? validJobTypes : jobTypesToFilter,
+ jobMetadata -> jobMetadata.get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE)
+ .equals(tableNameWithType)));
}
-
return result;
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ServerRebalanceJobStatusResponse.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ServerRebalanceJobStatusResponse.java
index 3018b4877e..3551c9c811 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ServerRebalanceJobStatusResponse.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ServerRebalanceJobStatusResponse.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.controller.api.resources;
+import org.apache.pinot.controller.helix.core.rebalance.TableRebalanceContext;
import org.apache.pinot.controller.helix.core.rebalance.TableRebalanceProgressStats;
@@ -26,6 +27,8 @@ public class ServerRebalanceJobStatusResponse {
private TableRebalanceProgressStats _tableRebalanceProgressStats;
+ private TableRebalanceContext _tableRebalanceContext;
+
public void setTimeElapsedSinceStartInSeconds(Long timeElapsedSinceStart) {
_timeElapsedSinceStartInSeconds = timeElapsedSinceStart;
}
@@ -41,4 +44,12 @@ public class ServerRebalanceJobStatusResponse {
public long getTimeElapsedSinceStartInSeconds() {
return _timeElapsedSinceStartInSeconds;
}
+
+ public TableRebalanceContext getTableRebalanceContext() {
+ return _tableRebalanceContext;
+ }
+
+ public void setTableRebalanceContext(TableRebalanceContext tableRebalanceContext) {
+ _tableRebalanceContext = tableRebalanceContext;
+ }
}
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 bfb4918bf7..a85401d706 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
@@ -48,7 +48,9 @@ import java.util.TreeMap;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
import java.util.function.Function;
+import java.util.function.Predicate;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import javax.ws.rs.BadRequestException;
@@ -86,7 +88,6 @@ import org.apache.helix.model.StateModelDefinition;
import org.apache.helix.model.builder.HelixConfigScopeBuilder;
import org.apache.helix.store.zk.ZkHelixPropertyStore;
import org.apache.helix.zookeeper.datamodel.ZNRecord;
-import org.apache.helix.zookeeper.zkclient.exception.ZkNoNodeException;
import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils;
import org.apache.pinot.common.assignment.InstancePartitions;
import org.apache.pinot.common.assignment.InstancePartitionsUtils;
@@ -146,6 +147,7 @@ import org.apache.pinot.controller.helix.core.lineage.LineageManagerFactory;
import org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
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.TableRebalanceContext;
import org.apache.pinot.controller.helix.core.rebalance.TableRebalancer;
import org.apache.pinot.controller.helix.core.rebalance.ZkBasedTableRebalanceObserver;
import org.apache.pinot.controller.helix.core.util.ZKMetadataUtils;
@@ -1981,43 +1983,38 @@ public class PinotHelixResourceManager {
/**
* Returns the ZK metdata for the given jobId and jobType
* @param jobId the id of the job
- * @param jobType Job Path
+ * @param jobType the type of the job to figure out where job metadata is kept in ZK
* @return Map representing the job's ZK properties
*/
@Nullable
public Map<String, String> getControllerJobZKMetadata(String jobId, ControllerJobType jobType) {
- String controllerJobResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
- ZNRecord taskResourceZnRecord = _propertyStore.get(controllerJobResourcePath, null, AccessOption.PERSISTENT);
- if (taskResourceZnRecord != null) {
- return taskResourceZnRecord.getMapFields().get(jobId);
- }
- return null;
+ String jobResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
+ ZNRecord jobsZnRecord = _propertyStore.get(jobResourcePath, null, AccessOption.PERSISTENT);
+ return jobsZnRecord != null ? jobsZnRecord.getMapFields().get(jobId) : null;
}
/**
- * Returns a Map of jobId to job's ZK metadata for the given table
- * @param tableNameWithType the table for which jobs are to be fetched
+ * Returns a Map of jobId to job's ZK metadata that passes the checker, like for specific tables.
* @return A Map of jobId to job properties
*/
- public Map<String, Map<String, String>> getAllJobsForTable(String tableNameWithType,
- Set<ControllerJobType> jobTypes) {
+ public Map<String, Map<String, String>> getAllJobs(Set<ControllerJobType> jobTypes,
+ Predicate<Map<String, String>> jobMetadataChecker) {
Map<String, Map<String, String>> controllerJobs = new HashMap<>();
for (ControllerJobType jobType : jobTypes) {
- String jobsResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
- try {
- ZNRecord znRecord = _propertyStore.get(jobsResourcePath, null, -1);
- if (znRecord != null) {
- Map<String, Map<String, String>> tableJobsRecord = znRecord.getMapFields();
- for (Map.Entry<String, Map<String, String>> tableEntry : tableJobsRecord.entrySet()) {
- if (tableEntry.getValue().get(CommonConstants.ControllerJob.JOB_TYPE).equals(jobType.name())
- && tableEntry.getValue().get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE)
- .equals(tableNameWithType)) {
- controllerJobs.put(tableEntry.getKey(), tableEntry.getValue());
- }
- }
+ String jobResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
+ ZNRecord jobsZnRecord = _propertyStore.get(jobResourcePath, null, AccessOption.PERSISTENT);
+ if (jobsZnRecord == null) {
+ continue;
+ }
+ Map<String, Map<String, String>> jobMetadataMap = jobsZnRecord.getMapFields();
+ for (Map.Entry<String, Map<String, String>> jobMetadataEntry : jobMetadataMap.entrySet()) {
+ String jobId = jobMetadataEntry.getKey();
+ Map<String, String> jobMetadata = jobMetadataEntry.getValue();
+ Preconditions.checkState(jobMetadata.get(CommonConstants.ControllerJob.JOB_TYPE).equals(jobType.name()),
+ "Got unexpected jobType: %s at jobResourcePath: %s with jobId: %s", jobType.name(), jobResourcePath, jobId);
+ if (jobMetadataChecker.test(jobMetadata)) {
+ controllerJobs.put(jobId, jobMetadata);
}
- } catch (ZkNoNodeException e) {
- LOGGER.warn("Could not find controller job node for table : {} jobType: {}", tableNameWithType, jobType, e);
}
}
return controllerJobs;
@@ -2041,8 +2038,7 @@ public class PinotHelixResourceManager {
jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, Long.toString(jobSubmissionTimeMs));
jobMetadata.put(CommonConstants.ControllerJob.MESSAGE_COUNT, Integer.toString(numMessagesSent));
jobMetadata.put(CommonConstants.ControllerJob.SEGMENT_RELOAD_JOB_SEGMENT_NAME, segmentName);
- return addControllerJobToZK(jobId, jobMetadata,
- ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.RELOAD_SEGMENT));
+ return addControllerJobToZK(jobId, jobMetadata, ControllerJobType.RELOAD_SEGMENT);
}
public boolean addNewForceCommitJob(String tableNameWithType, String jobId, long jobSubmissionTimeMs,
@@ -2055,8 +2051,7 @@ public class PinotHelixResourceManager {
jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, Long.toString(jobSubmissionTimeMs));
jobMetadata.put(CommonConstants.ControllerJob.CONSUMING_SEGMENTS_FORCE_COMMITTED_LIST,
JsonUtils.objectToString(consumingSegmentsCommitted));
- return addControllerJobToZK(jobId, jobMetadata,
- ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.FORCE_COMMIT));
+ return addControllerJobToZK(jobId, jobMetadata, ControllerJobType.FORCE_COMMIT);
}
/**
@@ -2075,32 +2070,81 @@ public class PinotHelixResourceManager {
jobMetadata.put(CommonConstants.ControllerJob.JOB_TYPE, ControllerJobType.RELOAD_SEGMENT.toString());
jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, Long.toString(jobSubmissionTimeMs));
jobMetadata.put(CommonConstants.ControllerJob.MESSAGE_COUNT, Integer.toString(numberOfMessagesSent));
- return addControllerJobToZK(jobId, jobMetadata,
- ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.RELOAD_SEGMENT));
+ return addControllerJobToZK(jobId, jobMetadata, ControllerJobType.RELOAD_SEGMENT);
}
- public boolean addControllerJobToZK(String jobId, Map<String, String> jobMetadata, String jobResourcePath) {
+ /**
+ * Adds a new job metadata for controller job like table rebalance or reload into ZK
+ * @param jobId job's UUID
+ * @param jobMetadata the job metadata
+ * @param jobType the type of the job to figure out where job metadata is kept in ZK
+ * @return boolean representing success / failure of the ZK write step
+ */
+ public boolean addControllerJobToZK(String jobId, Map<String, String> jobMetadata, ControllerJobType jobType) {
+ return addControllerJobToZK(jobId, jobMetadata, jobType, prev -> true);
+ }
+
+ /**
+ * Adds a new job metadata for controller job like table rebalance or reload into ZK
+ * @param jobId job's UUID
+ * @param jobMetadata the job metadata
+ * @param jobType the type of the job to figure out where job metadata is kept in ZK
+ * @param prevJobMetadataChecker to check the previous job metadata before adding new one
+ * @return boolean representing success / failure of the ZK write step
+ */
+ public boolean addControllerJobToZK(String jobId, Map<String, String> jobMetadata, ControllerJobType jobType,
+ Predicate<Map<String, String>> prevJobMetadataChecker) {
Preconditions.checkState(jobMetadata.get(CommonConstants.ControllerJob.SUBMISSION_TIME_MS) != null,
"Submission Time in JobMetadata record not set. Cannot expire these records");
+ String jobResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
Stat stat = new Stat();
- ZNRecord tableJobsZnRecord = _propertyStore.get(jobResourcePath, stat, AccessOption.PERSISTENT);
- if (tableJobsZnRecord != null) {
- Map<String, Map<String, String>> tasks = tableJobsZnRecord.getMapFields();
- tasks.put(jobId, jobMetadata);
- if (tasks.size() > CommonConstants.ControllerJob.MAXIMUM_CONTROLLER_JOBS_IN_ZK) {
- tasks = tasks.entrySet().stream().sorted((v1, v2) -> Long.compare(
+ ZNRecord jobsZnRecord = _propertyStore.get(jobResourcePath, stat, AccessOption.PERSISTENT);
+ if (jobsZnRecord != null) {
+ Map<String, Map<String, String>> jobMetadataMap = jobsZnRecord.getMapFields();
+ Map<String, String> prevJobMetadata = jobMetadataMap.get(jobId);
+ if (!prevJobMetadataChecker.test(prevJobMetadata)) {
+ return false;
+ }
+ jobMetadataMap.put(jobId, jobMetadata);
+ if (jobMetadataMap.size() > CommonConstants.ControllerJob.MAXIMUM_CONTROLLER_JOBS_IN_ZK) {
+ jobMetadataMap = jobMetadataMap.entrySet().stream().sorted((v1, v2) -> Long.compare(
Long.parseLong(v2.getValue().get(CommonConstants.ControllerJob.SUBMISSION_TIME_MS)),
Long.parseLong(v1.getValue().get(CommonConstants.ControllerJob.SUBMISSION_TIME_MS))))
.collect(Collectors.toList()).subList(0, CommonConstants.ControllerJob.MAXIMUM_CONTROLLER_JOBS_IN_ZK)
.stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
}
- tableJobsZnRecord.setMapFields(tasks);
- return _propertyStore.set(jobResourcePath, tableJobsZnRecord, stat.getVersion(), AccessOption.PERSISTENT);
+ jobsZnRecord.setMapFields(jobMetadataMap);
+ return _propertyStore.set(jobResourcePath, jobsZnRecord, stat.getVersion(), AccessOption.PERSISTENT);
} else {
- tableJobsZnRecord = new ZNRecord(jobResourcePath);
- tableJobsZnRecord.setMapField(jobId, jobMetadata);
- return _propertyStore.set(jobResourcePath, tableJobsZnRecord, AccessOption.PERSISTENT);
+ jobsZnRecord = new ZNRecord(jobResourcePath);
+ jobsZnRecord.setMapField(jobId, jobMetadata);
+ return _propertyStore.set(jobResourcePath, jobsZnRecord, AccessOption.PERSISTENT);
+ }
+ }
+
+ /**
+ * Update existing job metadata belong to the table
+ * @param tableNameWithType whose job metadata to be updated
+ * @param jobType the type of the job to figure out where job metadata is kept in ZK
+ * @param updater to modify the job metadata in place
+ * @return boolean representing success / failure of the ZK write step
+ */
+ public boolean updateJobsForTable(String tableNameWithType, ControllerJobType jobType,
+ Consumer<Map<String, String>> updater) {
+ String jobResourcePath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(jobType);
+ Stat stat = new Stat();
+ ZNRecord jobsZnRecord = _propertyStore.get(jobResourcePath, stat, AccessOption.PERSISTENT);
+ if (jobsZnRecord == null) {
+ return true;
+ }
+ Map<String, Map<String, String>> jobMetadataMap = jobsZnRecord.getMapFields();
+ for (Map<String, String> jobMetadata : jobMetadataMap.values()) {
+ if (jobMetadata.get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE).equals(tableNameWithType)) {
+ updater.accept(jobMetadata);
+ }
}
+ jobsZnRecord.setMapFields(jobMetadataMap);
+ return _propertyStore.set(jobResourcePath, jobsZnRecord, stat.getVersion(), AccessOption.PERSISTENT);
}
@VisibleForTesting
@@ -3112,12 +3156,19 @@ public class PinotHelixResourceManager {
throw new TableNotFoundException("Failed to find table config for table: " + tableNameWithType);
}
Preconditions.checkState(rebalanceJobId != null, "RebalanceId not populated in the rebalanceConfig");
- if (rebalanceConfig.isUpdateTargetTier()) {
- updateTargetTier(rebalanceJobId, tableNameWithType, tableConfig);
- }
ZkBasedTableRebalanceObserver zkBasedTableRebalanceObserver = null;
if (trackRebalanceProgress) {
- zkBasedTableRebalanceObserver = new ZkBasedTableRebalanceObserver(tableNameWithType, rebalanceJobId, this);
+ zkBasedTableRebalanceObserver = new ZkBasedTableRebalanceObserver(tableNameWithType, rebalanceJobId,
+ TableRebalanceContext.forInitialAttempt(rebalanceJobId, rebalanceConfig), this);
+ }
+ return rebalanceTable(tableNameWithType, tableConfig, rebalanceJobId, rebalanceConfig,
+ zkBasedTableRebalanceObserver);
+ }
+
+ public RebalanceResult rebalanceTable(String tableNameWithType, TableConfig tableConfig, String rebalanceJobId,
+ RebalanceConfig rebalanceConfig, @Nullable ZkBasedTableRebalanceObserver zkBasedTableRebalanceObserver) {
+ if (rebalanceConfig.isUpdateTargetTier()) {
+ updateTargetTier(rebalanceJobId, tableNameWithType, tableConfig);
}
TableRebalancer tableRebalancer =
new TableRebalancer(_helixZkManager, zkBasedTableRebalanceObserver, _controllerMetrics);
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/NoOpTableRebalanceObserver.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/NoOpTableRebalanceObserver.java
index 9bb9179c90..a435e35ead 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/NoOpTableRebalanceObserver.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/NoOpTableRebalanceObserver.java
@@ -36,4 +36,14 @@ public class NoOpTableRebalanceObserver implements TableRebalanceObserver {
@Override
public void onError(String errorMsg) {
}
+
+ @Override
+ public boolean isStopped() {
+ return false;
+ }
+
+ @Override
+ public RebalanceResult.Status getStopStatus() {
+ throw new UnsupportedOperationException();
+ }
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceChecker.java
new file mode 100644
index 0000000000..93e4b04dcb
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceChecker.java
@@ -0,0 +1,331 @@
+/**
+ * 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.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import org.apache.commons.lang3.RandomUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pinot.common.metadata.controllerjob.ControllerJobType;
+import org.apache.pinot.common.metrics.ControllerGauge;
+import org.apache.pinot.common.metrics.ControllerMeter;
+import org.apache.pinot.common.metrics.ControllerMetrics;
+import org.apache.pinot.controller.ControllerConf;
+import org.apache.pinot.controller.LeadControllerManager;
+import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
+import org.apache.pinot.controller.helix.core.periodictask.ControllerPeriodicTask;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Periodic task to check whether a user triggered rebalance job is completed or not, and retry if failed. The retry
+ * job is started with the same rebalance configs provided by the user and does best effort to stop the other jobs
+ * for the same table. This task can be configured to just check failures and report metrics, and not to do retry.
+ */
+public class RebalanceChecker extends ControllerPeriodicTask<Void> {
+ private static final Logger LOGGER = LoggerFactory.getLogger(RebalanceChecker.class);
+ private static final double RETRY_DELAY_SCALE_FACTOR = 2.0;
+ private final ExecutorService _executorService;
+
+ public RebalanceChecker(PinotHelixResourceManager pinotHelixResourceManager,
+ LeadControllerManager leadControllerManager, ControllerConf config, ControllerMetrics controllerMetrics,
+ ExecutorService executorService) {
+ super(RebalanceChecker.class.getSimpleName(), config.getRebalanceCheckerFrequencyInSeconds(),
+ config.getRebalanceCheckerInitialDelayInSeconds(), pinotHelixResourceManager, leadControllerManager,
+ controllerMetrics);
+ _executorService = executorService;
+ }
+
+ @Override
+ protected void processTables(List<String> tableNamesWithType, Properties periodicTaskProperties) {
+ int numTables = tableNamesWithType.size();
+ LOGGER.info("Processing {} tables in task: {}", numTables, _taskName);
+ int numTablesProcessed = retryRebalanceTables(new HashSet<>(tableNamesWithType));
+ _controllerMetrics.setValueOfGlobalGauge(ControllerGauge.PERIODIC_TASK_NUM_TABLES_PROCESSED, _taskName,
+ numTablesProcessed);
+ LOGGER.info("Finish processing {}/{} tables in task: {}", numTablesProcessed, numTables, _taskName);
+ }
+
+ /**
+ * Rare but the task may be executed by more than one threads because user can trigger the periodic task to run
+ * immediately, in addition to the one scheduled to run periodically. So make this method synchronized to be simple.
+ */
+ private synchronized int retryRebalanceTables(Set<String> tableNamesWithType) {
+ // Get all jobMetadata for all the given tables with a single ZK read.
+ Map<String, Map<String, String>> allJobMetadataByJobId =
+ _pinotHelixResourceManager.getAllJobs(Collections.singleton(ControllerJobType.TABLE_REBALANCE),
+ jobMetadata -> tableNamesWithType.contains(
+ jobMetadata.get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE)));
+ Map<String, Map<String, Map<String, String>>> tableJobMetadataMap = new HashMap<>();
+ allJobMetadataByJobId.forEach((jobId, jobMetadata) -> {
+ String tableNameWithType = jobMetadata.get(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE);
+ tableJobMetadataMap.computeIfAbsent(tableNameWithType, k -> new HashMap<>()).put(jobId, jobMetadata);
+ });
+ int numTablesProcessed = 0;
+ for (String tableNameWithType : tableNamesWithType) {
+ try {
+ LOGGER.info("Start to retry rebalance for table: {}", tableNameWithType);
+ // Get all rebalance jobs as tracked in ZK for this table to check if there is any failure to retry.
+ Map<String, Map<String, String>> allJobMetadata = tableJobMetadataMap.get(tableNameWithType);
+ if (allJobMetadata.isEmpty()) {
+ LOGGER.info("No rebalance job has been triggered for table: {}. Skip retry", tableNameWithType);
+ continue;
+ }
+ retryRebalanceTable(tableNameWithType, allJobMetadata);
+ numTablesProcessed++;
+ } catch (Exception e) {
+ LOGGER.error("Failed to retry rebalance for table: {}", tableNameWithType, e);
+ _controllerMetrics.addMeteredTableValue(tableNameWithType + "." + _taskName,
+ ControllerMeter.PERIODIC_TASK_ERROR, 1L);
+ }
+ }
+ return numTablesProcessed;
+ }
+
+ @VisibleForTesting
+ void retryRebalanceTable(String tableNameWithType, Map<String, Map<String, String>> allJobMetadata)
+ throws Exception {
+ // Skip retry for the table if rebalance job is still running or has completed, in specific:
+ // 1) Skip retry if any rebalance job is actively running. Being actively running means the job is at IN_PROGRESS
+ // status, and has updated its status kept in ZK within the heartbeat timeout. It's possible that more than one
+ // rebalance jobs are running for the table, but that's fine with idempotent rebalance algorithm.
+ // 2) Skip retry if the most recently started rebalance job has completed with DONE or NO_OP. It's possible that
+ // jobs started earlier may be still running, but they are ignored here.
+ //
+ // Otherwise, we can get a list of failed rebalance jobs, i.e. those at FAILED status; or IN_PROGRESS status but
+ // haven't updated their status kept in ZK within the heartbeat timeout. For those candidate jobs to retry:
+ // 1) Firstly, group them by the original jobIds they retry for so that we can skip those exceeded maxRetry.
+ // 2) For the remaining jobs, we take the one started most recently and retry it with its original configs.
+ // 3) If configured, we can abort the other rebalance jobs for the table by setting their status to FAILED.
+ Map<String/*original jobId*/, Set<Pair<TableRebalanceContext/*job attempts*/, Long
+ /*startTime*/>>> candidateJobs = getCandidateJobs(tableNameWithType, allJobMetadata);
+ if (candidateJobs.isEmpty()) {
+ LOGGER.info("Found no failed rebalance jobs for table: {}. Skip retry", tableNameWithType);
+ return;
+ }
+ _controllerMetrics.addMeteredTableValue(tableNameWithType, ControllerMeter.TABLE_REBALANCE_FAILURE_DETECTED, 1L);
+ Pair<TableRebalanceContext, Long> jobContextAndStartTime = getLatestJob(candidateJobs);
+ if (jobContextAndStartTime == null) {
+ LOGGER.info("Rebalance has been retried too many times for table: {}. Skip retry", tableNameWithType);
+ _controllerMetrics.addMeteredTableValue(tableNameWithType, ControllerMeter.TABLE_REBALANCE_RETRY_TOO_MANY_TIMES,
+ 1L);
+ return;
+ }
+ TableRebalanceContext jobCtx = jobContextAndStartTime.getLeft();
+ String prevJobId = jobCtx.getJobId();
+ RebalanceConfig rebalanceConfig = jobCtx.getConfig();
+ long jobStartTimeMs = jobContextAndStartTime.getRight();
+ long retryDelayMs = getRetryDelayInMs(rebalanceConfig.getRetryInitialDelayInMs(), jobCtx.getAttemptId());
+ if (jobStartTimeMs + retryDelayMs > System.currentTimeMillis()) {
+ LOGGER.info("Delay retry for failed rebalance job: {} that started at: {}, by: {}ms", prevJobId, jobStartTimeMs,
+ retryDelayMs);
+ return;
+ }
+ abortExistingJobs(tableNameWithType, _pinotHelixResourceManager);
+ // Get tableConfig only when the table needs to retry rebalance, and get it before submitting rebalance to another
+ // thread, in order to avoid unnecessary ZK reads and making too many ZK reads in a short time.
+ TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType);
+ Preconditions.checkState(tableConfig != null, "Failed to find table config for table: {}", tableNameWithType);
+ _executorService.submit(() -> {
+ // Retry rebalance in another thread as rebalance can take time.
+ try {
+ retryRebalanceTableWithContext(tableNameWithType, tableConfig, jobCtx);
+ } catch (Throwable t) {
+ LOGGER.error("Failed to retry rebalance for table: {} asynchronously", tableNameWithType, t);
+ }
+ });
+ }
+
+ private void retryRebalanceTableWithContext(String tableNameWithType, TableConfig tableConfig,
+ TableRebalanceContext jobCtx) {
+ String prevJobId = jobCtx.getJobId();
+ RebalanceConfig rebalanceConfig = jobCtx.getConfig();
+ TableRebalanceContext retryCtx =
+ TableRebalanceContext.forRetry(jobCtx.getOriginalJobId(), rebalanceConfig, jobCtx.getAttemptId() + 1);
+ String attemptJobId = retryCtx.getJobId();
+ LOGGER.info("Retry rebalance job: {} for table: {} with attempt job: {}", prevJobId, tableNameWithType,
+ attemptJobId);
+ _controllerMetrics.addMeteredTableValue(tableNameWithType, ControllerMeter.TABLE_REBALANCE_RETRY, 1L);
+ ZkBasedTableRebalanceObserver observer =
+ new ZkBasedTableRebalanceObserver(tableNameWithType, attemptJobId, retryCtx, _pinotHelixResourceManager);
+ RebalanceResult result =
+ _pinotHelixResourceManager.rebalanceTable(tableNameWithType, tableConfig, attemptJobId, rebalanceConfig,
+ observer);
+ LOGGER.info("New attempt: {} for table: {} is done with result status: {}", attemptJobId, tableNameWithType,
+ result.getStatus());
+ }
+
+ @VisibleForTesting
+ static long getRetryDelayInMs(long initDelayMs, int attemptId) {
+ // TODO: Just exponential backoff by factor 2 for now. Can add other retry polices as needed.
+ // The attemptId starts from 1, so minus one as the exponent.
+ double minDelayMs = initDelayMs * Math.pow(RETRY_DELAY_SCALE_FACTOR, attemptId - 1);
+ double maxDelayMs = minDelayMs * RETRY_DELAY_SCALE_FACTOR;
+ return RandomUtils.nextLong((long) minDelayMs, (long) maxDelayMs);
+ }
+
+ private static void abortExistingJobs(String tableNameWithType, PinotHelixResourceManager pinotHelixResourceManager) {
+ boolean updated = pinotHelixResourceManager.updateJobsForTable(tableNameWithType, ControllerJobType.TABLE_REBALANCE,
+ jobMetadata -> {
+ String jobId = jobMetadata.get(CommonConstants.ControllerJob.JOB_ID);
+ try {
+ String jobStatsInStr = jobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS);
+ TableRebalanceProgressStats jobStats =
+ JsonUtils.stringToObject(jobStatsInStr, TableRebalanceProgressStats.class);
+ if (jobStats.getStatus() != RebalanceResult.Status.IN_PROGRESS) {
+ return;
+ }
+ LOGGER.info("Abort rebalance job: {} for table: {}", jobId, tableNameWithType);
+ jobStats.setStatus(RebalanceResult.Status.ABORTED);
+ jobMetadata.put(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS,
+ JsonUtils.objectToString(jobStats));
+ } catch (Exception e) {
+ LOGGER.error("Failed to abort rebalance job: {} for table: {}", jobId, tableNameWithType, e);
+ }
+ });
+ LOGGER.info("Tried to abort existing jobs at best effort and done: {}", updated);
+ }
+
+ @VisibleForTesting
+ static Pair<TableRebalanceContext, Long> getLatestJob(
+ Map<String, Set<Pair<TableRebalanceContext, Long>>> candidateJobs) {
+ Pair<TableRebalanceContext, Long> candidateJobRun = null;
+ for (Map.Entry<String, Set<Pair<TableRebalanceContext, Long>>> entry : candidateJobs.entrySet()) {
+ // The job configs from all retry jobs are same, as the same set of job configs is used to do retry.
+ // The job metadata kept in ZK is cleaned by submission time order gradually, so we can't compare Set.size()
+ // against maxAttempts, but check retryNum of each run to see if retries have exceeded limit.
+ Set<Pair<TableRebalanceContext, Long>> jobRuns = entry.getValue();
+ int maxAttempts = jobRuns.iterator().next().getLeft().getConfig().getMaxAttempts();
+ Pair<TableRebalanceContext, Long> latestJobRun = null;
+ for (Pair<TableRebalanceContext, Long> jobRun : jobRuns) {
+ if (jobRun.getLeft().getAttemptId() >= maxAttempts) {
+ latestJobRun = null;
+ break;
+ }
+ if (latestJobRun == null || latestJobRun.getRight() < jobRun.getRight()) {
+ latestJobRun = jobRun;
+ }
+ }
+ if (latestJobRun == null) {
+ LOGGER.info("Rebalance job: {} had exceeded maxAttempts: {}. Skip retry", entry.getKey(), maxAttempts);
+ continue;
+ }
+ if (candidateJobRun == null || candidateJobRun.getRight() < latestJobRun.getRight()) {
+ candidateJobRun = latestJobRun;
+ }
+ }
+ return candidateJobRun;
+ }
+
+ @VisibleForTesting
+ static Map<String, Set<Pair<TableRebalanceContext, Long>>> getCandidateJobs(String tableNameWithType,
+ Map<String, Map<String, String>> allJobMetadata)
+ throws Exception {
+ long nowMs = System.currentTimeMillis();
+ Map<String, Set<Pair<TableRebalanceContext, Long>>> candidates = new HashMap<>();
+ // If the job started most recently has already completed, then skip retry for the table.
+ Pair<String, Long> latestStartedJob = null;
+ Pair<String, Long> latestCompletedJob = null;
+ // The processing order of job metadata from the given Map is not deterministic. Track the completed original
+ // jobs so that we can simply skip the retry jobs belonging to the completed original jobs.
+ Map<String, String> completedOriginalJobs = new HashMap<>();
+ Set<String> cancelledOriginalJobs = new HashSet<>();
+ for (Map.Entry<String, Map<String, String>> entry : allJobMetadata.entrySet()) {
+ String jobId = entry.getKey();
+ Map<String, String> jobMetadata = entry.getValue();
+ long statsUpdatedAt = Long.parseLong(jobMetadata.get(CommonConstants.ControllerJob.SUBMISSION_TIME_MS));
+ String jobStatsInStr = jobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS);
+ if (StringUtils.isEmpty(jobStatsInStr)) {
+ LOGGER.info("Skip rebalance job: {} as it has no job progress stats", jobId);
+ continue;
+ }
+ String jobCtxInStr = jobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_CONTEXT);
+ if (StringUtils.isEmpty(jobCtxInStr)) {
+ LOGGER.info("Skip rebalance job: {} as it has no job context", jobId);
+ continue;
+ }
+ TableRebalanceProgressStats jobStats = JsonUtils.stringToObject(jobStatsInStr, TableRebalanceProgressStats.class);
+ TableRebalanceContext jobCtx = JsonUtils.stringToObject(jobCtxInStr, TableRebalanceContext.class);
+ long jobStartTimeMs = jobStats.getStartTimeMs();
+ if (latestStartedJob == null || latestStartedJob.getRight() < jobStartTimeMs) {
+ latestStartedJob = Pair.of(jobId, jobStartTimeMs);
+ }
+ String originalJobId = jobCtx.getOriginalJobId();
+ RebalanceResult.Status jobStatus = jobStats.getStatus();
+ if (jobStatus == RebalanceResult.Status.DONE || jobStatus == RebalanceResult.Status.NO_OP) {
+ LOGGER.info("Skip rebalance job: {} as it has completed with status: {}", jobId, jobStatus);
+ completedOriginalJobs.put(originalJobId, jobId);
+ if (latestCompletedJob == null || latestCompletedJob.getRight() < jobStartTimeMs) {
+ latestCompletedJob = Pair.of(jobId, jobStartTimeMs);
+ }
+ continue;
+ }
+ if (jobStatus == RebalanceResult.Status.FAILED || jobStatus == RebalanceResult.Status.ABORTED) {
+ LOGGER.info("Found rebalance job: {} for original job: {} has been stopped with status: {}", jobId,
+ originalJobId, jobStatus);
+ candidates.computeIfAbsent(originalJobId, (k) -> new HashSet<>()).add(Pair.of(jobCtx, jobStartTimeMs));
+ continue;
+ }
+ if (jobStatus == RebalanceResult.Status.CANCELLED) {
+ LOGGER.info("Found cancelled rebalance job: {} for original job: {}", jobId, originalJobId);
+ cancelledOriginalJobs.add(originalJobId);
+ continue;
+ }
+ // Check if an IN_PROGRESS job is still actively running.
+ long heartbeatTimeoutMs = jobCtx.getConfig().getHeartbeatTimeoutInMs();
+ if (nowMs - statsUpdatedAt < heartbeatTimeoutMs) {
+ LOGGER.info("Rebalance job: {} is actively running with status updated at: {} within timeout: {}. Skip "
+ + "retry for table: {}", jobId, statsUpdatedAt, heartbeatTimeoutMs, tableNameWithType);
+ return Collections.emptyMap();
+ }
+ // The job is considered failed, but it's possible it is still running, then we might end up with more than one
+ // rebalance jobs running in parallel for a table. The rebalance algorithm is idempotent, so this should be fine
+ // for the correctness.
+ LOGGER.info("Found stuck rebalance job: {} for original job: {}", jobId, originalJobId);
+ candidates.computeIfAbsent(originalJobId, (k) -> new HashSet<>()).add(Pair.of(jobCtx, jobStartTimeMs));
+ }
+ if (latestCompletedJob != null && latestCompletedJob.getLeft().equals(latestStartedJob.getLeft())) {
+ LOGGER.info("Rebalance job: {} started most recently has already done. Skip retry for table: {}",
+ latestCompletedJob.getLeft(), tableNameWithType);
+ return Collections.emptyMap();
+ }
+ for (String jobId : cancelledOriginalJobs) {
+ LOGGER.info("Skip original job: {} as it's cancelled", jobId);
+ candidates.remove(jobId);
+ }
+ for (Map.Entry<String, String> entry : completedOriginalJobs.entrySet()) {
+ LOGGER.info("Skip original job: {} as it's completed by attempt: {}", entry.getKey(), entry.getValue());
+ candidates.remove(entry.getKey());
+ }
+ return candidates;
+ }
+}
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 78e3b19bca..94709f065c 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
@@ -84,6 +84,24 @@ public class RebalanceConfig {
@ApiModelProperty(example = "false")
private boolean _updateTargetTier = false;
+ // Update job status every this interval as heartbeat, to indicate the job is still actively running.
+ @JsonProperty("heartbeatIntervalInMs")
+ @ApiModelProperty(example = "300000")
+ private long _heartbeatIntervalInMs = 300000L;
+
+ // The job is considered as failed if not updating its status by this timeout, even though it's IN_PROGRESS status.
+ @JsonProperty("heartbeatTimeoutInMs")
+ @ApiModelProperty(example = "3600000")
+ private long _heartbeatTimeoutInMs = 3600000L;
+
+ @JsonProperty("maxAttempts")
+ @ApiModelProperty(example = "3")
+ private int _maxAttempts = 3;
+
+ @JsonProperty("retryInitialDelayInMs")
+ @ApiModelProperty(example = "300000")
+ private long _retryInitialDelayInMs = 300000L;
+
public boolean isDryRun() {
return _dryRun;
}
@@ -164,6 +182,50 @@ public class RebalanceConfig {
_updateTargetTier = updateTargetTier;
}
+ public long getHeartbeatIntervalInMs() {
+ return _heartbeatIntervalInMs;
+ }
+
+ public void setHeartbeatIntervalInMs(long heartbeatIntervalInMs) {
+ _heartbeatIntervalInMs = heartbeatIntervalInMs;
+ }
+
+ public long getHeartbeatTimeoutInMs() {
+ return _heartbeatTimeoutInMs;
+ }
+
+ public void setHeartbeatTimeoutInMs(long heartbeatTimeoutInMs) {
+ _heartbeatTimeoutInMs = heartbeatTimeoutInMs;
+ }
+
+ public int getMaxAttempts() {
+ return _maxAttempts;
+ }
+
+ public void setMaxAttempts(int maxAttempts) {
+ _maxAttempts = maxAttempts;
+ }
+
+ public long getRetryInitialDelayInMs() {
+ return _retryInitialDelayInMs;
+ }
+
+ public void setRetryInitialDelayInMs(long retryInitialDelayInMs) {
+ _retryInitialDelayInMs = retryInitialDelayInMs;
+ }
+
+ @Override
+ public String toString() {
+ return "RebalanceConfig{" + "_dryRun=" + _dryRun + ", _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="
+ + _retryInitialDelayInMs + '}';
+ }
+
public static RebalanceConfig copy(RebalanceConfig cfg) {
RebalanceConfig rc = new RebalanceConfig();
rc._dryRun = cfg._dryRun;
@@ -176,6 +238,10 @@ public class RebalanceConfig {
rc._externalViewCheckIntervalInMs = cfg._externalViewCheckIntervalInMs;
rc._externalViewStabilizationTimeoutInMs = cfg._externalViewStabilizationTimeoutInMs;
rc._updateTargetTier = cfg._updateTargetTier;
+ rc._heartbeatIntervalInMs = cfg._heartbeatIntervalInMs;
+ rc._heartbeatTimeoutInMs = cfg._heartbeatTimeoutInMs;
+ rc._maxAttempts = cfg._maxAttempts;
+ rc._retryInitialDelayInMs = cfg._retryInitialDelayInMs;
return rc;
}
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceJobConstants.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceJobConstants.java
index 8070d44972..68314bbba2 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceJobConstants.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceJobConstants.java
@@ -24,4 +24,6 @@ public class RebalanceJobConstants {
// Progress status of the rebalance operartion
public static final String JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS = "REBALANCE_PROGRESS_STATS";
+ // Configs to retry the rebalance operartion
+ public static final String JOB_METADATA_KEY_REBALANCE_CONTEXT = "REBALANCE_CONTEXT";
}
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 ed3ad624d8..2be7fc7753 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
@@ -87,6 +87,10 @@ public class RebalanceResult {
}
public enum Status {
- NO_OP, DONE, FAILED, IN_PROGRESS
+ // FAILED if the job has ended with known exceptions;
+ // ABORTED if the job is stopped by others but retry is still allowed;
+ // CANCELLED if the job is stopped by user, and retry is cancelled too;
+ // UNKNOWN_ERROR if the job hits on an unexpected exception.
+ NO_OP, DONE, FAILED, IN_PROGRESS, ABORTED, CANCELLED, UNKNOWN_ERROR
}
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceContext.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceContext.java
new file mode 100644
index 0000000000..48e43558cc
--- /dev/null
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceContext.java
@@ -0,0 +1,94 @@
+/**
+ * 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;
+
+/**
+ * Track the job configs and attempt numbers as part of the job ZK metadata to retry failed rebalance.
+ */
+public class TableRebalanceContext {
+ private static final int INITIAL_ATTEMPT_ID = 1;
+ private String _jobId;
+ private String _originalJobId;
+ private RebalanceConfig _config;
+ private int _attemptId;
+
+ public static TableRebalanceContext forInitialAttempt(String originalJobId, RebalanceConfig config) {
+ return new TableRebalanceContext(originalJobId, config, INITIAL_ATTEMPT_ID);
+ }
+
+ public static TableRebalanceContext forRetry(String originalJobId, RebalanceConfig config, int attemptId) {
+ return new TableRebalanceContext(originalJobId, config, attemptId);
+ }
+
+ public TableRebalanceContext() {
+ // For JSON deserialization.
+ }
+
+ private TableRebalanceContext(String originalJobId, RebalanceConfig config, int attemptId) {
+ _jobId = createAttemptJobId(originalJobId, attemptId);
+ _originalJobId = originalJobId;
+ _config = config;
+ _attemptId = attemptId;
+ }
+
+ public int getAttemptId() {
+ return _attemptId;
+ }
+
+ public void setAttemptId(int attemptId) {
+ _attemptId = attemptId;
+ }
+
+ public String getOriginalJobId() {
+ return _originalJobId;
+ }
+
+ public void setOriginalJobId(String originalJobId) {
+ _originalJobId = originalJobId;
+ }
+
+ public String getJobId() {
+ return _jobId;
+ }
+
+ public void setJobId(String jobId) {
+ _jobId = jobId;
+ }
+
+ public RebalanceConfig getConfig() {
+ return _config;
+ }
+
+ public void setConfig(RebalanceConfig config) {
+ _config = config;
+ }
+
+ @Override
+ public String toString() {
+ return "TableRebalanceContext{" + "_jobId='" + _jobId + '\'' + ", _originalJobId='" + _originalJobId + '\''
+ + ", _config=" + _config + ", _attemptId=" + _attemptId + '}';
+ }
+
+ private static String createAttemptJobId(String originalJobId, int attemptId) {
+ if (attemptId == INITIAL_ATTEMPT_ID) {
+ return originalJobId;
+ }
+ return originalJobId + "_" + attemptId;
+ }
+}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceObserver.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceObserver.java
index d52dce1f6e..e9c5c299cf 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceObserver.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceObserver.java
@@ -43,4 +43,8 @@ public interface TableRebalanceObserver {
void onSuccess(String msg);
void onError(String errorMsg);
+
+ boolean isStopped();
+
+ RebalanceResult.Status getStopStatus();
}
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceProgressStats.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceProgressStats.java
index 6d6f70b5c7..aad369ae33 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceProgressStats.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalanceProgressStats.java
@@ -42,7 +42,7 @@ public class TableRebalanceProgressStats {
}
// Done/In_progress/Failed
- private String _status;
+ private RebalanceResult.Status _status;
// When did Rebalance start
private long _startTimeMs;
// How long did rebalance take
@@ -62,7 +62,7 @@ public class TableRebalanceProgressStats {
_initialToTargetStateConvergence = new RebalanceStateStats();
}
- public void setStatus(String status) {
+ public void setStatus(RebalanceResult.Status status) {
_status = status;
}
@@ -90,7 +90,7 @@ public class TableRebalanceProgressStats {
_completionStatusMsg = completionStatusMsg;
}
- public String getStatus() {
+ public RebalanceResult.Status getStatus() {
return _status;
}
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 a86c3e27a1..a27c4ffbfb 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
@@ -29,6 +29,7 @@ import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.UUID;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.ToIntFunction;
import javax.annotation.Nullable;
@@ -47,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.metrics.ControllerMetrics;
+import org.apache.pinot.common.metrics.ControllerTimer;
import org.apache.pinot.common.tier.PinotServerTierStorage;
import org.apache.pinot.common.tier.Tier;
import org.apache.pinot.common.tier.TierFactory;
@@ -138,6 +140,24 @@ public class TableRebalancer {
public RebalanceResult rebalance(TableConfig tableConfig, RebalanceConfig rebalanceConfig,
@Nullable String rebalanceJobId) {
+ long startTime = System.currentTimeMillis();
+ String tableNameWithType = tableConfig.getTableName();
+ RebalanceResult.Status status = RebalanceResult.Status.UNKNOWN_ERROR;
+ try {
+ RebalanceResult result = doRebalance(tableConfig, rebalanceConfig, rebalanceJobId);
+ status = result.getStatus();
+ return result;
+ } finally {
+ if (_controllerMetrics != null) {
+ _controllerMetrics.addTimedTableValue(String.format("%s.%s", tableNameWithType, status.toString()),
+ ControllerTimer.TABLE_REBALANCE_EXECUTION_TIME_MS, System.currentTimeMillis() - startTime,
+ TimeUnit.MILLISECONDS);
+ }
+ }
+ }
+
+ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig rebalanceConfig,
+ @Nullable String rebalanceJobId) {
long startTimeMs = System.currentTimeMillis();
String tableNameWithType = tableConfig.getTableName();
if (rebalanceJobId == null) {
@@ -355,6 +375,11 @@ public class TableRebalancer {
"For rebalanceId: %s, caught exception while waiting for ExternalView to converge for table: %s, "
+ "aborting the rebalance", rebalanceJobId, tableNameWithType);
LOGGER.warn(errorMsg, e);
+ if (_tableRebalanceObserver.isStopped()) {
+ return new RebalanceResult(rebalanceJobId, _tableRebalanceObserver.getStopStatus(),
+ "Caught exception while waiting for ExternalView to converge: " + e, instancePartitionsMap,
+ tierToInstancePartitionsMap, targetAssignment);
+ }
_tableRebalanceObserver.onError(errorMsg);
return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED,
"Caught exception while waiting for ExternalView to converge: " + e, instancePartitionsMap,
@@ -434,6 +459,11 @@ public class TableRebalancer {
// Record change of current ideal state and the new target
_tableRebalanceObserver.onTrigger(TableRebalanceObserver.Trigger.IDEAL_STATE_CHANGE_TRIGGER, currentAssignment,
targetAssignment);
+ if (_tableRebalanceObserver.isStopped()) {
+ return new RebalanceResult(rebalanceJobId, _tableRebalanceObserver.getStopStatus(),
+ "Rebalance has stopped already before updating the IdealState", instancePartitionsMap,
+ tierToInstancePartitionsMap, targetAssignment);
+ }
Map<String, Map<String, String>> nextAssignment =
getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup);
LOGGER.info("For rebalanceId: {}, got the next assignment for table: {} with number of segments to be moved to "
@@ -691,6 +721,11 @@ public class TableRebalancer {
_tableRebalanceObserver.onTrigger(
TableRebalanceObserver.Trigger.EXTERNAL_VIEW_TO_IDEAL_STATE_CONVERGENCE_TRIGGER,
externalView.getRecord().getMapFields(), idealState.getRecord().getMapFields());
+ if (_tableRebalanceObserver.isStopped()) {
+ throw new RuntimeException(
+ String.format("Rebalance for table: %s has already stopped with status: %s", tableNameWithType,
+ _tableRebalanceObserver.getStopStatus()));
+ }
if (isExternalViewConverged(tableNameWithType, externalView.getRecord().getMapFields(),
idealState.getRecord().getMapFields(), bestEfforts, segmentsToMonitor)) {
LOGGER.info("ExternalView converged for table: {}", tableNameWithType);
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/ZkBasedTableRebalanceObserver.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/ZkBasedTableRebalanceObserver.java
index 1978df35f3..7b57147ec0 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/ZkBasedTableRebalanceObserver.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/ZkBasedTableRebalanceObserver.java
@@ -19,10 +19,10 @@
package org.apache.pinot.controller.helix.core.rebalance;
import com.fasterxml.jackson.core.JsonProcessingException;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.util.HashMap;
import java.util.Map;
-import org.apache.pinot.common.metadata.ZKMetadataProvider;
import org.apache.pinot.common.metadata.controllerjob.ControllerJobType;
import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
import org.apache.pinot.spi.utils.CommonConstants;
@@ -30,6 +30,7 @@ import org.apache.pinot.spi.utils.JsonUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+
/**
* <code>ZkBasedTableRebalanceObserver</code> observes rebalance progress and tracks rebalance status,
* stats in Zookeeper. This will be used to show the progress of rebalance to users via rebalanceStatus API.
@@ -39,12 +40,16 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
private final String _tableNameWithType;
private final String _rebalanceJobId;
private final PinotHelixResourceManager _pinotHelixResourceManager;
- private TableRebalanceProgressStats _tableRebalanceProgressStats;
+ private final TableRebalanceProgressStats _tableRebalanceProgressStats;
+ private final TableRebalanceContext _tableRebalanceContext;
+ private long _lastUpdateTimeMs;
// Keep track of number of updates. Useful during debugging.
private int _numUpdatesToZk;
+ private boolean _isStopped = false;
+ private RebalanceResult.Status _stopStatus;
public ZkBasedTableRebalanceObserver(String tableNameWithType, String rebalanceJobId,
- PinotHelixResourceManager pinotHelixResourceManager) {
+ TableRebalanceContext tableRebalanceContext, PinotHelixResourceManager pinotHelixResourceManager) {
Preconditions.checkState(tableNameWithType != null, "Table name cannot be null");
Preconditions.checkState(rebalanceJobId != null, "rebalanceId cannot be null");
Preconditions.checkState(pinotHelixResourceManager != null, "PinotHelixManager cannot be null");
@@ -52,16 +57,19 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
_rebalanceJobId = rebalanceJobId;
_pinotHelixResourceManager = pinotHelixResourceManager;
_tableRebalanceProgressStats = new TableRebalanceProgressStats();
+ _tableRebalanceContext = tableRebalanceContext;
_numUpdatesToZk = 0;
}
@Override
public void onTrigger(Trigger trigger, Map<String, Map<String, String>> currentState,
Map<String, Map<String, String>> targetState) {
+ boolean updatedStatsInZk = false;
switch (trigger) {
case START_TRIGGER:
updateOnStart(currentState, targetState);
trackStatsInZk();
+ updatedStatsInZk = true;
break;
// Write to Zk if there's change since previous stats computation
case IDEAL_STATE_CHANGE_TRIGGER:
@@ -71,6 +79,7 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
latest)) {
_tableRebalanceProgressStats.setCurrentToTargetConvergence(latest);
trackStatsInZk();
+ updatedStatsInZk = true;
}
break;
case EXTERNAL_VIEW_TO_IDEAL_STATE_CONVERGENCE_TRIGGER:
@@ -79,18 +88,28 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
_tableRebalanceProgressStats.getExternalViewToIdealStateConvergence(), latest)) {
_tableRebalanceProgressStats.setExternalViewToIdealStateConvergence(latest);
trackStatsInZk();
+ updatedStatsInZk = true;
}
break;
default:
throw new IllegalArgumentException("Unimplemented trigger: " + trigger);
}
+ // The onTrigger method is mainly driven by the while loop of waiting for external view to converge to ideal
+ // state. That while loop wait for at least externalViewCheckIntervalInMs. So the real interval to send out
+ // heartbeat is the max(heartbeat_interval, externalViewCheckIntervalInMs);
+ long heartbeatIntervalInMs = _tableRebalanceContext.getConfig().getHeartbeatIntervalInMs();
+ if (!updatedStatsInZk && System.currentTimeMillis() - _lastUpdateTimeMs > heartbeatIntervalInMs) {
+ LOGGER.debug("Update status of rebalance job: {} for table: {} after {}ms as heartbeat", _rebalanceJobId,
+ _tableNameWithType, heartbeatIntervalInMs);
+ trackStatsInZk();
+ }
}
private void updateOnStart(Map<String, Map<String, String>> currentState,
Map<String, Map<String, String>> targetState) {
- Preconditions.checkState(_tableRebalanceProgressStats.getStatus() != RebalanceResult.Status.IN_PROGRESS.toString(),
+ Preconditions.checkState(RebalanceResult.Status.IN_PROGRESS != _tableRebalanceProgressStats.getStatus(),
"Rebalance Observer onStart called multiple times");
- _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.IN_PROGRESS.toString());
+ _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.IN_PROGRESS);
_tableRebalanceProgressStats.setInitialToTargetStateConvergence(
getDifferenceBetweenTableRebalanceStates(targetState, currentState));
_tableRebalanceProgressStats.setStartTimeMs(System.currentTimeMillis());
@@ -98,13 +117,12 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
@Override
public void onSuccess(String msg) {
- Preconditions.checkState(_tableRebalanceProgressStats.getStatus() != RebalanceResult.Status.DONE.toString(),
+ Preconditions.checkState(RebalanceResult.Status.DONE != _tableRebalanceProgressStats.getStatus(),
"Table Rebalance already completed");
- long timeToFinishInSeconds =
- (System.currentTimeMillis() - _tableRebalanceProgressStats.getStartTimeMs()) / 1000L;
+ long timeToFinishInSeconds = (System.currentTimeMillis() - _tableRebalanceProgressStats.getStartTimeMs()) / 1000L;
_tableRebalanceProgressStats.setCompletionStatusMsg(msg);
_tableRebalanceProgressStats.setTimeToFinishInSeconds(timeToFinishInSeconds);
- _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.DONE.toString());
+ _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.DONE);
// Zero out the in_progress convergence stats
TableRebalanceProgressStats.RebalanceStateStats stats = new TableRebalanceProgressStats.RebalanceStateStats();
_tableRebalanceProgressStats.setExternalViewToIdealStateConvergence(stats);
@@ -114,42 +132,97 @@ public class ZkBasedTableRebalanceObserver implements TableRebalanceObserver {
@Override
public void onError(String errorMsg) {
- long timeToFinishInSeconds =
- (long) (System.currentTimeMillis() - _tableRebalanceProgressStats.getStartTimeMs()) / 1000;
+ long timeToFinishInSeconds = (System.currentTimeMillis() - _tableRebalanceProgressStats.getStartTimeMs()) / 1000;
_tableRebalanceProgressStats.setTimeToFinishInSeconds(timeToFinishInSeconds);
- _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.FAILED.toString());
+ _tableRebalanceProgressStats.setStatus(RebalanceResult.Status.FAILED);
_tableRebalanceProgressStats.setCompletionStatusMsg(errorMsg);
trackStatsInZk();
}
+ @Override
+ public boolean isStopped() {
+ return _isStopped;
+ }
+
+ @Override
+ public RebalanceResult.Status getStopStatus() {
+ return _stopStatus;
+ }
+
public int getNumUpdatesToZk() {
return _numUpdatesToZk;
}
+ @VisibleForTesting
+ TableRebalanceContext getTableRebalanceContext() {
+ return _tableRebalanceContext;
+ }
+
private void trackStatsInZk() {
+ Map<String, String> jobMetadata =
+ createJobMetadata(_tableNameWithType, _rebalanceJobId, _tableRebalanceProgressStats, _tableRebalanceContext);
+ _pinotHelixResourceManager.addControllerJobToZK(_rebalanceJobId, jobMetadata, ControllerJobType.TABLE_REBALANCE,
+ prevJobMetadata -> {
+ // In addition to updating job progress status, the observer also checks if the job status is IN_PROGRESS.
+ // If not, then no need to update the job status, and we keep this status to end the job promptly.
+ if (prevJobMetadata == null) {
+ return true;
+ }
+ String prevStatsInStr = prevJobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS);
+ TableRebalanceProgressStats prevStats;
+ try {
+ prevStats = JsonUtils.stringToObject(prevStatsInStr, TableRebalanceProgressStats.class);
+ } catch (JsonProcessingException ignore) {
+ return true;
+ }
+ if (prevStats == null || RebalanceResult.Status.IN_PROGRESS == prevStats.getStatus()) {
+ return true;
+ }
+ _isStopped = true;
+ _stopStatus = prevStats.getStatus();
+ LOGGER.warn("Rebalance job: {} for table: {} has already stopped with status: {}", _rebalanceJobId,
+ _tableNameWithType, _stopStatus);
+ // No need to update job status if job has ended. This also keeps the last status from being overwritten.
+ return false;
+ });
+ _numUpdatesToZk++;
+ _lastUpdateTimeMs = System.currentTimeMillis();
+ LOGGER.debug("Made {} ZK updates for rebalance job: {} of table: {}", _numUpdatesToZk, _rebalanceJobId,
+ _tableNameWithType);
+ }
+
+ @VisibleForTesting
+ static Map<String, String> createJobMetadata(String tableNameWithType, String jobId,
+ TableRebalanceProgressStats tableRebalanceProgressStats, TableRebalanceContext tableRebalanceContext) {
Map<String, String> jobMetadata = new HashMap<>();
- jobMetadata.put(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE, _tableNameWithType);
- jobMetadata.put(CommonConstants.ControllerJob.JOB_ID, _rebalanceJobId);
+ jobMetadata.put(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE, tableNameWithType);
+ jobMetadata.put(CommonConstants.ControllerJob.JOB_ID, jobId);
jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, Long.toString(System.currentTimeMillis()));
jobMetadata.put(CommonConstants.ControllerJob.JOB_TYPE, ControllerJobType.TABLE_REBALANCE.name());
try {
jobMetadata.put(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS,
- JsonUtils.objectToString(_tableRebalanceProgressStats));
+ JsonUtils.objectToString(tableRebalanceProgressStats));
} catch (JsonProcessingException e) {
- LOGGER.error("Error serialising rebalance stats to JSON for persisting to ZK {}", _rebalanceJobId, e);
+ LOGGER.error("Error serialising stats for rebalance job: {} of table: {} to keep in ZK", jobId, tableNameWithType,
+ e);
}
- _pinotHelixResourceManager.addControllerJobToZK(_rebalanceJobId, jobMetadata,
- ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.TABLE_REBALANCE));
- _numUpdatesToZk++;
- LOGGER.debug("Number of updates to Zk: {} for rebalanceJob: {} ", _numUpdatesToZk, _rebalanceJobId);
+ try {
+ jobMetadata.put(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_CONTEXT,
+ JsonUtils.objectToString(tableRebalanceContext));
+ } catch (JsonProcessingException e) {
+ LOGGER.error("Error serialising retry configs for rebalance job: {} of table: {} to keep in ZK", jobId,
+ tableNameWithType, e);
+ }
+ return jobMetadata;
}
/**
- * Takes in targetState and sourceState and computes stats based on the comparison
- * between sourceState and targetState.This captures how far the source state is from
- * the target state. Example - if there are 4 segments and 16 replicas in the source state
- * not matching the target state, _segmentsToRebalance is 4 and _replicasToRebalance is 16.
- * @param targetState- The state that we want to get to
+ * Takes in targetState and sourceState and computes stats based on the comparison between sourceState and
+ * targetState.This captures how far the source state is from the target state. Example - if there are 4 segments and
+ * 16 replicas in the source state not matching the target state, _segmentsToRebalance is 4 and _replicasToRebalance
+ * is 16.
+ *
+ * @param targetState - The state that we want to get to
* @param sourceState - A given state that needs to converge to targetState
* @return RebalanceStats
*/
diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/ZkBasedTenantRebalanceObserver.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/ZkBasedTenantRebalanceObserver.java
index 6c26985539..9fabc64d9d 100644
--- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/ZkBasedTenantRebalanceObserver.java
+++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/ZkBasedTenantRebalanceObserver.java
@@ -25,7 +25,6 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
-import org.apache.pinot.common.metadata.ZKMetadataProvider;
import org.apache.pinot.common.metadata.controllerjob.ControllerJobType;
import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
import org.apache.pinot.controller.helix.core.rebalance.RebalanceJobConstants;
@@ -108,8 +107,7 @@ public class ZkBasedTenantRebalanceObserver implements TenantRebalanceObserver {
} catch (JsonProcessingException e) {
LOGGER.error("Error serialising rebalance stats to JSON for persisting to ZK {}", _jobId, e);
}
- _pinotHelixResourceManager.addControllerJobToZK(_jobId, jobMetadata,
- ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.TENANT_REBALANCE));
+ _pinotHelixResourceManager.addControllerJobToZK(_jobId, jobMetadata, ControllerJobType.TENANT_REBALANCE);
_numUpdatesToZk++;
LOGGER.debug("Number of updates to Zk: {} for rebalanceJob: {} ", _numUpdatesToZk, _jobId);
}
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java
index 6172c67def..a1b086c8fa 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerPeriodicTaskStarterStatelessTest.java
@@ -57,7 +57,7 @@ public class ControllerPeriodicTaskStarterStatelessTest extends ControllerTest {
}
private class MockControllerStarter extends ControllerStarter {
- private static final int NUM_PERIODIC_TASKS = 10;
+ private static final int NUM_PERIODIC_TASKS = 11;
public MockControllerStarter() {
super();
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceCheckerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceCheckerTest.java
new file mode 100644
index 0000000000..22663fa333
--- /dev/null
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceCheckerTest.java
@@ -0,0 +1,388 @@
+/**
+ * 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.core.JsonProcessingException;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.util.concurrent.MoreExecutors;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.helix.AccessOption;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixDataAccessor;
+import org.apache.helix.HelixManager;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.metadata.ZKMetadataProvider;
+import org.apache.pinot.common.metadata.controllerjob.ControllerJobType;
+import org.apache.pinot.common.metrics.ControllerMetrics;
+import org.apache.pinot.controller.ControllerConf;
+import org.apache.pinot.controller.LeadControllerManager;
+import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+public class RebalanceCheckerTest {
+
+ @Test
+ public void testGetRetryDelayInMs() {
+ assertEquals(RebalanceChecker.getRetryDelayInMs(0, 1), 0);
+ assertEquals(RebalanceChecker.getRetryDelayInMs(0, 2), 0);
+ assertEquals(RebalanceChecker.getRetryDelayInMs(0, 3), 0);
+
+ for (long initDelayMs : new long[]{1, 30000, 3600000}) {
+ long delayMs = RebalanceChecker.getRetryDelayInMs(initDelayMs, 1);
+ assertTrue(delayMs >= initDelayMs && delayMs < initDelayMs * 2);
+ delayMs = RebalanceChecker.getRetryDelayInMs(initDelayMs, 2);
+ assertTrue(delayMs >= initDelayMs * 2 && delayMs < initDelayMs * 4);
+ delayMs = RebalanceChecker.getRetryDelayInMs(initDelayMs, 3);
+ assertTrue(delayMs >= initDelayMs * 4 && delayMs < initDelayMs * 8);
+ }
+ }
+
+ @Test
+ public void testGetCandidateJobs()
+ throws Exception {
+ String tableName = "table01";
+ Map<String, Map<String, String>> allJobMetadata = new HashMap<>();
+
+ // Original job run as job1, and all its retry jobs failed too.
+ RebalanceConfig jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ TableRebalanceProgressStats stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(1000);
+ TableRebalanceContext jobCtx = TableRebalanceContext.forInitialAttempt("job1", jobCfg);
+ Map<String, String> jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job1", stats, jobCtx);
+ allJobMetadata.put("job1", jobMetadata);
+ // 3 failed retry runs for job1
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 2, 1100, RebalanceResult.Status.FAILED);
+ allJobMetadata.put("job1_2", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 3, 1200, RebalanceResult.Status.ABORTED);
+ allJobMetadata.put("job1_3", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 4, 1300, RebalanceResult.Status.FAILED);
+ allJobMetadata.put("job1_4", jobMetadata);
+
+ // Original job run as job2, and its retry job job2_1 completed.
+ jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(2000);
+ jobCtx = TableRebalanceContext.forInitialAttempt("job2", jobCfg);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job2", stats, jobCtx);
+ allJobMetadata.put("job2", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job2", 2, 2100, RebalanceResult.Status.DONE);
+ allJobMetadata.put("job2_2", jobMetadata);
+
+ // Original job run as job3, and failed to send out heartbeat in time.
+ jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.IN_PROGRESS);
+ stats.setStartTimeMs(3000);
+ jobCtx = TableRebalanceContext.forInitialAttempt("job3", jobCfg);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job3", stats, jobCtx);
+ jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, "3000");
+ allJobMetadata.put("job3", jobMetadata);
+
+ // Original job run as job4, which didn't have retryJobCfg as from old version of the code.
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(4000);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job4", stats, null);
+ jobMetadata.remove(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_CONTEXT);
+ allJobMetadata.put("job4", jobMetadata);
+
+ // Only need to retry job1 and job3, as job2 is completed and job4 is from old version of code.
+ Map<String, Set<Pair<TableRebalanceContext, Long>>> jobs =
+ RebalanceChecker.getCandidateJobs(tableName, allJobMetadata);
+ assertEquals(jobs.size(), 2);
+ assertTrue(jobs.containsKey("job1"));
+ assertTrue(jobs.containsKey("job3"));
+ assertEquals(jobs.get("job1").size(), 4); // four runs including job1,job1_1,job1_2,job1_3
+ assertEquals(jobs.get("job3").size(), 1); // just a single run job3
+
+ // Abort job1 and cancel its retries, then only job3 is retry candidate.
+ jobMetadata = allJobMetadata.get("job1_4");
+ cancelRebalanceJob(jobMetadata);
+ jobs = RebalanceChecker.getCandidateJobs(tableName, allJobMetadata);
+ assertEquals(jobs.size(), 1);
+ assertTrue(jobs.containsKey("job3"));
+ assertEquals(jobs.get("job3").size(), 1); // just a single run job3
+
+ // Add latest job5 that's already done, thus no need to retry for table.
+ jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.DONE);
+ stats.setStartTimeMs(5000);
+ jobCtx = TableRebalanceContext.forInitialAttempt("job5", jobCfg);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job5", stats, jobCtx);
+ allJobMetadata.put("job5", jobMetadata);
+ jobs = RebalanceChecker.getCandidateJobs(tableName, allJobMetadata);
+ assertEquals(jobs.size(), 0);
+ }
+
+ @Test
+ public void testGetLatestJob() {
+ Map<String, Set<Pair<TableRebalanceContext, Long>>> jobs = new HashMap<>();
+ // The most recent job run is job1_3, and within 3 maxAttempts.
+ jobs.put("job1",
+ ImmutableSet.of(Pair.of(createDummyJobCtx("job1", 1), 10L), Pair.of(createDummyJobCtx("job1", 2), 20L),
+ Pair.of(createDummyJobCtx("job1", 3), 1020L)));
+ jobs.put("job2", ImmutableSet.of(Pair.of(createDummyJobCtx("job2", 1), 1000L)));
+ Pair<TableRebalanceContext, Long> jobTime = RebalanceChecker.getLatestJob(jobs);
+ assertNotNull(jobTime);
+ assertEquals(jobTime.getLeft().getJobId(), "job1_3");
+
+ // The most recent job run is job1_4, but reached 3 maxAttempts.
+ jobs.put("job1",
+ ImmutableSet.of(Pair.of(createDummyJobCtx("job1", 1), 10L), Pair.of(createDummyJobCtx("job1", 2), 20L),
+ Pair.of(createDummyJobCtx("job1", 3), 1020L), Pair.of(createDummyJobCtx("job1", 4), 2020L)));
+ jobTime = RebalanceChecker.getLatestJob(jobs);
+ assertNotNull(jobTime);
+ assertEquals(jobTime.getLeft().getJobId(), "job2");
+
+ // Add job3 that's started more recently.
+ jobs.put("job3", ImmutableSet.of(Pair.of(createDummyJobCtx("job3", 1), 3000L)));
+ jobTime = RebalanceChecker.getLatestJob(jobs);
+ assertNotNull(jobTime);
+ assertEquals(jobTime.getLeft().getJobId(), "job3");
+
+ // Remove job2 and job3, and we'd have no job to retry then.
+ jobs.remove("job2");
+ jobs.remove("job3");
+ jobTime = RebalanceChecker.getLatestJob(jobs);
+ assertNull(jobTime);
+ }
+
+ @Test
+ public void testRetryRebalance()
+ throws Exception {
+ String tableName = "table01";
+ LeadControllerManager leadController = mock(LeadControllerManager.class);
+ ControllerMetrics metrics = mock(ControllerMetrics.class);
+ ExecutorService exec = MoreExecutors.newDirectExecutorService();
+ ControllerConf cfg = new ControllerConf();
+
+ Map<String, Map<String, String>> allJobMetadata = new HashMap<>();
+ // Original job run as job1, and all its retry jobs failed too.
+ RebalanceConfig jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ TableRebalanceProgressStats stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(1000);
+ TableRebalanceContext jobCtx = TableRebalanceContext.forInitialAttempt("job1", jobCfg);
+ Map<String, String> jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job1", stats, jobCtx);
+ allJobMetadata.put("job1", jobMetadata);
+ // 3 failed retry runs for job1
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 2, 1100, RebalanceResult.Status.FAILED);
+ allJobMetadata.put("job1_2", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 3, 1200, RebalanceResult.Status.FAILED);
+ allJobMetadata.put("job1_3", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job1", 4, 5300, RebalanceResult.Status.FAILED);
+ allJobMetadata.put("job1_4", jobMetadata);
+
+ // Original job run as job2, and its retry job job2_1 completed.
+ jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(2000);
+ jobCtx = TableRebalanceContext.forInitialAttempt("job2", jobCfg);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job2", stats, jobCtx);
+ allJobMetadata.put("job2", jobMetadata);
+ jobMetadata = createDummyJobMetadata(tableName, "job2", 2, 2100, RebalanceResult.Status.DONE);
+ allJobMetadata.put("job2_2", jobMetadata);
+
+ // Original job run as job3, and failed to send out heartbeat in time.
+ jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.IN_PROGRESS);
+ stats.setStartTimeMs(3000);
+ jobCtx = TableRebalanceContext.forInitialAttempt("job3", jobCfg);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job3", stats, jobCtx);
+ jobMetadata.put(CommonConstants.ControllerJob.SUBMISSION_TIME_MS, "3000");
+ allJobMetadata.put("job3", jobMetadata);
+
+ TableConfig tableConfig = mock(TableConfig.class);
+ PinotHelixResourceManager helixManager = mock(PinotHelixResourceManager.class);
+ when(helixManager.getTableConfig(tableName)).thenReturn(tableConfig);
+ when(helixManager.getAllJobs(any(), any())).thenReturn(allJobMetadata);
+ RebalanceChecker checker = new RebalanceChecker(helixManager, leadController, cfg, metrics, exec);
+ // Although job1_3 was submitted most recently but job1 had exceeded maxAttempts. Chose job3 to retry, which got
+ // stuck at in progress status.
+ checker.retryRebalanceTable(tableName, allJobMetadata);
+ // The new retry job is for job3 and attemptId is increased to 2.
+ ArgumentCaptor<ZkBasedTableRebalanceObserver> observerCaptor =
+ ArgumentCaptor.forClass(ZkBasedTableRebalanceObserver.class);
+ verify(helixManager, times(1)).rebalanceTable(eq(tableName), any(), anyString(), any(), observerCaptor.capture());
+ ZkBasedTableRebalanceObserver observer = observerCaptor.getValue();
+ jobCtx = observer.getTableRebalanceContext();
+ assertEquals(jobCtx.getOriginalJobId(), "job3");
+ assertEquals(jobCtx.getAttemptId(), 2);
+ }
+
+ @Test
+ public void testRetryRebalanceWithBackoff()
+ throws Exception {
+ String tableName = "table01";
+ LeadControllerManager leadController = mock(LeadControllerManager.class);
+ ControllerMetrics metrics = mock(ControllerMetrics.class);
+ ExecutorService exec = MoreExecutors.newDirectExecutorService();
+ ControllerConf cfg = new ControllerConf();
+
+ Map<String, Map<String, String>> allJobMetadata = new HashMap<>();
+ // Original job run as job1, and all its retry jobs failed too.
+ RebalanceConfig jobCfg = new RebalanceConfig();
+ jobCfg.setMaxAttempts(4);
+ long nowMs = System.currentTimeMillis();
+ TableRebalanceProgressStats stats = new TableRebalanceProgressStats();
+ stats.setStatus(RebalanceResult.Status.FAILED);
+ stats.setStartTimeMs(nowMs);
+ TableRebalanceContext jobCtx = TableRebalanceContext.forInitialAttempt("job1", jobCfg);
+ Map<String, String> jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job1", stats, jobCtx);
+ allJobMetadata.put("job1", jobMetadata);
+
+ PinotHelixResourceManager helixManager = mock(PinotHelixResourceManager.class);
+ TableConfig tableConfig = mock(TableConfig.class);
+ when(helixManager.getTableConfig(tableName)).thenReturn(tableConfig);
+ RebalanceChecker checker = new RebalanceChecker(helixManager, leadController, cfg, metrics, exec);
+ checker.retryRebalanceTable(tableName, allJobMetadata);
+ // Retry for job1 is delayed with 5min backoff.
+ ArgumentCaptor<ZkBasedTableRebalanceObserver> observerCaptor =
+ ArgumentCaptor.forClass(ZkBasedTableRebalanceObserver.class);
+ verify(helixManager, times(0)).rebalanceTable(eq(tableName), any(), anyString(), any(), observerCaptor.capture());
+
+ // Set initial delay to 0 to disable retry backoff.
+ jobCfg.setRetryInitialDelayInMs(0);
+ jobMetadata = ZkBasedTableRebalanceObserver.createJobMetadata(tableName, "job1", stats, jobCtx);
+ allJobMetadata.put("job1", jobMetadata);
+ checker.retryRebalanceTable(tableName, allJobMetadata);
+ // Retry for job1 is delayed with 0 backoff.
+ observerCaptor = ArgumentCaptor.forClass(ZkBasedTableRebalanceObserver.class);
+ verify(helixManager, times(1)).rebalanceTable(eq(tableName), any(), anyString(), any(), observerCaptor.capture());
+ }
+
+ @Test
+ public void testAddUpdateControllerJobsForTable() {
+ ControllerConf cfg = new ControllerConf();
+ cfg.setZkStr("localhost:2181");
+ cfg.setHelixClusterName("cluster01");
+ PinotHelixResourceManager pinotHelixManager = new PinotHelixResourceManager(cfg);
+ HelixManager helixZkManager = mock(HelixManager.class);
+ ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class);
+ String zkPath = ZKMetadataProvider.constructPropertyStorePathForControllerJob(ControllerJobType.TABLE_REBALANCE);
+ ZNRecord jobsZnRecord = new ZNRecord("jobs");
+ when(propertyStore.get(eq(zkPath), any(), eq(AccessOption.PERSISTENT))).thenReturn(jobsZnRecord);
+ when(helixZkManager.getClusterManagmentTool()).thenReturn(mock(HelixAdmin.class));
+ when(helixZkManager.getHelixPropertyStore()).thenReturn(propertyStore);
+ when(helixZkManager.getHelixDataAccessor()).thenReturn(mock(HelixDataAccessor.class));
+ pinotHelixManager.start(helixZkManager, null);
+
+ pinotHelixManager.addControllerJobToZK("job1",
+ ImmutableMap.of("jobId", "job1", "submissionTimeMs", "1000", "tableName", "table01"),
+ ControllerJobType.TABLE_REBALANCE, jmd -> true);
+ pinotHelixManager.addControllerJobToZK("job2",
+ ImmutableMap.of("jobId", "job2", "submissionTimeMs", "2000", "tableName", "table01"),
+ ControllerJobType.TABLE_REBALANCE, jmd -> false);
+ pinotHelixManager.addControllerJobToZK("job3",
+ ImmutableMap.of("jobId", "job3", "submissionTimeMs", "3000", "tableName", "table02"),
+ ControllerJobType.TABLE_REBALANCE, jmd -> true);
+ pinotHelixManager.addControllerJobToZK("job4",
+ ImmutableMap.of("jobId", "job4", "submissionTimeMs", "4000", "tableName", "table02"),
+ ControllerJobType.TABLE_REBALANCE, jmd -> true);
+ Map<String, Map<String, String>> jmds = jobsZnRecord.getMapFields();
+ assertEquals(jmds.size(), 3);
+ assertTrue(jmds.containsKey("job1"));
+ assertTrue(jmds.containsKey("job3"));
+ assertTrue(jmds.containsKey("job4"));
+
+ Set<String> expectedJobs01 = new HashSet<>();
+ pinotHelixManager.updateJobsForTable("table01", ControllerJobType.TABLE_REBALANCE,
+ jmd -> expectedJobs01.add(jmd.get("jobId")));
+ assertEquals(expectedJobs01.size(), 1);
+ assertTrue(expectedJobs01.contains("job1"));
+
+ Set<String> expectedJobs02 = new HashSet<>();
+ pinotHelixManager.updateJobsForTable("table02", ControllerJobType.TABLE_REBALANCE,
+ jmd -> expectedJobs02.add(jmd.get("jobId")));
+ assertEquals(expectedJobs02.size(), 2);
+ assertTrue(expectedJobs02.contains("job3"));
+ assertTrue(expectedJobs02.contains("job4"));
+ }
+
+ private static TableRebalanceContext createDummyJobCtx(String originalJobId, int attemptId) {
+ TableRebalanceContext jobCtx = new TableRebalanceContext();
+ RebalanceConfig cfg = new RebalanceConfig();
+ cfg.setMaxAttempts(4);
+ jobCtx.setJobId(originalJobId);
+ if (attemptId > 1) {
+ jobCtx.setJobId(originalJobId + "_" + attemptId);
+ }
+ jobCtx.setOriginalJobId(originalJobId);
+ jobCtx.setConfig(cfg);
+ jobCtx.setAttemptId(attemptId);
+ return jobCtx;
+ }
+
+ private static Map<String, String> createDummyJobMetadata(String tableName, String originalJobId, int attemptId,
+ long startTimeMs, RebalanceResult.Status status) {
+ RebalanceConfig cfg = new RebalanceConfig();
+ cfg.setMaxAttempts(4);
+ TableRebalanceProgressStats stats = new TableRebalanceProgressStats();
+ stats.setStatus(status);
+ stats.setStartTimeMs(startTimeMs);
+ TableRebalanceContext jobCtx = TableRebalanceContext.forRetry(originalJobId, cfg, attemptId);
+ String attemptJobId = originalJobId + "_" + attemptId;
+ return ZkBasedTableRebalanceObserver.createJobMetadata(tableName, attemptJobId, stats, jobCtx);
+ }
+
+ private static void cancelRebalanceJob(Map<String, String> jobMetadata)
+ throws JsonProcessingException {
+ String jobStatsInStr = jobMetadata.get(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS);
+ TableRebalanceProgressStats jobStats = JsonUtils.stringToObject(jobStatsInStr, TableRebalanceProgressStats.class);
+ jobStats.setStatus(RebalanceResult.Status.CANCELLED);
+ jobMetadata.put(RebalanceJobConstants.JOB_METADATA_KEY_REBALANCE_PROGRESS_STATS,
+ JsonUtils.objectToString(jobStats));
+ }
+}
diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TestZkBasedTableRebalanceObserver.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TestZkBasedTableRebalanceObserver.java
index 74b8a4367f..143caebf2b 100644
--- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TestZkBasedTableRebalanceObserver.java
+++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TestZkBasedTableRebalanceObserver.java
@@ -35,14 +35,16 @@ import static org.testng.Assert.assertEquals;
public class TestZkBasedTableRebalanceObserver {
+ // This is a test to verify if Zk stats are pushed out correctly
@Test
- // This is a test to verify if Zk stats are pushed out correctly
void testZkObserverTracking() {
PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class);
// Mocking this. We will verify using numZkUpdate stat
when(pinotHelixResourceManager.addControllerJobToZK(any(), any(), any())).thenReturn(true);
+ TableRebalanceContext retryCtx = new TableRebalanceContext();
+ retryCtx.setConfig(new RebalanceConfig());
ZkBasedTableRebalanceObserver observer =
- new ZkBasedTableRebalanceObserver("dummy", "dummyId", pinotHelixResourceManager);
+ new ZkBasedTableRebalanceObserver("dummy", "dummyId", retryCtx, pinotHelixResourceManager);
Map<String, Map<String, String>> source = new TreeMap<>();
Map<String, Map<String, String>> target = new TreeMap<>();
target.put("segment1",
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/auth/Actions.java b/pinot-core/src/main/java/org/apache/pinot/core/auth/Actions.java
index e72d066bc9..db6941bf1f 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/auth/Actions.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/auth/Actions.java
@@ -90,6 +90,7 @@ public class Actions {
// Action names for table
public static class Table {
public static final String BUILD_ROUTING = "BuildRouting";
+ public static final String CANCEL_REBALANCE = "CancelRebalance";
public static final String CREATE_INSTANCE_PARTITIONS = "CreateInstancePartitions";
public static final String CREATE_SCHEMA = "CreateSchema";
public static final String CREATE_TABLE = "CreateTable";
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
index 0f88a36d06..4a543556f7 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java
@@ -370,9 +370,8 @@ public class PartialUpsertTableRebalanceIntegrationTest extends BaseClusterInteg
ServerRebalanceJobStatusResponse serverRebalanceJobStatusResponse =
JsonUtils.stringToObject(httpResponse.getResponse(), ServerRebalanceJobStatusResponse.class);
- String status = serverRebalanceJobStatusResponse.getTableRebalanceProgressStats().getStatus();
- return status.equals(RebalanceResult.Status.DONE.toString()) || status.equals(
- RebalanceResult.Status.FAILED.toString());
+ RebalanceResult.Status status = serverRebalanceJobStatusResponse.getTableRebalanceProgressStats().getStatus();
+ return status != RebalanceResult.Status.IN_PROGRESS;
} catch (HttpErrorStatusException | URISyntaxException e) {
throw new IOException(e);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org