feat: Implement ReadChangeStreamResumptionStrategy

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -82,6 +82,7 @@
 import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
 import com.google.cloud.bigtable.data.v2.stub.changestream.ChangeStreamRecordMergingCallable;
 import com.google.cloud.bigtable.data.v2.stub.changestream.ListChangeStreamPartitionsUserCallable;
+import com.google.cloud.bigtable.data.v2.stub.changestream.ReadChangeStreamResumptionStrategy;
 import com.google.cloud.bigtable.data.v2.stub.changestream.ReadChangeStreamUserCallable;
 import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerStreamingCallable;
 import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerUnaryCallable;
@@ -898,7 +899,7 @@ public Map<String, String> extract(
    *       case of mutations, it will merge the {@link ReadChangeStreamResponse.DataChange}s into
    *       {@link ChangeStreamMutation}. The actual change stream record implementation can be
    *       configured by the {@code changeStreamRecordAdapter} parameter.
-   *   <li>TODO: Retry/resume on failure.
+   *   <li>Retry/resume on failure.
    *   <li>Add tracing & metrics.
    * </ul>
    */
@@ -939,7 +940,8 @@ public Map<String, String> extract(
     // Copy idle timeout settings for watchdog.
     ServerStreamingCallSettings<ReadChangeStreamRequest, ChangeStreamRecordT> innerSettings =
         ServerStreamingCallSettings.<ReadChangeStreamRequest, ChangeStreamRecordT>newBuilder()
-            // TODO: setResumptionStrategy.
+            .setResumptionStrategy(
+                new ReadChangeStreamResumptionStrategy<>(changeStreamRecordAdapter))
             .setRetryableCodes(settings.readChangeStreamSettings().getRetryableCodes())
             .setRetrySettings(settings.readChangeStreamSettings().getRetrySettings())
             .setIdleTimeout(settings.readChangeStreamSettings().getIdleTimeout())
@@ -951,8 +953,6 @@ public Map<String, String> extract(
     ServerStreamingCallable<ReadChangeStreamRequest, ChangeStreamRecordT> withBigtableTracer =
         new BigtableTracerStreamingCallable<>(watched);
 
-    // TODO: Add ReadChangeStreamRetryCompletedCallable.
-
     ServerStreamingCallable<ReadChangeStreamRequest, ChangeStreamRecordT> readChangeStreamCallable =
         Callables.retrying(withBigtableTracer, innerSettings, clientContext);
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamResumptionStrategy.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.data.v2.stub.changestream;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.retrying.StreamResumptionStrategy;
+import com.google.bigtable.v2.ReadChangeStreamRequest;
+import com.google.bigtable.v2.ReadChangeStreamRequest.Builder;
+import com.google.bigtable.v2.StreamContinuationToken;
+import com.google.bigtable.v2.StreamContinuationTokens;
+import com.google.cloud.bigtable.data.v2.models.ChangeStreamRecordAdapter;
+
+/**
+ * An implementation of a {@link StreamResumptionStrategy} for change stream records. This class
+ * tracks the continuation token and upon retry can build a request to resume the stream from where
+ * it left off.
+ *
+ * <p>This class is considered an internal implementation detail and not meant to be used by
+ * applications.
+ */
+@InternalApi
+public class ReadChangeStreamResumptionStrategy<ChangeStreamRecordT>
+    implements StreamResumptionStrategy<ReadChangeStreamRequest, ChangeStreamRecordT> {
+  private final ChangeStreamRecordAdapter<ChangeStreamRecordT> changeStreamRecordAdapter;
+  private String token = null;
+
+  public ReadChangeStreamResumptionStrategy(
+      ChangeStreamRecordAdapter<ChangeStreamRecordT> changeStreamRecordAdapter) {
+    this.changeStreamRecordAdapter = changeStreamRecordAdapter;
+  }
+
+  @Override
+  public boolean canResume() {
+    return true;
+  }
+
+  @Override
+  public StreamResumptionStrategy<ReadChangeStreamRequest, ChangeStreamRecordT> createNew() {
+    return new ReadChangeStreamResumptionStrategy<>(changeStreamRecordAdapter);
+  }
+
+  @Override
+  public ChangeStreamRecordT processResponse(ChangeStreamRecordT response) {
+    // Update the token from a Heartbeat or a ChangeStreamMutation.
+    // We don't worry about resumption after CloseStream, since the server
+    // will return an OK status right after sending a CloseStream.
+    if (changeStreamRecordAdapter.isHeartbeat(response)) {
+      this.token = changeStreamRecordAdapter.getTokenFromHeartbeat(response);
+    }
+    if (changeStreamRecordAdapter.isChangeStreamMutation(response)) {
+      this.token = changeStreamRecordAdapter.getTokenFromChangeStreamMutation(response);
+    }
+    return response;
+  }
+
+  /**
+   * {@inheritDoc}
+   *
+   * <p>Given a request, this implementation will narrow that request to include data changes that
+   * come after {@link #token}.
+   */
+  @Override
+  public ReadChangeStreamRequest getResumeRequest(ReadChangeStreamRequest originalRequest) {
+    // A null token means that we have not successfully read a Heartbeat nor a ChangeStreamMutation,
+    // so start from the beginning.
+    if (this.token == null) {
+      return originalRequest;
+    }
+
+    Builder builder = originalRequest.toBuilder();
+    // We need to clear the start_from and use the updated continuation_tokens
+    // to resume the request.
+    // The partition should always be the same as the one from the original request,
+    // otherwise we would receive a CloseStream with different
+    // partitions(which indicates tablet split/merge events).
+    builder.clearStartFrom();
+    builder.setContinuationTokens(
+        StreamContinuationTokens.newBuilder()
+            .addTokens(
+                StreamContinuationToken.newBuilder()
+                    .setPartition(originalRequest.getPartition())
+                    .setToken(this.token)
+                    .build())
+            .build());
+
+    return builder.build();
+  }
+}