@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;
  }
@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);
  }