chore: extend resumption strategy (#2489)

Author: mutianf

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2025 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;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.retrying.StreamResumptionStrategy;
+
+@InternalApi
+/** Expand StreamResumptionStrategy to also process the error. */
+public abstract class BigtableStreamResumptionStrategy<RequestT, ResponseT>
+    implements StreamResumptionStrategy<RequestT, ResponseT> {
+
+  public abstract Throwable processError(Throwable throwable);
+}

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

@@ -22,20 +22,21 @@
 import com.google.bigtable.v2.RowSet;
 import com.google.cloud.bigtable.data.v2.internal.RowSetUtil;
 import com.google.cloud.bigtable.data.v2.models.RowAdapter;
+import com.google.cloud.bigtable.data.v2.stub.BigtableStreamResumptionStrategy;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.ByteString;
 
 /**
- * An implementation of a {@link StreamResumptionStrategy} for merged rows. This class tracks the
- * last complete row seen and upon retry can build a request to resume the stream from where it left
- * off.
+ * An implementation of a {@link BigtableStreamResumptionStrategy} for merged rows. This class
+ * tracks the last complete row seen 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 ReadRowsResumptionStrategy<RowT>
-    implements StreamResumptionStrategy<ReadRowsRequest, RowT> {
+    extends BigtableStreamResumptionStrategy<ReadRowsRequest, RowT> {
   private final RowAdapter<RowT> rowAdapter;
   private ByteString lastKey = ByteString.EMPTY;
   // Number of rows processed excluding Marker row.
@@ -69,6 +70,12 @@ public RowT processResponse(RowT response) {
     return response;
   }
 
+  @Override
+  public Throwable processError(Throwable throwable) {
+    // Noop
+    return throwable;
+  }
+
   /**
    * {@inheritDoc}
    *

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ServerStreamingAttemptCallable.java

@@ -25,6 +25,7 @@
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StateCheckingResponseObserver;
 import com.google.api.gax.rpc.StreamController;
+import com.google.cloud.bigtable.data.v2.stub.BigtableStreamResumptionStrategy;
 import com.google.common.base.Preconditions;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CancellationException;
@@ -345,6 +346,10 @@ private void onAttemptError(Throwable throwable) {
       localCancellationCause = cancellationCause;
     }
 
+    if (resumptionStrategy instanceof BigtableStreamResumptionStrategy) {
+      throwable = ((BigtableStreamResumptionStrategy) resumptionStrategy).processError(throwable);
+    }
+
     if (localCancellationCause != null) {
       // Take special care to preserve the cancellation's stack trace.
       innerAttemptFuture.setException(localCancellationCause);