From 7673b5d99a95310432a6f60078a9739c175a93f3 Mon Sep 17 00:00:00 2001
From: Jack Dingilian <jackdingilian@google.com>
Date: Tue, 27 Jun 2023 16:54:19 -0400
Subject: [PATCH] Cleanup for GA launch of ReadChangeStream (#27249)

This includes a series of small changes to release the connector:
- Remove option for users to override heartbeat duration
- Add option to skip creating metadata table and standalone utility to create metadata table
- Add hard timeout to mutateRow requests as workaround for hanging requests
- Remove excessive RCSP logging for debugging during preview
- Increase read rows timeout to better accomodate large tables
- Remove side effects from InitializeDoFn
- Lower ReadChangeStream deadline to closer align with checkpoint duration
- Add release note to CHANGES.md
---
 CHANGES.md                                    |   1 +
 .../beam/sdk/io/gcp/bigtable/BigtableIO.java  | 133 ++++++++++++++----
 .../action/ChangeStreamAction.java            |  41 +-----
 .../ReadChangeStreamPartitionAction.java      |  31 +---
 .../dao/BigtableChangeStreamAccessor.java     |  44 ++++--
 .../changestreams/dao/ChangeStreamDao.java    |  14 +-
 .../changestreams/dao/DaoFactory.java         |  15 +-
 .../dao/MetadataTableAdminDao.java            |   5 +
 .../changestreams/dao/MetadataTableDao.java   |  47 ++++++-
 .../changestreams/dofn/InitializeDoFn.java    |  35 +----
 .../dofn/ReadChangeStreamPartitionDoFn.java   |  11 +-
 .../action/ChangeStreamActionTest.java        |  12 +-
 .../ReadChangeStreamPartitionActionTest.java  |  33 +++--
 .../dao/MetadataTableDaoTest.java             |  31 ++++
 .../dofn/InitializeDoFnTest.java              |  71 +---------
 .../ReadChangeStreamPartitionDoFnTest.java    |   5 +-
 16 files changed, 271 insertions(+), 258 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 6c55e276f91f..aee9b96d9f4a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -59,6 +59,7 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Support for Bigtable Change Streams added in Java `BigtableIO.ReadChangeStream` ([#27183](https://github.com/apache/beam/issues/27183))
 
 ## New Features / Improvements
 
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
index 319c31390217..06497458a66a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/BigtableIO.java
@@ -48,6 +48,7 @@
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamMetrics;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.UniqueIdGenerator;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.action.ActionFactory;
+import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.BigtableChangeStreamAccessor;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.DaoFactory;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao.MetadataTableAdminDao;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.dofn.DetectNewPartitionsDoFn;
@@ -310,7 +311,6 @@ public static Write write() {
    *
    * <ul>
    *   <li>{@link BigtableIO.ReadChangeStream#withStartTime} which defaults to now.
-   *   <li>{@link BigtableIO.ReadChangeStream#withHeartbeatDuration} with defaults to 1 seconds.
    *   <li>{@link BigtableIO.ReadChangeStream#withMetadataTableProjectId} which defaults to value
    *       from {@link BigtableIO.ReadChangeStream#withProjectId}
    *   <li>{@link BigtableIO.ReadChangeStream#withMetadataTableInstanceId} which defaults to value
@@ -1797,8 +1797,6 @@ public enum ExistingPipelineOptions {
     RESUME_OR_NEW,
     // Same as RESUME_OR_NEW except if previous pipeline doesn't exist, don't start.
     RESUME_OR_FAIL,
-    // Start a new pipeline. Overriding existing pipeline with the same name.
-    NEW,
     // This skips cleaning up previous pipeline metadata and starts a new pipeline. This should
     // only be used to skip cleanup in tests
     @VisibleForTesting
@@ -1827,8 +1825,6 @@ static ReadChangeStream create() {
 
     abstract @Nullable Instant getEndTime();
 
-    abstract @Nullable Duration getHeartbeatDuration();
-
     abstract @Nullable String getChangeStreamName();
 
     abstract @Nullable ExistingPipelineOptions getExistingPipelineOptions();
@@ -1837,6 +1833,8 @@ static ReadChangeStream create() {
 
     abstract @Nullable String getMetadataTableId();
 
+    abstract @Nullable Boolean getCreateOrUpdateMetadataTable();
+
     abstract ReadChangeStream.Builder toBuilder();
 
     /**
@@ -1909,16 +1907,6 @@ ReadChangeStream withEndTime(Instant endTime) {
       return toBuilder().setEndTime(endTime).build();
     }
 
-    /**
-     * Returns a new {@link BigtableIO.ReadChangeStream} that will send heartbeat messages at
-     * specified interval.
-     *
-     * <p>Does not modify this object.
-     */
-    public ReadChangeStream withHeartbeatDuration(Duration interval) {
-      return toBuilder().setHeartbeatDuration(interval).build();
-    }
-
     /**
      * Returns a new {@link BigtableIO.ReadChangeStream} that uses changeStreamName as prefix for
      * the metadata table.
@@ -2000,6 +1988,19 @@ public ReadChangeStream withMetadataTableAppProfileId(String appProfileId) {
           .build();
     }
 
+    /**
+     * Returns a new {@link BigtableIO.ReadChangeStream} that, if set to true, will create or update
+     * metadata table before launching pipeline. Otherwise, it is expected that a metadata table
+     * with correct schema exists.
+     *
+     * <p>Optional: defaults to true
+     *
+     * <p>Does not modify this object.
+     */
+    public ReadChangeStream withCreateOrUpdateMetadataTable(boolean shouldCreate) {
+      return toBuilder().setCreateOrUpdateMetadataTable(shouldCreate).build();
+    }
+
     @Override
     public PCollection<KV<ByteString, ChangeStreamMutation>> expand(PBegin input) {
       checkArgument(
@@ -2040,10 +2041,6 @@ public PCollection<KV<ByteString, ChangeStreamMutation>> expand(PBegin input) {
       if (startTime == null) {
         startTime = Instant.now();
       }
-      Duration heartbeatDuration = getHeartbeatDuration();
-      if (heartbeatDuration == null) {
-        heartbeatDuration = Duration.standardSeconds(1);
-      }
       String changeStreamName = getChangeStreamName();
       if (changeStreamName == null || changeStreamName.isEmpty()) {
         changeStreamName = UniqueIdGenerator.generateRowKeyPrefix();
@@ -2053,21 +2050,55 @@ public PCollection<KV<ByteString, ChangeStreamMutation>> expand(PBegin input) {
         existingPipelineOptions = ExistingPipelineOptions.FAIL_IF_EXISTS;
       }
 
+      boolean shouldCreateOrUpdateMetadataTable = true;
+      if (getCreateOrUpdateMetadataTable() != null) {
+        shouldCreateOrUpdateMetadataTable = getCreateOrUpdateMetadataTable();
+      }
+
       ActionFactory actionFactory = new ActionFactory();
+      ChangeStreamMetrics metrics = new ChangeStreamMetrics();
       DaoFactory daoFactory =
           new DaoFactory(
               bigtableConfig, metadataTableConfig, getTableId(), metadataTableId, changeStreamName);
-      ChangeStreamMetrics metrics = new ChangeStreamMetrics();
+
+      try {
+        MetadataTableAdminDao metadataTableAdminDao = daoFactory.getMetadataTableAdminDao();
+        checkArgument(metadataTableAdminDao != null);
+        checkArgument(
+            metadataTableAdminDao.isAppProfileSingleClusterAndTransactional(
+                metadataTableConfig.getAppProfileId().get()),
+            "App profile id '"
+                + metadataTableConfig.getAppProfileId().get()
+                + "' provided to access metadata table needs to use single-cluster routing policy"
+                + " and allow single-row transactions.");
+
+        // Only try to create or update metadata table if option is set to true. Otherwise, just
+        // check if the table exists.
+        if (shouldCreateOrUpdateMetadataTable && metadataTableAdminDao.createMetadataTable()) {
+          LOG.info("Created metadata table: " + metadataTableAdminDao.getTableId());
+        }
+        checkArgument(
+            metadataTableAdminDao.doesMetadataTableExist(),
+            "Metadata table does not exist: " + metadataTableAdminDao.getTableId());
+
+        try (BigtableChangeStreamAccessor bigtableChangeStreamAccessor =
+            BigtableChangeStreamAccessor.getOrCreate(bigtableConfig)) {
+          checkArgument(
+              bigtableChangeStreamAccessor.getTableAdminClient().exists(getTableId()),
+              "Change Stream table does not exist");
+        }
+      } catch (Exception e) {
+        throw new RuntimeException(e);
+      } finally {
+        daoFactory.close();
+      }
+
       InitializeDoFn initializeDoFn =
-          new InitializeDoFn(
-              daoFactory,
-              metadataTableConfig.getAppProfileId().get(),
-              startTime,
-              existingPipelineOptions);
+          new InitializeDoFn(daoFactory, startTime, existingPipelineOptions);
       DetectNewPartitionsDoFn detectNewPartitionsDoFn =
           new DetectNewPartitionsDoFn(getEndTime(), actionFactory, daoFactory, metrics);
       ReadChangeStreamPartitionDoFn readChangeStreamPartitionDoFn =
-          new ReadChangeStreamPartitionDoFn(heartbeatDuration, daoFactory, actionFactory, metrics);
+          new ReadChangeStreamPartitionDoFn(daoFactory, actionFactory, metrics);
 
       PCollection<KV<ByteString, ChangeStreamRecord>> readChangeStreamOutput =
           input
@@ -2101,14 +2132,60 @@ abstract ReadChangeStream.Builder setMetadataTableBigtableConfig(
 
       abstract ReadChangeStream.Builder setEndTime(Instant endTime);
 
-      abstract ReadChangeStream.Builder setHeartbeatDuration(Duration interval);
-
       abstract ReadChangeStream.Builder setChangeStreamName(String changeStreamName);
 
       abstract ReadChangeStream.Builder setExistingPipelineOptions(
           ExistingPipelineOptions existingPipelineOptions);
 
+      abstract ReadChangeStream.Builder setCreateOrUpdateMetadataTable(boolean shouldCreate);
+
       abstract ReadChangeStream build();
     }
   }
+
+  /**
+   * Utility method to create or update Read Change Stream metadata table. This requires Bigtable
+   * table create permissions. This method is useful if the pipeline isn't granted permissions to
+   * create Bigtable tables. Run this method with correct permissions to create the metadata table,
+   * which is required to read Bigtable change streams. This method only needs to be run once, and
+   * the metadata table can be reused for all pipelines.
+   *
+   * @param projectId project id of the metadata table, usually the same as the project of the table
+   *     being streamed
+   * @param instanceId instance id of the metadata table, usually the same as the instance of the
+   *     table being streamed
+   * @param tableId name of the metadata table, leave it null or empty to use default.
+   * @return true if the table was successfully created. Otherwise, false.
+   */
+  public static boolean createOrUpdateReadChangeStreamMetadataTable(
+      String projectId, String instanceId, @Nullable String tableId) throws IOException {
+    BigtableConfig bigtableConfig =
+        BigtableConfig.builder()
+            .setValidate(true)
+            .setProjectId(StaticValueProvider.of(projectId))
+            .setInstanceId(StaticValueProvider.of(instanceId))
+            .setAppProfileId(
+                StaticValueProvider.of(
+                    "default")) // App profile is not used. It's only required for data API.
+            .build();
+
+    if (tableId == null || tableId.isEmpty()) {
+      tableId = MetadataTableAdminDao.DEFAULT_METADATA_TABLE_NAME;
+    }
+
+    DaoFactory daoFactory = new DaoFactory(null, bigtableConfig, null, tableId, null);
+
+    try {
+      MetadataTableAdminDao metadataTableAdminDao = daoFactory.getMetadataTableAdminDao();
+
+      // Only try to create or update metadata table if option is set to true. Otherwise, just
+      // check if the table exists.
+      if (metadataTableAdminDao.createMetadataTable()) {
+        LOG.info("Created metadata table: " + metadataTableAdminDao.getTableId());
+      }
+      return metadataTableAdminDao.doesMetadataTableExist();
+    } finally {
+      daoFactory.close();
+    }
+  }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java
index 3ab3b1973b97..adb817a4c08a 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamAction.java
@@ -28,7 +28,6 @@
 import com.google.cloud.bigtable.data.v2.models.Range;
 import com.google.protobuf.ByteString;
 import java.util.Optional;
-import java.util.stream.Collectors;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamMetrics;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.estimator.BytesThroughputEstimator;
@@ -108,8 +107,7 @@ public Optional<DoFn.ProcessContinuation> run(
       RestrictionTracker<StreamProgress, StreamProgress> tracker,
       DoFn.OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver,
       ManualWatermarkEstimator<Instant> watermarkEstimator,
-      BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator,
-      boolean shouldDebug) {
+      BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator) {
     if (record instanceof Heartbeat) {
       Heartbeat heartbeat = (Heartbeat) record;
       final Instant watermark = toJodaTime(heartbeat.getEstimatedLowWatermark());
@@ -129,24 +127,11 @@ public Optional<DoFn.ProcessContinuation> run(
               true);
       watermarkEstimator.setWatermark(watermark);
 
-      if (shouldDebug) {
-        LOG.info(
-            "RCSP {}: Heartbeat partition: {} token: {} watermark: {}",
-            formatByteStringRange(partitionRecord.getPartition()),
-            formatByteStringRange(heartbeat.getChangeStreamContinuationToken().getPartition()),
-            heartbeat.getChangeStreamContinuationToken().getToken(),
-            heartbeat.getEstimatedLowWatermark());
-      }
       // 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)) {
-        if (shouldDebug) {
-          LOG.info(
-              "RCSP {}: Checkpoint heart beat tracker",
-              formatByteStringRange(partitionRecord.getPartition()));
-        }
         return Optional.of(DoFn.ProcessContinuation.stop());
       }
       metrics.incHeartbeatCount();
@@ -163,30 +148,11 @@ public Optional<DoFn.ProcessContinuation> run(
       CloseStream closeStream = (CloseStream) record;
       StreamProgress streamProgress = new StreamProgress(closeStream);
 
-      if (shouldDebug) {
-        LOG.info(
-            "RCSP {}: CloseStream: {}",
-            formatByteStringRange(partitionRecord.getPartition()),
-            closeStream.getChangeStreamContinuationTokens().stream()
-                .map(
-                    c ->
-                        "{partition: "
-                            + formatByteStringRange(c.getPartition())
-                            + " token: "
-                            + c.getToken()
-                            + "}")
-                .collect(Collectors.joining(", ", "[", "]")));
-      }
       // 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)) {
-        if (shouldDebug) {
-          LOG.info(
-              "RCSP {}: Checkpoint close stream tracker",
-              formatByteStringRange(partitionRecord.getPartition()));
-        }
         return Optional.of(DoFn.ProcessContinuation.stop());
       }
       metrics.incClosestreamCount();
@@ -217,11 +183,6 @@ public Optional<DoFn.ProcessContinuation> run(
       // a checkpoint. See ReadChangeStreamPartitionProgressTracker for more information regarding
       // runner initiated checkpoints.
       if (!tracker.tryClaim(streamProgress)) {
-        if (shouldDebug) {
-          LOG.info(
-              "RCSP {}: Checkpoint data change tracker",
-              formatByteStringRange(partitionRecord.getPartition()));
-        }
         return Optional.of(DoFn.ProcessContinuation.stop());
       }
       if (changeStreamMutation.getType() == ChangeStreamMutation.MutationType.GARBAGE_COLLECTION) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java
index 87273c9894d9..8638339ac5c0 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionAction.java
@@ -138,23 +138,8 @@ public ProcessContinuation run(
       OutputReceiver<KV<ByteString, ChangeStreamRecord>> receiver,
       ManualWatermarkEstimator<Instant> watermarkEstimator)
       throws IOException {
-    // Watermark being delayed beyond 5 minutes signals a possible problem.
-    boolean shouldDebug =
-        watermarkEstimator.getState().plus(Duration.standardMinutes(5)).isBeforeNow();
     BytesThroughputEstimator<KV<ByteString, ChangeStreamRecord>> throughputEstimator =
         new BytesThroughputEstimator<>(sizeEstimator, Instant.now());
-
-    if (shouldDebug) {
-      LOG.info(
-          "RCSP {}: Partition: "
-              + partitionRecord
-              + "\n Watermark: "
-              + watermarkEstimator.getState()
-              + "\n RestrictionTracker: "
-              + tracker.currentRestriction(),
-          formatByteStringRange(partitionRecord.getPartition()));
-    }
-
     // Lock the partition
     if (tracker.currentRestriction().isEmpty()) {
       boolean lockedPartition = metadataTableDao.lockAndRecordPartition(partitionRecord);
@@ -266,12 +251,10 @@ public ProcessContinuation run(
             new NewPartition(
                 childPartition, Collections.singletonList(token), watermarkEstimator.getState()));
       }
-      if (shouldDebug) {
-        LOG.info(
-            "RCSP {}: Split/Merge into {}",
-            formatByteStringRange(partitionRecord.getPartition()),
-            partitionsToString(childPartitions));
-      }
+      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",
@@ -299,8 +282,7 @@ public ProcessContinuation run(
               partitionRecord,
               tracker.currentRestriction(),
               partitionRecord.getEndTime(),
-              heartbeatDuration,
-              shouldDebug);
+              heartbeatDuration);
       for (ChangeStreamRecord record : stream) {
         Optional<ProcessContinuation> result =
             changeStreamAction.run(
@@ -309,8 +291,7 @@ public ProcessContinuation run(
                 tracker,
                 receiver,
                 watermarkEstimator,
-                throughputEstimator,
-                shouldDebug);
+                throughputEstimator);
         // changeStreamAction will usually return Optional.empty() except for when a checkpoint
         // (either runner or pipeline initiated) is required.
         if (result.isPresent()) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java
index 69c6efba1de8..ecf9a7039598 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/BigtableChangeStreamAccessor.java
@@ -47,7 +47,8 @@
  * backend and the jobs on the same machine shares the same sets of connections.
  */
 @Internal
-class BigtableChangeStreamAccessor {
+public class BigtableChangeStreamAccessor implements AutoCloseable {
+  static final Duration MUTATE_ROW_DEADLINE = Duration.ofSeconds(30);
   // Create one bigtable data/admin client per bigtable config (project/instance/table/app profile)
   private static final ConcurrentHashMap<BigtableConfig, BigtableChangeStreamAccessor>
       bigtableAccessors = new ConcurrentHashMap<>();
@@ -55,14 +56,31 @@ class BigtableChangeStreamAccessor {
   private final BigtableDataClient dataClient;
   private final BigtableTableAdminClient tableAdminClient;
   private final BigtableInstanceAdminClient instanceAdminClient;
+  private final BigtableConfig bigtableConfig;
 
   private BigtableChangeStreamAccessor(
       BigtableDataClient dataClient,
       BigtableTableAdminClient tableAdminClient,
-      BigtableInstanceAdminClient instanceAdminClient) {
+      BigtableInstanceAdminClient instanceAdminClient,
+      BigtableConfig bigtableConfig) {
     this.dataClient = dataClient;
     this.tableAdminClient = tableAdminClient;
     this.instanceAdminClient = instanceAdminClient;
+    this.bigtableConfig = bigtableConfig;
+  }
+
+  @Override
+  public synchronized void close() {
+    if (dataClient != null) {
+      dataClient.close();
+    }
+    if (tableAdminClient != null) {
+      tableAdminClient.close();
+    }
+    if (instanceAdminClient != null) {
+      instanceAdminClient.close();
+    }
+    bigtableAccessors.remove(bigtableConfig);
   }
 
   /**
@@ -146,9 +164,10 @@ private static BigtableChangeStreamAccessor createAccessor(@NonNull BigtableConf
         .readRowsSettings()
         .setRetrySettings(
             readRowsRetrySettings
-                .setInitialRpcTimeout(Duration.ofSeconds(30))
-                .setTotalTimeout(Duration.ofSeconds(30))
-                .setMaxRpcTimeout(Duration.ofSeconds(30))
+                // metadata table scans can get quite large, so use a higher deadline
+                .setInitialRpcTimeout(Duration.ofMinutes(3))
+                .setTotalTimeout(Duration.ofMinutes(3))
+                .setMaxRpcTimeout(Duration.ofMinutes(3))
                 .setMaxAttempts(10)
                 .build());
 
@@ -159,9 +178,9 @@ private static BigtableChangeStreamAccessor createAccessor(@NonNull BigtableConf
         .mutateRowSettings()
         .setRetrySettings(
             mutateRowRetrySettings
-                .setInitialRpcTimeout(Duration.ofSeconds(30))
-                .setTotalTimeout(Duration.ofSeconds(30))
-                .setMaxRpcTimeout(Duration.ofSeconds(30))
+                .setInitialRpcTimeout(MUTATE_ROW_DEADLINE)
+                .setTotalTimeout(MUTATE_ROW_DEADLINE)
+                .setMaxRpcTimeout(MUTATE_ROW_DEADLINE)
                 .setMaxAttempts(10)
                 .build());
 
@@ -185,9 +204,9 @@ private static BigtableChangeStreamAccessor createAccessor(@NonNull BigtableConf
         .readChangeStreamSettings()
         .setRetrySettings(
             readChangeStreamRetrySettings
-                .setInitialRpcTimeout(Duration.ofSeconds(60))
-                .setTotalTimeout(Duration.ofSeconds(60))
-                .setMaxRpcTimeout(Duration.ofSeconds(60))
+                .setInitialRpcTimeout(Duration.ofSeconds(15))
+                .setTotalTimeout(Duration.ofSeconds(15))
+                .setMaxRpcTimeout(Duration.ofSeconds(15))
                 .setMaxAttempts(10)
                 .build());
 
@@ -196,7 +215,8 @@ private static BigtableChangeStreamAccessor createAccessor(@NonNull BigtableConf
         BigtableTableAdminClient.create(tableAdminSettingsBuilder.build());
     BigtableInstanceAdminClient instanceAdminClient =
         BigtableInstanceAdminClient.create(instanceAdminSettingsBuilder.build());
-    return new BigtableChangeStreamAccessor(dataClient, tableAdminClient, instanceAdminClient);
+    return new BigtableChangeStreamAccessor(
+        dataClient, tableAdminClient, instanceAdminClient, bigtableConfig);
   }
 
   public BigtableDataClient getDataClient() {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java
index 1c24e09db7da..5974af22f863 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/ChangeStreamDao.java
@@ -17,7 +17,6 @@
  */
 package org.apache.beam.sdk.io.gcp.bigtable.changestreams.dao;
 
-import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ByteStringRangeHelper.formatByteStringRange;
 import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.TimestampConverter.toThreetenInstant;
 
 import com.google.api.gax.rpc.ServerStream;
@@ -36,14 +35,10 @@
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.restriction.StreamProgress;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /** Data access object to list and read stream partitions of a table. */
 @Internal
 public class ChangeStreamDao {
-  private static final Logger LOG = LoggerFactory.getLogger(ChangeStreamDao.class);
-
   private final BigtableDataClient dataClient;
   private final String tableId;
 
@@ -75,8 +70,7 @@ public ServerStream<ChangeStreamRecord> readChangeStreamPartition(
       PartitionRecord partition,
       StreamProgress streamProgress,
       @Nullable Instant endTime,
-      Duration heartbeatDuration,
-      boolean shouldDebug)
+      Duration heartbeatDuration)
       throws IOException {
     ReadChangeStreamQuery query =
         ReadChangeStreamQuery.create(tableId).streamPartition(partition.getPartition());
@@ -98,12 +92,6 @@ public ServerStream<ChangeStreamRecord> readChangeStreamPartition(
       query.endTime(TimestampConverter.toThreetenInstant(endTime));
     }
     query.heartbeatDuration(org.threeten.bp.Duration.ofMillis(heartbeatDuration.getMillis()));
-    if (shouldDebug) {
-      LOG.info(
-          "RCSP {} ReadChangeStreamRequest: {}",
-          formatByteStringRange(partition.getPartition()),
-          query);
-    }
     return dataClient.readChangeStream(query);
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java
index acecc22b8919..35ac8ed29c3e 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/DaoFactory.java
@@ -30,7 +30,7 @@
 // Allows transient fields to be intialized later
 @SuppressWarnings("initialization.fields.uninitialized")
 @Internal
-public class DaoFactory implements Serializable {
+public class DaoFactory implements Serializable, AutoCloseable {
   private static final long serialVersionUID = 3732208768248394205L;
 
   private transient ChangeStreamDao changeStreamDao;
@@ -58,6 +58,19 @@ public DaoFactory(
     this.metadataTableId = metadataTableId;
   }
 
+  @Override
+  public void close() {
+    try {
+      if (metadataTableAdminDao != null || metadataTableDao != null) {
+        BigtableChangeStreamAccessor.getOrCreate(metadataTableConfig).close();
+      }
+      if (changeStreamDao != null) {
+        BigtableChangeStreamAccessor.getOrCreate(changeStreamConfig).close();
+      }
+    } catch (Exception ignored) {
+    }
+  }
+
   public String getChangeStreamName() {
     return changeStreamName;
   }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java
index 5194ca49d62e..73bfbe5865b4 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableAdminDao.java
@@ -128,6 +128,11 @@ public boolean isAppProfileSingleClusterAndTransactional(String appProfileId) {
     return false;
   }
 
+  /** @return true if metadata table exists, otherwise false. */
+  public boolean doesMetadataTableExist() {
+    return tableAdminClient.exists(tableId);
+  }
+
   /**
    * Create the metadata table if it does not exist yet. If the table does exist, verify all the
    * column families exists, if not add those column families. This table only need to be created
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java
index d4938e6bef26..418c29d1de76 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDao.java
@@ -30,6 +30,7 @@
 import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder.parseWatermarkFromRow;
 import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder.parseWatermarkLastUpdatedFromRow;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.models.ChangeStreamContinuationToken;
@@ -50,6 +51,9 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.apache.beam.repackaged.core.org.apache.commons.lang3.SerializationException;
@@ -60,6 +64,7 @@
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.NewPartition;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.PartitionRecord;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.StreamPartitionWithWatermark;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Longs;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -319,7 +324,7 @@ public void writeNewPartition(NewPartition newPartition) {
             .deleteCells(
                 MetadataTableAdminDao.CF_SHOULD_DELETE,
                 ByteStringRange.serializeToByteString(parentPartition));
-    dataClient.mutateRow(rowMutation);
+    mutateRowWithHardTimeout(rowMutation);
   }
 
   /**
@@ -352,7 +357,7 @@ public void markNewPartitionForDeletion(NewPartition newPartition) {
           ByteStringRange.serializeToByteString(token.getPartition()),
           1);
     }
-    dataClient.mutateRow(rowMutation);
+    mutateRowWithHardTimeout(rowMutation);
   }
 
   /**
@@ -511,7 +516,7 @@ private void writeToMdTableWatermarkHelper(
           MetadataTableAdminDao.QUALIFIER_DEFAULT,
           currentToken.getToken());
     }
-    dataClient.mutateRow(rowMutation);
+    mutateRowWithHardTimeout(rowMutation);
   }
 
   /**
@@ -705,6 +710,10 @@ tableId, convertPartitionToStreamPartitionRowKey(partitionRecord.getPartition())
 
     boolean lockAcquired = !dataClient.checkAndMutateRow(rowMutation);
     if (lockAcquired) {
+      LOG.info(
+          "RCSP: {} acquired lock for uid: {}",
+          formatByteStringRange(partitionRecord.getPartition()),
+          partitionRecord.getUuid());
       return true;
     } else {
       // If the lock is already held we need to check if it was acquired by a duplicate
@@ -724,7 +733,7 @@ public void writeDetectNewPartitionVersion() {
                 MetadataTableAdminDao.CF_VERSION,
                 MetadataTableAdminDao.QUALIFIER_DEFAULT,
                 MetadataTableAdminDao.CURRENT_METADATA_TABLE_VERSION);
-    dataClient.mutateRow(rowMutation);
+    mutateRowWithHardTimeout(rowMutation);
   }
 
   /**
@@ -778,6 +787,34 @@ public void writeDetectNewPartitionMissingPartitions(
                 MetadataTableAdminDao.CF_MISSING_PARTITIONS,
                 ByteString.copyFromUtf8(MetadataTableAdminDao.QUALIFIER_DEFAULT),
                 ByteString.copyFrom(serializedMissingPartition));
-    dataClient.mutateRow(rowMutation);
+    mutateRowWithHardTimeout(rowMutation);
+  }
+
+  /**
+   * This adds a hard timeout of 40 seconds to mutate row futures. These requests already have a
+   * 30-second deadline. This is a workaround for an extremely rare issue we see with requests not
+   * respecting their deadlines. This can be removed once we've pinpointed the cause.
+   *
+   * @param rowMutation Bigtable RowMutation to apply
+   */
+  @VisibleForTesting
+  void mutateRowWithHardTimeout(RowMutation rowMutation) {
+    ApiFuture<Void> mutateRowFuture = dataClient.mutateRowAsync(rowMutation);
+    try {
+      mutateRowFuture.get(
+          BigtableChangeStreamAccessor.MUTATE_ROW_DEADLINE.getSeconds() + 10, TimeUnit.SECONDS);
+    } catch (TimeoutException timeoutException) {
+      mutateRowFuture.cancel(true);
+      throw new RuntimeException(
+          "Cancelled mutateRow request after exceeding deadline", timeoutException);
+    } catch (ExecutionException executionException) {
+      if (executionException.getCause() instanceof RuntimeException) {
+        throw (RuntimeException) executionException.getCause();
+      }
+      throw new RuntimeException(executionException);
+    } catch (InterruptedException interruptedException) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException(interruptedException);
+    }
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java
index 07826ee1db73..c8ee12772cc0 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFn.java
@@ -36,20 +36,14 @@
 @Internal
 public class InitializeDoFn extends DoFn<byte[], InitialPipelineState> implements Serializable {
   private static final long serialVersionUID = 1868189906451252363L;
-
   private static final Logger LOG = LoggerFactory.getLogger(InitializeDoFn.class);
   private final DaoFactory daoFactory;
-  private final String metadataTableAppProfileId;
   private Instant startTime;
   private final ExistingPipelineOptions existingPipelineOptions;
 
   public InitializeDoFn(
-      DaoFactory daoFactory,
-      String metadataTableAppProfileId,
-      Instant startTime,
-      ExistingPipelineOptions existingPipelineOptions) {
+      DaoFactory daoFactory, Instant startTime, ExistingPipelineOptions existingPipelineOptions) {
     this.daoFactory = daoFactory;
-    this.metadataTableAppProfileId = metadataTableAppProfileId;
     this.startTime = startTime;
     this.existingPipelineOptions = existingPipelineOptions;
   }
@@ -59,35 +53,12 @@ public void processElement(OutputReceiver<InitialPipelineState> receiver) throws
     LOG.info(daoFactory.getStreamTableDebugString());
     LOG.info(daoFactory.getMetadataTableDebugString());
     LOG.info("ChangeStreamName: " + daoFactory.getChangeStreamName());
-    if (!daoFactory
-        .getMetadataTableAdminDao()
-        .isAppProfileSingleClusterAndTransactional(this.metadataTableAppProfileId)) {
-      LOG.error(
-          "App profile id '"
-              + metadataTableAppProfileId
-              + "' provided to access metadata table needs to use single-cluster routing policy"
-              + " and allow single-row transactions.");
-      // Terminate this pipeline now.
-      return;
-    }
-    if (daoFactory.getMetadataTableAdminDao().createMetadataTable()) {
-      LOG.info("Created metadata table: " + daoFactory.getMetadataTableAdminDao().getTableId());
-    } else {
-      LOG.info(
-          "Reusing existing metadata table: " + daoFactory.getMetadataTableAdminDao().getTableId());
-    }
 
     boolean resume = false;
     DetectNewPartitionsState detectNewPartitionsState =
         daoFactory.getMetadataTableDao().readDetectNewPartitionsState();
 
     switch (existingPipelineOptions) {
-      case NEW:
-        // clean up table
-        LOG.info(
-            "Cleaning up an old pipeline with the same change stream name to start a new pipeline with the same name.");
-        daoFactory.getMetadataTableAdminDao().cleanUpPrefix();
-        break;
       case RESUME_OR_NEW:
         // perform resumption.
         if (detectNewPartitionsState != null) {
@@ -98,7 +69,6 @@ public void processElement(OutputReceiver<InitialPipelineState> receiver) throws
           LOG.info(
               "Attempted to resume, but previous watermark does not exist, starting at {}",
               startTime);
-          daoFactory.getMetadataTableAdminDao().cleanUpPrefix();
         }
         break;
       case RESUME_OR_FAIL:
@@ -118,9 +88,6 @@ public void processElement(OutputReceiver<InitialPipelineState> receiver) throws
               "A previous pipeline exists with the same change stream name and existingPipelineOption is set to FAIL_IF_EXISTS.");
           return;
         }
-        // We still want to clean up any existing prefixes in case there are lingering metadata that
-        // would interfere with the new run.
-        daoFactory.getMetadataTableAdminDao().cleanUpPrefix();
         break;
       case SKIP_CLEANUP:
         if (detectNewPartitionsState != null) {
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java
index 91210caa5972..e4c911f926a3 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFn.java
@@ -58,9 +58,8 @@ public class ReadChangeStreamPartitionDoFn
   private static final long serialVersionUID = 4418739381635104479L;
   private static final BigDecimal MAX_DOUBLE = BigDecimal.valueOf(Double.MAX_VALUE);
   private static final Logger LOG = LoggerFactory.getLogger(ReadChangeStreamPartitionDoFn.class);
-  public static final int THROUGHPUT_ESTIMATION_WINDOW_SECONDS = 10;
+  private static final Duration HEARTBEAT_DURATION = Duration.standardSeconds(1);
 
-  private final Duration heartbeatDuration;
   private final DaoFactory daoFactory;
   private final ChangeStreamMetrics metrics;
   private final ActionFactory actionFactory;
@@ -68,11 +67,7 @@ public class ReadChangeStreamPartitionDoFn
   private ReadChangeStreamPartitionAction readChangeStreamPartitionAction;
 
   public ReadChangeStreamPartitionDoFn(
-      Duration heartbeatDuration,
-      DaoFactory daoFactory,
-      ActionFactory actionFactory,
-      ChangeStreamMetrics metrics) {
-    this.heartbeatDuration = heartbeatDuration;
+      DaoFactory daoFactory, ActionFactory actionFactory, ChangeStreamMetrics metrics) {
     this.daoFactory = daoFactory;
     this.metrics = metrics;
     this.actionFactory = actionFactory;
@@ -171,7 +166,7 @@ public void setup() throws IOException {
             changeStreamDao,
             metrics,
             changeStreamAction,
-            heartbeatDuration,
+            HEARTBEAT_DURATION,
             sizeEstimator);
   }
 
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java
index 2e340c2b30b5..b2c9d6ac1318 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ChangeStreamActionTest.java
@@ -103,8 +103,7 @@ public void testHeartBeat() {
             tracker,
             receiver,
             watermarkEstimator,
-            throughputEstimator,
-            false);
+            throughputEstimator);
 
     assertFalse(result.isPresent());
     verify(metrics).incHeartbeatCount();
@@ -152,8 +151,7 @@ public void testCloseStreamResume() {
             tracker,
             receiver,
             watermarkEstimator,
-            throughputEstimator,
-            false);
+            throughputEstimator);
 
     assertTrue(result.isPresent());
     assertEquals(DoFn.ProcessContinuation.resume(), result.get());
@@ -187,8 +185,7 @@ public void testChangeStreamMutationUser() {
             tracker,
             receiver,
             watermarkEstimator,
-            throughputEstimator,
-            false);
+            throughputEstimator);
 
     assertFalse(result.isPresent());
     verify(metrics).incChangeStreamMutationUserCounter();
@@ -243,8 +240,7 @@ public void testChangeStreamMutationGc() {
             tracker,
             receiver,
             watermarkEstimator,
-            throughputEstimator,
-            false);
+            throughputEstimator);
 
     assertFalse(result.isPresent());
     verify(metrics).incChangeStreamMutationGcCounter();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java
index 305be1568851..43419b7147f6 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/action/ReadChangeStreamPartitionActionTest.java
@@ -20,7 +20,6 @@
 import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.ChangeStreamContinuationTokenHelper.getTokenWithCorrectPartition;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -128,10 +127,10 @@ public void testLockingRowSucceed() throws IOException {
     Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class);
     when(responseIterator.next()).thenReturn(mockHeartBeat);
     when(responseIterator.hasNext()).thenReturn(true);
-    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any()))
         .thenReturn(responses);
 
-    when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamAction.run(any(), any(), any(), any(), any(), any()))
         .thenReturn(Optional.of(DoFn.ProcessContinuation.stop()));
 
     final DoFn.ProcessContinuation result =
@@ -140,7 +139,7 @@ public void testLockingRowSucceed() throws IOException {
     assertEquals(DoFn.ProcessContinuation.stop(), result);
     // Verify that on successful lock, we don't tryClaim on the tracker.
     verify(tracker, never()).tryClaim(any());
-    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean());
+    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any());
   }
 
   @Test
@@ -157,7 +156,7 @@ public void testLockingRowFailsStops() throws IOException {
     StreamProgress streamProgress = new StreamProgress();
     streamProgress.setFailToLock(true);
     verify(tracker).tryClaim(streamProgress);
-    verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean());
+    verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any());
   }
 
   @Test
@@ -174,10 +173,10 @@ public void testLockingRowNotNeededAfterFirstRun() throws IOException {
     Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class);
     when(responseIterator.next()).thenReturn(mockHeartBeat);
     when(responseIterator.hasNext()).thenReturn(true);
-    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any()))
         .thenReturn(responses);
 
-    when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamAction.run(any(), any(), any(), any(), any(), any()))
         .thenReturn(Optional.of(DoFn.ProcessContinuation.stop()));
 
     final DoFn.ProcessContinuation result =
@@ -186,7 +185,7 @@ public void testLockingRowNotNeededAfterFirstRun() throws IOException {
     assertEquals(DoFn.ProcessContinuation.stop(), result);
     // Verify that on successful lock, we don't tryClaim on the tracker.
     verify(tracker, never()).tryClaim(any());
-    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean());
+    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any());
   }
 
   @Test
@@ -205,7 +204,7 @@ public void testLockingRowNotNeededAfterFirstRunNotSame() throws IOException {
     StreamProgress streamProgress = new StreamProgress();
     streamProgress.setFailToLock(true);
     verify(tracker).tryClaim(streamProgress);
-    verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any(), anyBoolean());
+    verify(changeStreamAction, never()).run(any(), any(), any(), any(), any(), any());
   }
 
   @Test
@@ -221,16 +220,16 @@ public void testThatChangeStreamWorkerCounterIsIncrementedOnInitialRun() throws
     Heartbeat mockHeartBeat = Mockito.mock(Heartbeat.class);
     when(responseIterator.next()).thenReturn(mockHeartBeat);
     when(responseIterator.hasNext()).thenReturn(true);
-    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any()))
         .thenReturn(responses);
 
-    when(changeStreamAction.run(any(), any(), any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamAction.run(any(), any(), any(), any(), any(), any()))
         .thenReturn(Optional.of(DoFn.ProcessContinuation.stop()));
 
     final DoFn.ProcessContinuation result =
         action.run(partitionRecord, tracker, receiver, watermarkEstimator);
     assertEquals(DoFn.ProcessContinuation.stop(), result);
-    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any(), anyBoolean());
+    verify(changeStreamAction).run(any(), any(), any(), any(), any(), any());
   }
 
   @Test
@@ -247,7 +246,7 @@ public void testCloseStreamTerminateOKStatus() throws IOException {
         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(), anyBoolean());
+    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.
@@ -273,7 +272,7 @@ public void testCloseStreamTerminateNotOutOfRangeStatus() throws IOException {
         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(), anyBoolean());
+    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.
@@ -308,7 +307,7 @@ public void testCloseStreamWritesContinuationTokens() throws IOException {
         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(), anyBoolean());
+    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.
@@ -360,7 +359,7 @@ public void testCloseStreamNewPartitionMerge() throws IOException {
         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(), anyBoolean());
+    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.
@@ -406,7 +405,7 @@ public void testCloseStreamMergeWithoutNewPartitionsField() throws IOException {
         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(), anyBoolean());
+    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.
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java
index c82e3eac2999..867117b4d392 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dao/MetadataTableDaoTest.java
@@ -23,10 +23,14 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.api.core.ApiFuture;
 import com.google.cloud.bigtable.admin.v2.BigtableTableAdminClient;
 import com.google.cloud.bigtable.admin.v2.BigtableTableAdminSettings;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
@@ -45,6 +49,9 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.UniqueIdGenerator;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.encoder.MetadataTableEncoder;
 import org.apache.beam.sdk.io.gcp.bigtable.changestreams.model.NewPartition;
@@ -743,4 +750,28 @@ public void testLockPartitionRecordsMetadata() {
             .get(0)
             .getValue());
   }
+
+  @Test
+  public void mutateRowWithHardTimeoutErrorHandling()
+      throws ExecutionException, InterruptedException, TimeoutException {
+    BigtableDataClient mockClient = Mockito.mock(BigtableDataClient.class);
+    MetadataTableDao daoWithMock =
+        new MetadataTableDao(mockClient, "test-table", ByteString.copyFromUtf8("test"));
+    ApiFuture<Void> mockFuture = mock(ApiFuture.class);
+    when(mockClient.mutateRowAsync(any())).thenReturn(mockFuture);
+
+    when(mockFuture.get(40, TimeUnit.SECONDS))
+        .thenThrow(TimeoutException.class)
+        .thenThrow(InterruptedException.class)
+        .thenThrow(ExecutionException.class);
+    assertThrows(
+        RuntimeException.class,
+        () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow()));
+    assertThrows(
+        RuntimeException.class,
+        () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow()));
+    assertThrows(
+        RuntimeException.class,
+        () -> daoWithMock.mutateRowWithHardTimeout(RowMutation.create("test", "test").deleteRow()));
+  }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java
index a9fa54d4c493..f0337504282f 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/InitializeDoFnTest.java
@@ -20,7 +20,6 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -88,8 +87,6 @@ public void setUp() throws IOException {
     metadataTableAdminDao =
         spy(new MetadataTableAdminDao(adminClient, null, changeStreamName, tableId));
     metadataTableAdminDao.createMetadataTable();
-    doReturn(true).when(metadataTableAdminDao).isAppProfileSingleClusterAndTransactional(any());
-    when(daoFactory.getMetadataTableAdminDao()).thenReturn(metadataTableAdminDao);
     metadataTableDao =
         new MetadataTableDao(
             dataClient, tableId, metadataTableAdminDao.getChangeStreamNamePrefix());
@@ -102,10 +99,7 @@ public void testInitializeDefault() throws IOException {
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
         new InitializeDoFn(
-            daoFactory,
-            "app-profile",
-            startTime,
-            BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
+            daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
     initializeDoFn.processElement(outputReceiver);
     verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
   }
@@ -116,10 +110,7 @@ public void testInitializeStopWithExistingPipeline() throws IOException {
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
         new InitializeDoFn(
-            daoFactory,
-            "app-profile",
-            startTime,
-            BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
+            daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
     initializeDoFn.processElement(outputReceiver);
     verify(outputReceiver, never()).output(any());
   }
@@ -139,51 +130,7 @@ public void testInitializeStopWithoutDNP() throws IOException {
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
         new InitializeDoFn(
-            daoFactory,
-            "app-profile",
-            startTime,
-            BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
-    initializeDoFn.processElement(outputReceiver);
-    verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
-    assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix()));
-  }
-
-  @Test
-  public void testInitializeNewWithoutDNP() throws IOException {
-    // DNP row doesn't exist, so we output with provided start time and we clean up any existing
-    // rows under the prefix.
-    dataClient.mutateRow(
-        RowMutation.create(
-                tableId,
-                metadataTableAdminDao
-                    .getChangeStreamNamePrefix()
-                    .concat(ByteString.copyFromUtf8("existing_row")))
-            .setCell(
-                MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
-    Instant startTime = Instant.now();
-    InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.NEW);
-    initializeDoFn.processElement(outputReceiver);
-    verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
-    assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix()));
-  }
-
-  @Test
-  public void testInitializeNewWithDNP() throws IOException {
-    metadataTableDao.updateDetectNewPartitionWatermark(Instant.now());
-    dataClient.mutateRow(
-        RowMutation.create(
-                tableId,
-                metadataTableAdminDao
-                    .getChangeStreamNamePrefix()
-                    .concat(ByteString.copyFromUtf8("existing_row")))
-            .setCell(
-                MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
-    Instant startTime = Instant.now();
-    InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.NEW);
+            daoFactory, startTime, BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS);
     initializeDoFn.processElement(outputReceiver);
     verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
     assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix()));
@@ -201,8 +148,7 @@ public void testInitializeResumeWithoutDNP() throws IOException {
                 MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW);
+        new InitializeDoFn(daoFactory, startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW);
     initializeDoFn.processElement(outputReceiver);
     // We want to resume but there's no DNP row, so we resume from the startTime provided.
     verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
@@ -222,8 +168,7 @@ public void testInitializeResumeWithDNP() throws IOException {
                 MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW);
+        new InitializeDoFn(daoFactory, startTime, BigtableIO.ExistingPipelineOptions.RESUME_OR_NEW);
     initializeDoFn.processElement(outputReceiver);
     verify(outputReceiver, times(1)).output(new InitialPipelineState(resumeTime, true));
     assertNull(dataClient.readRow(tableId, metadataTableAdminDao.getChangeStreamNamePrefix()));
@@ -241,8 +186,7 @@ public void testInitializeSkipCleanupWithoutDNP() throws IOException {
                 MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, ExistingPipelineOptions.SKIP_CLEANUP);
+        new InitializeDoFn(daoFactory, startTime, ExistingPipelineOptions.SKIP_CLEANUP);
     initializeDoFn.processElement(outputReceiver);
     // Skip cleanup will always resume from startTime
     verify(outputReceiver, times(1)).output(new InitialPipelineState(startTime, false));
@@ -264,8 +208,7 @@ public void testInitializeSkipCleanupWithDNP() throws IOException {
                 MetadataTableAdminDao.CF_WATERMARK, MetadataTableAdminDao.QUALIFIER_DEFAULT, 123));
     Instant startTime = Instant.now();
     InitializeDoFn initializeDoFn =
-        new InitializeDoFn(
-            daoFactory, "app-profile", startTime, ExistingPipelineOptions.SKIP_CLEANUP);
+        new InitializeDoFn(daoFactory, startTime, ExistingPipelineOptions.SKIP_CLEANUP);
     initializeDoFn.processElement(outputReceiver);
     // We don't want the pipeline to resume to avoid duplicates
     verify(outputReceiver, never()).output(any());
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java
index cb82fc6d19ca..b89b2bf15aa3 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigtable/changestreams/dofn/ReadChangeStreamPartitionDoFnTest.java
@@ -20,7 +20,6 @@
 import static org.apache.beam.sdk.io.gcp.bigtable.changestreams.TimestampConverter.toThreetenInstant;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -94,7 +93,7 @@ public void setup() throws IOException {
             sizeEstimator))
         .thenReturn(readChangeStreamPartitionAction);
 
-    doFn = new ReadChangeStreamPartitionDoFn(heartbeatDuration, daoFactory, actionFactory, metrics);
+    doFn = new ReadChangeStreamPartitionDoFn(daoFactory, actionFactory, metrics);
     doFn.setSizeEstimator(sizeEstimator);
   }
 
@@ -136,7 +135,7 @@ public void testProcessElementAndGetSize() throws IOException, InterruptedExcept
     when(mockResponses.hasNext()).thenReturn(true, true, true);
     when(mockResponses.next()).thenReturn(mockMutation, mockMutation, mockMutation);
     when(mockStream.iterator()).thenReturn(mockResponses);
-    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any(), anyBoolean()))
+    when(changeStreamDao.readChangeStreamPartition(any(), any(), any(), any()))
         .thenReturn(mockStream);
 
     when(watermarkEstimator.getState()).thenReturn(tenSecondsAgo);