fix: fix retry so it won't fail when rows read == rows limit (#2925)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/java-bigtable-hbase/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> ☕️

Author: mutianf

bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/ReadRowsRequestManager.java

@@ -52,10 +52,6 @@ class ReadRowsRequestManager {
     this.originalRequest = originalRequest;
   }
 
-  public ByteString getLastFoundKey() {
-    return lastFoundKey;
-  }
-
   void updateLastFoundKey(ByteString key) {
     this.lastFoundKey = key;
   }
@@ -64,6 +60,21 @@ void incrementRowCount(int count) {
     rowCount += count;
   }
 
+  boolean isConsumed() {
+    // A non-null lastFoundKey implies that we had at least one row that was returned and the
+    // filtered request should be truncated. In this case an empty filtered row set implies that all
+    // the rows have been seen and this stream is complete.
+    if (this.lastFoundKey != null && filterRows().equals(RowSet.getDefaultInstance())) {
+      return true;
+    }
+    // If the request has a rowsLimit and rowCount has reached the limit, it implies that the
+    // stream is complete.
+    if (originalRequest.getRowsLimit() > 0 && originalRequest.getRowsLimit() == rowCount) {
+      return true;
+    }
+    return false;
+  }
+
   ReadRowsRequest buildUpdatedRequest() {
     ReadRowsRequest.Builder newRequest =
         ReadRowsRequest.newBuilder()

bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java

@@ -20,7 +20,6 @@
 import com.google.api.core.InternalApi;
 import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.bigtable.v2.ReadRowsResponse;
-import com.google.bigtable.v2.RowSet;
 import com.google.cloud.bigtable.config.RetryOptions;
 import com.google.cloud.bigtable.grpc.DeadlineGenerator;
 import com.google.cloud.bigtable.grpc.async.AbstractRetryingOperation;
@@ -204,13 +203,11 @@ public void onClose(Status status, Metadata trailers) {
     // scan. To mitigate this, we will just mask the status as an ok after we are certain that we
     // have received all of the data.
     // This mitigation must only be activated after at least one row is received so that we can
-    // distinguish from a perfectly valid initial full table scan.
-    if (!status.isOk() && requestManager.getLastFoundKey() != null) {
-      ReadRowsRequest retryRequest = requestManager.buildUpdatedRequest();
-      boolean isFullTableScan = retryRequest.getRows().equals(RowSet.getDefaultInstance());
-      if (isFullTableScan) {
-        status = Status.OK;
-      }
+    // distinguish from a perfectly valid initial full table scan. And we verify the ReadRows stream
+    // is done by checking if all the keys in the request are read, or the number of rows returned
+    // reached the rowsLimit.
+    if (!status.isOk() && requestManager.isConsumed()) {
+      status = Status.OK;
     }
 
     if (status.getCause() instanceof StreamWaitTimeoutException) {

bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java

@@ -449,6 +449,41 @@ public void testErrorAfterComplete() throws UnsupportedEncodingException {
     Assert.assertTrue(underTest.getRowMerger().isComplete());
   }
 
+  @Test
+  public void testErrorAfterCompleteWithRowsLimit() throws UnsupportedEncodingException {
+    ByteString key1 = ByteString.copyFromUtf8("SomeKey1");
+    ByteString key2 = ByteString.copyFromUtf8("SomeKey2");
+
+    ReadRowsRequest req =
+        ReadRowsRequest.newBuilder()
+            .setRows(RowSet.newBuilder().addRowKeys(key1).addRowKeys(key2))
+            .setRowsLimit(1)
+            .build();
+    RetryingReadRowsOperation underTest =
+        createOperation(DeadlineGenerator.DEFAULT, req, mockFlatRowObserver);
+
+    start(underTest);
+    // Not all the rows are read yet, but because the rows limit has reached, this will still return
+    // an OK status
+    underTest.onMessage(buildResponse(key1));
+    underTest.onClose(Status.CANCELLED, new Metadata());
+
+    verify(mockFlatRowObserver, times(1)).onCompleted();
+    Assert.assertFalse(underTest.inRetryMode());
+    Assert.assertTrue(underTest.getRowMerger().isComplete());
+  }
+
+  @Test
+  public void testFullTableScanRetried() {
+    ReadRowsRequest req = ReadRowsRequest.newBuilder().build();
+    RetryingReadRowsOperation underTest =
+        createOperation(DeadlineGenerator.DEFAULT, req, mockFlatRowObserver);
+
+    start(underTest);
+    underTest.onClose(Status.DEADLINE_EXCEEDED, new Metadata());
+    Assert.assertTrue(underTest.inRetryMode());
+  }
+
   @SuppressWarnings("unchecked")
   @Test
   public void testImmediateOnClose() {