@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: %s", 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);
      }
    });
  }
@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());
  }