public ProcessContinuation run(
      PartitionRecord partitionRecord,
      RestrictionTracker<StreamProgress, StreamProgress> tracker,
      OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver,
      ManualWatermarkEstimator<Instant> watermarkEstimator)
      throws IOException {
    BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator =
        new BytesThroughputEstimator<>(sizeEstimator, Instant.now());
    // Lock the partition
    if (tracker.currentRestriction().isEmpty()) {
      boolean lockedPartition = metadataTableDao.lockAndRecordPartition(partitionRecord);
      // Clean up NewPartition on the first run regardless of locking result. If locking fails it
      // means this partition is being streamed, then cleaning up NewPartitions avoids lingering
      // NewPartitions.
      for (NewPartition newPartition : partitionRecord.getParentPartitions()) {
        metadataTableDao.deleteNewPartition(newPartition);
      }
      if (!lockedPartition) {
        LOG.info(
            "RCSP  {} : Could not acquire lock with uid: {}, because this is a "
                + "duplicate and another worker is working  on this partition already.",
            formatByteStringRange(partitionRecord.getPartition()),
            partitionRecord.getUuid());
        StreamProgress streamProgress = new StreamProgress();
        streamProgress.setFailToLock(true);
        metrics.decPartitionStreamCount();
        tracker.tryClaim(streamProgress);
        return ProcessContinuation.stop();
      }
    } else if (tracker.currentRestriction().getCloseStream() == null
        && !metadataTableDao.doHoldLock(
            partitionRecord.getPartition(), partitionRecord.getUuid())) {
      // We only verify the lock if we are not holding CloseStream because if this is a retry of
      // CloseStream we might have already cleaned up the lock in a previous attempt.
      // Failed correctness check on this worker holds the lock on this partition. This shouldn't
      // fail because there's a restriction tracker which means this worker has already acquired the
      // lock and once it has acquired the lock it shouldn't fail the lock check.
      LOG.warn(
          "RCSP  {} : Subsequent run that doesn't hold the lock {}. This is not unexpected and "
              + "should probably be reviewed.",
          formatByteStringRange(partitionRecord.getPartition()),
          partitionRecord.getUuid());
      StreamProgress streamProgress = new StreamProgress();
      streamProgress.setFailToLock(true);
      metrics.decPartitionStreamCount();
      tracker.tryClaim(streamProgress);
      return ProcessContinuation.stop();
    }

    // Process CloseStream if it exists
    CloseStream closeStream = tracker.currentRestriction().getCloseStream();
    if (closeStream != null) {
      LOG.debug("RCSP: Processing CloseStream");
      metrics.decPartitionStreamCount();
      if (closeStream.getStatus().getCode() == Status.Code.OK) {
        // We need to update watermark here. We're terminating this stream because we have reached
        // endTime. Instant.now is greater or equal to endTime. The goal here is
        // DNP will need to know this stream has passed the endTime so DNP can eventually terminate.
        Instant terminatingWatermark = Instant.ofEpochMilli(Long.MAX_VALUE);
        Instant endTime = partitionRecord.getEndTime();
        if (endTime != null) {
          terminatingWatermark = endTime;
        }
        watermarkEstimator.setWatermark(terminatingWatermark);
        metadataTableDao.updateWatermark(
            partitionRecord.getPartition(), watermarkEstimator.currentWatermark(), null);
        LOG.info(
            "RCSP {}: Reached end time, terminating...",
            formatByteStringRange(partitionRecord.getPartition()));
        return ProcessContinuation.stop();
      }
      if (closeStream.getStatus().getCode() != Status.Code.OUT_OF_RANGE) {
        LOG.error(
            "RCSP {}: Reached unexpected terminal state: {}",
            formatByteStringRange(partitionRecord.getPartition()),
            closeStream.getStatus());
        return ProcessContinuation.stop();
      }
      // Release the lock only if the uuid matches. In normal operation this doesn't change
      // anything. However, it's possible for this RCSP to crash while processing CloseStream but
      // after the side effects of writing the new partitions to the metadata table. New partitions
      // can be created while this RCSP restarts from the previous checkpoint and processes the
      // CloseStream again. In certain race scenarios the child partitions may merge back to this
      // partition, but as a new RCSP. The new partition (same as this partition) would write the
      // exact same content to the metadata table but with a different uuid. We don't want to
      // accidentally delete the StreamPartition because it now belongs to the new RCSP.
      // If the uuid is the same (meaning this race scenario did not take place) we release the lock
      // and mark the StreamPartition to be deleted, so we can delete it after we have written the
      // NewPartitions.
      metadataTableDao.releaseStreamPartitionLockForDeletion(
          partitionRecord.getPartition(), partitionRecord.getUuid());
      // The partitions in the continuation tokens must cover the same key space as this partition.
      // If there's only 1 token, then the token's partition is equals to this partition.
      // If there are more than 1 tokens, then the tokens form a continuous row range equals to this
      // partition.
      List<ByteStringRange> childPartitions = new ArrayList<>();
      List<ByteStringRange> tokenPartitions = new ArrayList<>();
      // Check if NewPartitions field exists, if not we default to using just the
      // ChangeStreamContinuationTokens.
      boolean useNewPartitionsField =
          closeStream.getNewPartitions().size()
              == closeStream.getChangeStreamContinuationTokens().size();
      for (int i = 0; i < closeStream.getChangeStreamContinuationTokens().size(); i++) {
        ByteStringRange childPartition;
        if (useNewPartitionsField) {
          childPartition = closeStream.getNewPartitions().get(i);
        } else {
          childPartition = closeStream.getChangeStreamContinuationTokens().get(i).getPartition();
        }
        childPartitions.add(childPartition);
        ChangeStreamContinuationToken token =
            getTokenWithCorrectPartition(
                partitionRecord.getPartition(),
                closeStream.getChangeStreamContinuationTokens().get(i));
        tokenPartitions.add(token.getPartition());
        metadataTableDao.writeNewPartition(
            new NewPartition(
                childPartition, Collections.singletonList(token), watermarkEstimator.getState()));
      }
      LOG.info(
          "RCSP {}: Split/Merge into {}",
          formatByteStringRange(partitionRecord.getPartition()),
          partitionsToString(childPartitions));
      if (!coverSameKeySpace(tokenPartitions, partitionRecord.getPartition())) {
        LOG.warn(
            "RCSP {}: CloseStream has tokens {} that don't cover the entire keyspace",
            formatByteStringRange(partitionRecord.getPartition()),
            partitionsToString(tokenPartitions));
      }
      // Perform the real cleanup. This step is no op if the race mentioned above occurs (splits and
      // merges results back to this partition again) because when we register the "new" partition,
      // we unset the deletion bit.
      metadataTableDao.deleteStreamPartitionRow(partitionRecord.getPartition());
      return ProcessContinuation.stop();
    }

    // Update the metadata table with the watermark
    metadataTableDao.updateWatermark(
        partitionRecord.getPartition(),
        watermarkEstimator.getState(),
        tracker.currentRestriction().getCurrentToken());

    // Start to stream the partition.
    ServerStream<ChangeStreamRecord> stream = null;
    try {
      stream =
          changeStreamDao.readChangeStreamPartition(
              partitionRecord,
              tracker.currentRestriction(),
              partitionRecord.getEndTime(),
              heartbeatDuration);
      for (ChangeStreamRecord record : stream) {
        Optional<ProcessContinuation> result =
            changeStreamAction.run(
                partitionRecord,
                record,
                tracker,
                receiver,
                watermarkEstimator,
                throughputEstimator);
        // changeStreamAction will usually return Optional.empty() except for when a checkpoint
        // (either runner or pipeline initiated) is required.
        if (result.isPresent()) {
          return result.get();
        }
      }
    } catch (Exception e) {
      throw e;
    } finally {
      if (stream != null) {
        stream.cancel();
      }
    }
    return ProcessContinuation.resume();
  }
@Test
  public void testCloseStreamWritesContinuationTokens() throws IOException {
    // Force lock fail because CloseStream should not depend on locking
    when(metadataTableDao.doHoldLock(partition, uuid)).thenReturn(false);
    ChangeStreamContinuationToken tokenAB =
        ChangeStreamContinuationToken.create(ByteStringRange.create("A", "B"), "AB");
    ChangeStreamContinuationToken tokenBC =
        ChangeStreamContinuationToken.create(ByteStringRange.create("B", "C"), "BC");

    CloseStream mockCloseStream = Mockito.mock(CloseStream.class);
    Status statusProto = Status.newBuilder().setCode(11).build();
    Mockito.when(mockCloseStream.getStatus())
        .thenReturn(com.google.cloud.bigtable.common.Status.fromProto(statusProto));
    Mockito.when(mockCloseStream.getChangeStreamContinuationTokens())
        .thenReturn(Arrays.asList(tokenAB, tokenBC));
    Mockito.when(mockCloseStream.getNewPartitions())
        .thenReturn(Arrays.asList(tokenAB.getPartition(), tokenBC.getPartition()));

    when(restriction.getCloseStream()).thenReturn(mockCloseStream);
    when(restriction.isEmpty()).thenReturn(false);
    final DoFn.ProcessContinuation result =
        action.run(partitionRecord, tracker, receiver, watermarkEstimator);
    assertEquals(DoFn.ProcessContinuation.stop(), result);
    // Should terminate before reaching processing stream partition responses.
    verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any());
    // Should not try claim any restriction when processing CloseStream
    verify(tracker, (never())).tryClaim(any());
    // Should decrement the metric on termination.
    verify(metrics).decPartitionStreamCount();
    // Write the new partitions.
    NewPartition newPartitionAB =
        new NewPartition(
            tokenAB.getPartition(),
            Collections.singletonList(tokenAB),
            watermarkEstimator.getState());
    verify(metadataTableDao).writeNewPartition(newPartitionAB);
    NewPartition newPartitionBC =
        new NewPartition(
            tokenBC.getPartition(),
            Collections.singletonList(tokenBC),
            watermarkEstimator.getState());
    verify(metadataTableDao).writeNewPartition(newPartitionBC);
    verify(metadataTableDao, times(1))
        .releaseStreamPartitionLockForDeletion(
            partitionRecord.getPartition(), partitionRecord.getUuid());
    verify(metadataTableDao, times(1)).deleteStreamPartitionRow(partitionRecord.getPartition());
  }