public Optional<DoFn.ProcessContinuation> run(
      PartitionRecord partitionRecord,
      ChangeStreamRecord record,
      RestrictionTracker<StreamProgress, StreamProgress> tracker,
      DoFn.OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver,
      ManualWatermarkEstimator<Instant> watermarkEstimator,
      BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator) {
    if (record instanceof Heartbeat) {
      Heartbeat heartbeat = (Heartbeat) record;
      final Instant watermark = toJodaTime(heartbeat.getEstimatedLowWatermark());

      // These will be filtered so the key doesn't really matter but the most logical thing to
      // key a heartbeat by is the partition it corresponds to.
      ByteString heartbeatKey =
          Range.ByteStringRange.serializeToByteString(partitionRecord.getPartition());
      KV<ByteString, ChangeStreamRecord> outputRecord = KV.of(heartbeatKey, heartbeat);
      throughputEstimator.update(Instant.now(), outputRecord);
      StreamProgress streamProgress =
          new StreamProgress(
              heartbeat.getChangeStreamContinuationToken(),
              watermark,
              throughputEstimator.get(),
              Instant.now(),
              true);
      watermarkEstimator.setWatermark(watermark);

      // If the tracker fail to claim the streamProgress, it most likely means the runner initiated
      // a checkpoint. See {@link
      // org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.ReadChangeStreamPartitionProgressTracker}
      // for more information regarding runner initiated checkpoints.
      if (!tracker.tryClaim(streamProgress)) {
        return Optional.of(DoFn.ProcessContinuation.stop());
      }
      metrics.incHeartbeatCount();
      // We output heartbeats so that they are factored into throughput and can be used to
      // autoscale. These will be filtered in a downstream step and never returned to users. This is
      // to prevent autoscaler from scaling down when we have large tables with no throughput but
      // we need enough workers to keep up with heartbeats.

      // We are outputting elements with timestamp of 0 to prevent reliance on event time. This
      // limits the ability to window on commit time of any data changes. It is still possible to
      // window on processing time.
      receiver.outputWithTimestamp(outputRecord, Instant.EPOCH);
    } else if (record instanceof CloseStream) {
      CloseStream closeStream = (CloseStream) record;
      StreamProgress streamProgress = new StreamProgress(closeStream);

      // If the tracker fail to claim the streamProgress, it most likely means the runner initiated
      // a checkpoint. See {@link
      // org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.ReadChangeStreamPartitionProgressTracker}
      // for more information regarding runner initiated checkpoints.
      if (!tracker.tryClaim(streamProgress)) {
        return Optional.of(DoFn.ProcessContinuation.stop());
      }
      metrics.incClosestreamCount();
      return Optional.of(DoFn.ProcessContinuation.resume());
    } else if (record instanceof ChangeStreamMutation) {
      ChangeStreamMutation changeStreamMutation = (ChangeStreamMutation) record;
      final Instant watermark = toJodaTime(changeStreamMutation.getEstimatedLowWatermark());
      watermarkEstimator.setWatermark(watermark);
      // Build a new StreamProgress with the continuation token to be claimed.
      ChangeStreamContinuationToken changeStreamContinuationToken =
          ChangeStreamContinuationToken.create(
              Range.ByteStringRange.create(
                  partitionRecord.getPartition().getStart(),
                  partitionRecord.getPartition().getEnd()),
              changeStreamMutation.getToken());

      KV<ByteString, ChangeStreamRecord> outputRecord =
          KV.of(changeStreamMutation.getRowKey(), changeStreamMutation);
      throughputEstimator.update(Instant.now(), outputRecord);
      StreamProgress streamProgress =
          new StreamProgress(
              changeStreamContinuationToken,
              watermark,
              throughputEstimator.get(),
              Instant.now(),
              false);
      // If the tracker fail to claim the streamProgress, it most likely means the runner initiated
      // a checkpoint. See ReadChangeStreamPartitionProgressTracker for more information regarding
      // runner initiated checkpoints.
      if (!tracker.tryClaim(streamProgress)) {
        return Optional.of(DoFn.ProcessContinuation.stop());
      }
      if (changeStreamMutation.getType() == ChangeStreamMutation.MutationType.GARBAGE_COLLECTION) {
        metrics.incChangeStreamMutationGcCounter();
      } else if (changeStreamMutation.getType() == ChangeStreamMutation.MutationType.USER) {
        metrics.incChangeStreamMutationUserCounter();
      }
      Instant delay = toJodaTime(changeStreamMutation.getCommitTimestamp());
      metrics.updateProcessingDelayFromCommitTimestamp(
          Instant.now().getMillis() - delay.getMillis());

      // We are outputting elements with timestamp of 0 to prevent reliance on event time. This
      // limits the ability to window on commit time of any data changes. It is still possible to
      // window on processing time.
      receiver.outputWithTimestamp(outputRecord, Instant.EPOCH);
    } else {
      LOG.warn(
          "RCSP {}: Invalid response type", formatByteStringRange(partitionRecord.getPartition()));
    }
    return Optional.empty();
  }
@Test
  public void testHeartBeat() {
    final Instant lowWatermark = Instant.ofEpochSecond(1000);
    ChangeStreamContinuationToken changeStreamContinuationToken =
        ChangeStreamContinuationToken.create(ByteStringRange.create("a", "b"), "1234");
    Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class);
    Mockito.when(mockHeartBeat.getEstimatedLowWatermark())
        .thenReturn(toThreetenInstant(lowWatermark));
    Mockito.when(mockHeartBeat.getChangeStreamContinuationToken())
        .thenReturn(changeStreamContinuationToken);

    final Optional<DoFn.ProcessContinuation> result =
        action.run(
            partitionRecord,
            mockHeartBeat,
            tracker,
            receiver,
            watermarkEstimator,
            throughputEstimator);

    assertFalse(result.isPresent());
    verify(metrics).incHeartbeatCount();
    verify(watermarkEstimator).setWatermark(eq(lowWatermark));
    StreamProgress streamProgress =
        new StreamProgress(
            changeStreamContinuationToken,
            lowWatermark,
            BigDecimal.valueOf(1000),
            Instant.now(),
            true);
    verify(tracker).tryClaim(streamProgressArgumentCaptor.capture());
    assertEquals(
        streamProgress.getCurrentToken(),
        streamProgressArgumentCaptor.getValue().getCurrentToken());
    assertEquals(
        streamProgress.getThroughputEstimate(),
        streamProgressArgumentCaptor.getValue().getThroughputEstimate());
    assertEquals(
        streamProgress.getEstimatedLowWatermark(),
        streamProgressArgumentCaptor.getValue().getEstimatedLowWatermark());
    assertEquals(
        streamProgress.isHeartbeat(), streamProgressArgumentCaptor.getValue().isHeartbeat());
    KV<ByteString, ChangeStreamRecord> record =
        KV.of(ByteStringRange.serializeToByteString(partitionRecord.getPartition()), mockHeartBeat);
    verify(receiver).outputWithTimestamp(eq(record), eq(Instant.EPOCH));
    verify(throughputEstimator).update(any(), eq(record));
  }