feat: Implement ReadChangeStreamResumptionStrategy

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/ChangeStreamRecordAdapter.java

@@ -36,6 +36,10 @@ public interface ChangeStreamRecordAdapter<ChangeStreamRecordT> {
   @InternalApi("Used in Changestream beam pipeline.")
   boolean isHeartbeat(ChangeStreamRecordT record);
 
+  /** Checks if the given change stream record is a CloseStream. */
+  @InternalApi("Used in Changestream beam pipeline.")
+  boolean isCloseStream(ChangeStreamRecordT record);
+
   /**
    * Get the token from the given Heartbeat record. If the given record is not a Heartbeat, it will
    * throw an Exception.

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapter.java

@@ -48,6 +48,12 @@ public String getTokenFromHeartbeat(ChangeStreamRecord record) {
     return ((Heartbeat) record).getChangeStreamContinuationToken().getToken();
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public boolean isCloseStream(ChangeStreamRecord record) {
+    return record instanceof CloseStream;
+  }
+
   /** {@inheritDoc} */
   @Override
   public boolean isChangeStreamMutation(ChangeStreamRecord record) {

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

@@ -36,6 +36,7 @@ public class ReadChangeStreamResumptionStrategy<ChangeStreamRecordT>
     implements StreamResumptionStrategy<ReadChangeStreamRequest, ChangeStreamRecordT> {
   private final ChangeStreamRecordAdapter<ChangeStreamRecordT> changeStreamRecordAdapter;
   private String token = null;
+  private boolean canResume = true;
 
   public ReadChangeStreamResumptionStrategy(
       ChangeStreamRecordAdapter<ChangeStreamRecordT> changeStreamRecordAdapter) {
@@ -44,7 +45,7 @@ public ReadChangeStreamResumptionStrategy(
 
   @Override
   public boolean canResume() {
-    return true;
+    return canResume;
   }
 
   @Override
@@ -55,9 +56,14 @@ public StreamResumptionStrategy<ReadChangeStreamRequest, ChangeStreamRecordT> cr
   @Override
   public ChangeStreamRecordT processResponse(ChangeStreamRecordT response) {
     // Update the token from a Heartbeat or a ChangeStreamMutation.
+    // If we get a CloseStream, disable resumption and don't re-enable it, since
+    // the stream is supposed to be closed upon CloseStream.
     if (changeStreamRecordAdapter.isHeartbeat(response)) {
       this.token = changeStreamRecordAdapter.getTokenFromHeartbeat(response);
     }
+    if (changeStreamRecordAdapter.isCloseStream(response)) {
+      canResume = false;
+    }
     if (changeStreamRecordAdapter.isChangeStreamMutation(response)) {
       this.token = changeStreamRecordAdapter.getTokenFromChangeStreamMutation(response);
     }
@@ -79,8 +85,11 @@ public ReadChangeStreamRequest getResumeRequest(ReadChangeStreamRequest original
     }
 
     Builder builder = originalRequest.toBuilder();
-    // We need to clear either start_time or continuation_tokens.
+    // We need to clear both start_time and continuation_tokens.
     // And just use the StreamPartition and the token to resume the request.
+    // The partition should is always the same as the one from the original request,
+    // because otherwise we would have received a CloseStream with different
+    // partitions(which indicates tablet split/merge events).
     builder.clearStartFrom();
     builder.setContinuationTokens(
         StreamContinuationTokens.newBuilder()

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/DefaultChangeStreamRecordAdapterTest.java

@@ -75,6 +75,23 @@ public void getTokenFromHeartbeatTest() {
     Assert.assertEquals(adapter.getTokenFromHeartbeat(heartbeatRecord), "heartbeat-token");
   }
 
+  @Test
+  public void isCloseStreamTest() {
+    ChangeStreamRecord heartbeatRecord =
+        Heartbeat.fromProto(ReadChangeStreamResponse.Heartbeat.getDefaultInstance());
+    ChangeStreamRecord closeStreamRecord =
+        CloseStream.fromProto(ReadChangeStreamResponse.CloseStream.getDefaultInstance());
+    ChangeStreamRecord changeStreamMutationRecord =
+        ChangeStreamMutation.createGcMutation(
+                ByteString.copyFromUtf8("key"), Timestamp.getDefaultInstance(), 0)
+            .setToken("token")
+            .setLowWatermark(Timestamp.getDefaultInstance())
+            .build();
+    Assert.assertFalse(adapter.isCloseStream(heartbeatRecord));
+    Assert.assertTrue(adapter.isCloseStream(closeStreamRecord));
+    Assert.assertFalse(adapter.isCloseStream(changeStreamMutationRecord));
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void getTokenFromHeartbeatInvalidTypeTest() {
     ChangeStreamRecord closeStreamRecord =

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/changestream/ReadChangeStreamRetryTest.java

@@ -22,6 +22,7 @@
 import com.google.api.gax.rpc.FixedTransportChannelProvider;
 import com.google.api.gax.rpc.InternalException;
 import com.google.api.gax.rpc.ServerStream;
+import com.google.api.gax.rpc.UnavailableException;
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.Mutation;
 import com.google.bigtable.v2.ReadChangeStreamRequest;
@@ -58,6 +59,7 @@
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -76,6 +78,7 @@ public class ReadChangeStreamRetryTest {
   @Rule public GrpcServerRule serverRule = new GrpcServerRule();
   private TestBigtableService service;
   private BigtableDataClient client;
+  @Rule public ExpectedException expect = ExpectedException.none();
 
   @Before
   public void setUp() throws IOException {
@@ -271,8 +274,8 @@ public void errorAfterHeartbeatShouldResumeWithTokenTest() {
     Assert.assertTrue(actualResults.get(0) instanceof Heartbeat);
   }
 
-  // [{ReadChangeStreamResponse.CloseStream}, {UNAVAILABLE}] -> Resume with original request.
-  @Test
+  // [{ReadChangeStreamResponse.CloseStream}, {UNAVAILABLE}] -> Request not resumed.
+  @Test(expected = UnavailableException.class)
   public void errorAfterSingleCloseStreamShouldResumeWithOriginalRequestTest() {
     // CloseStream.
     ReadChangeStreamResponse closeStreamResponse =
@@ -282,10 +285,10 @@ public void errorAfterSingleCloseStreamShouldResumeWithOriginalRequestTest() {
             .expectInitialRequest()
             .respondWith(closeStreamResponse)
             .respondWithStatus(Code.UNAVAILABLE));
-    // Resume the request with the original request.
-    // We don't care about the response here so just do expectRequest.
-    service.expectations.add(RpcExpectation.create().expectInitialRequest());
+    // Request is not resumed. We'll get a CloseStream and then an UNAVAILABLE error.
     List<ChangeStreamRecord> actualResults = getResults();
+    expect.expect(UnavailableException.class);
+    // This is the CloseStream we get before UNAVAILABLE.
     Assert.assertEquals(actualResults.size(), 1);
     Assert.assertTrue(actualResults.get(0) instanceof CloseStream);
   }