@VisibleForTesting
  public ProcessContinuation run(
      RestrictionTracker<OffsetRange, Long> tracker,
      OutputReceiver<PartitionRecord> receiver,
      ManualWatermarkEstimator<Instant> watermarkEstimator,
      InitialPipelineState initialPipelineState)
      throws Exception {
    LOG.debug("DNP: Watermark: " + watermarkEstimator.getState());
    LOG.debug("DNP: CurrentTracker: " + tracker.currentRestriction().getFrom());
    if (tracker.currentRestriction().getFrom() == 0L) {
      if (!tracker.tryClaim(0L)) {
        LOG.error(
            "Could not claim initial DetectNewPartition restriction. No partitions are outputted.");
        return ProcessContinuation.stop();
      }
      watermarkEstimator.setWatermark(initialPipelineState.getStartTime());
      if (initialPipelineState.isResume()) {
        resumeFromPreviousPipelineAction.run(receiver);
      } else {
        generateInitialPartitionsAction.run(receiver, initialPipelineState.getStartTime());
      }
      return ProcessContinuation.resume();
    }

    // Create a new partition reconciler every run to reset the state each time.
    partitionReconciler = new PartitionReconciler(metadataTableDao, metrics);
    orphanedMetadataCleaner = new OrphanedMetadataCleaner();

    // Calculating the new value of watermark is a resource intensive process. We have to do a full
    // scan of the metadata table and then ensure we're not missing partitions and then calculate
    // the low watermark. This is usually a fairly fast process even with thousands of partitions.
    // However, sometimes this may take so long that the runner checkpoints before the watermark is
    // calculated. Because the checkpoint takes place before tryClaim, this forces the DoFn to
    // restart, wasting the resources spent calculating the watermark. On restart, we will try to
    // calculate the watermark again. The problem causing the slow watermark calculation can persist
    // leading to a crash loop. In order to ensure we persist the calculated watermark, we calculate
    // the watermark after successful tryClaim. Then we write to the metadata table the new
    // watermark. On the start of each run we read the watermark and update the DoFn's watermark.
    DetectNewPartitionsState detectNewPartitionsState =
        metadataTableDao.readDetectNewPartitionsState();
    if (detectNewPartitionsState != null) {
      watermarkEstimator.setWatermark(detectNewPartitionsState.getWatermark());
    }

    // Terminate if endTime <= watermark that means all partitions have read up to or beyond
    // watermark. We no longer need to manage splits and merges, we can terminate.
    if (endTime != null && !watermarkEstimator.currentWatermark().isBefore(endTime)) {
      tracker.tryClaim(tracker.currentRestriction().getTo());
      return ProcessContinuation.stop();
    }

    if (!tracker.tryClaim(tracker.currentRestriction().getFrom())) {
      LOG.warn("DNP: Checkpointing, stopping this run: " + tracker.currentRestriction());
      return ProcessContinuation.stop();
    }

    // Read StreamPartitions to calculate watermark.
    List<StreamPartitionWithWatermark> streamPartitionsWithWatermark = null;
    if (shouldUpdateWatermark(tracker.currentRestriction().getFrom(), detectNewPartitionsState)) {
      streamPartitionsWithWatermark = metadataTableDao.readStreamPartitionsWithWatermark();
    }

    // Process NewPartitions and track the ones successfully outputted.
    List<NewPartition> newPartitions = metadataTableDao.readNewPartitions();
    List<ByteStringRange> outputtedNewPartitions = new ArrayList<>();
    for (NewPartition newPartition : newPartitions) {
      if (processNewPartitionsAction.processNewPartition(newPartition, receiver)) {
        outputtedNewPartitions.add(newPartition.getPartition());
      } else if (streamPartitionsWithWatermark != null) {
        // streamPartitionsWithWatermark is not null on runs that we update watermark. We only run
        // reconciliation when we update watermark. Only add incompleteNewPartitions if
        // reconciliation is being run
        partitionReconciler.addIncompleteNewPartitions(newPartition);
        orphanedMetadataCleaner.addIncompleteNewPartitions(newPartition);
      }
    }

    // Process the watermark using read StreamPartitions and NewPartitions.
    if (streamPartitionsWithWatermark != null) {
      Optional<Instant> maybeWatermark =
          getNewWatermark(streamPartitionsWithWatermark, newPartitions);
      maybeWatermark.ifPresent(metadataTableDao::updateDetectNewPartitionWatermark);
      // Only start reconciling after the pipeline has been running for a while.
      if (tracker.currentRestriction().getFrom() > 50) {
        // Using NewPartitions and StreamPartitions, evaluate partitions that are possibly not being
        // streamed. This isn't perfect because there may be partitions moving between
        // StreamPartitions and NewPartitions while scanning the metadata table. Also, this does not
        // include NewPartitions marked as deleted from a previous DNP run not yet processed by
        // RCSP.
        List<ByteStringRange> existingPartitions =
            streamPartitionsWithWatermark.stream()
                .map(StreamPartitionWithWatermark::getPartition)
                .collect(Collectors.toList());
        existingPartitions.addAll(outputtedNewPartitions);
        List<ByteStringRange> missingStreamPartitions =
            getMissingPartitionsFromEntireKeySpace(existingPartitions);
        orphanedMetadataCleaner.addMissingPartitions(missingStreamPartitions);
        partitionReconciler.addMissingPartitions(missingStreamPartitions);
        processReconcilerPartitions(
            receiver, watermarkEstimator, initialPipelineState.getStartTime());
        cleanUpOrphanedMetadata();
      }
    }

    return ProcessContinuation.resume().withResumeDelay(Duration.millis(100));
  }
@Test
  public void testMissingPartitionReconciled() throws Exception {
    // We only start reconciling after 50.
    // We advance watermark on every 2 restriction tracker advancement
    OffsetRange offsetRange = new OffsetRange(52, Long.MAX_VALUE);
    when(tracker.currentRestriction()).thenReturn(offsetRange);
    when(tracker.tryClaim(offsetRange.getFrom())).thenReturn(true);

    // Write 2 partitions to the table, missing [a, b)
    ByteStringRange partitionEmptyA = ByteStringRange.create("", "a");
    Instant watermarkEmptyA = endTime.plus(Duration.millis(100));
    PartitionRecord partitionRecordEmptyA =
        new PartitionRecord(
            partitionEmptyA,
            watermarkEmptyA,
            UniqueIdGenerator.getNextId(),
            watermarkEmptyA,
            Collections.emptyList(),
            null);
    metadataTableDao.lockAndRecordPartition(partitionRecordEmptyA);
    ByteStringRange partitionBEmpty = ByteStringRange.create("b", "");
    Instant watermarkBEmpty = endTime.plus(Duration.millis(1));
    PartitionRecord partitionRecordBEmpty =
        new PartitionRecord(
            partitionBEmpty,
            watermarkBEmpty,
            UniqueIdGenerator.getNextId(),
            watermarkBEmpty,
            Collections.emptyList(),
            null);
    metadataTableDao.lockAndRecordPartition(partitionRecordBEmpty);

    HashMap<ByteStringRange, Instant> missingPartitionDurations = new HashMap<>();
    ByteStringRange partitionAB = ByteStringRange.create("a", "b");
    // Partition missing for 10 minutes less 1 second.
    missingPartitionDurations.put(
        partitionAB, Instant.now().minus(Duration.standardSeconds(20 * 60 - 1)));
    metadataTableDao.writeDetectNewPartitionMissingPartitions(missingPartitionDurations);

    // Since there's no NewPartition corresponding to the missing partition, we can't reconcile with
    // continuation tokens. In order to reconcile without continuation tokens, the partition needs
    // to have been missing for more than 10 minutes.
    assertEquals(
        DoFn.ProcessContinuation.resume().withResumeDelay(Duration.millis(100)),
        action.run(
            tracker, receiver, watermarkEstimator, new InitialPipelineState(startTime, false)));
    verify(receiver, never()).outputWithTimestamp(any(), any());
    assertEquals(1, metadataTableDao.readDetectNewPartitionMissingPartitions().size());

    // Sleep for more than 1 second, enough that the missing partition needs to be reconciled.
    Thread.sleep(1001);

    // We advance the restriction tracker by 1. Because it is not a multiple of 2, we don't
    // evaluate missing partitions, which means we don't perform reconciliation.
    offsetRange = new OffsetRange(53, Long.MAX_VALUE);
    when(tracker.currentRestriction()).thenReturn(offsetRange);
    when(tracker.tryClaim(offsetRange.getFrom())).thenReturn(true);

    assertEquals(
        DoFn.ProcessContinuation.resume().withResumeDelay(Duration.millis(100)),
        action.run(
            tracker, receiver, watermarkEstimator, new InitialPipelineState(startTime, false)));
    verify(receiver, never()).outputWithTimestamp(any(), any());
    assertEquals(1, metadataTableDao.readDetectNewPartitionMissingPartitions().size());

    // Multiple of 2, reconciliation should happen.
    offsetRange = new OffsetRange(54, Long.MAX_VALUE);
    when(tracker.currentRestriction()).thenReturn(offsetRange);
    when(tracker.tryClaim(offsetRange.getFrom())).thenReturn(true);

    assertEquals(
        DoFn.ProcessContinuation.resume().withResumeDelay(Duration.millis(100)),
        action.run(
            tracker, receiver, watermarkEstimator, new InitialPipelineState(startTime, false)));
    assertEquals(0, metadataTableDao.readDetectNewPartitionMissingPartitions().size());
    verify(receiver, times(1))
        .outputWithTimestamp(partitionRecordArgumentCaptor.capture(), eq(Instant.EPOCH));
    assertEquals(partitionAB, partitionRecordArgumentCaptor.getValue().getPartition());
    assertEquals(
        watermarkEstimator.currentWatermark(),
        partitionRecordArgumentCaptor.getValue().getParentLowWatermark());
    assertEquals(endTime, partitionRecordArgumentCaptor.getValue().getEndTime());
    assertNotNull(partitionRecordArgumentCaptor.getValue().getStartTime());
    // The startTime should be the startTime of the pipeline because it's less than 1 hour before
    // the low watermark.
    assertEquals(startTime, partitionRecordArgumentCaptor.getValue().getStartTime());
  }