fix: retry rst stream (port of #3000) (#3211)

kolea2 · web-flow · commit 9f6388228d94 · 2021-09-10T18:50:30.000Z
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AbstractRetryingOperation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/AbstractRetryingOperation.java
@@ -203,7 +203,7 @@ protected void onError(Status status, Metadata trailers) {
     String channelId = ChannelPool.extractIdentifier(trailers);
     // Non retry scenario
     if (!retryOptions.enableRetries()
-        || !retryOptions.isRetryable(code)
+        || !isStatusRetryable(status)
         // Unauthenticated is special because the request never made it to
         // to the server, so all requests are retryable
         || !(isRequestRetryable() || code == Code.UNAUTHENTICATED || code == Code.UNAVAILABLE)) {
@@ -263,6 +263,10 @@ protected boolean isRequestRetryable() {
     return rpc.isRetryable(getRetryRequest());
   }
 
+  protected boolean isStatusRetryable(Status status) {
+    return retryOptions.isRetryable(status.getCode());
+  }
+
   protected void setException(Exception exception) {
     completionFuture.setException(exception);
   }
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperation.java
@@ -31,6 +31,7 @@
 import com.google.protobuf.ByteString;
 import io.grpc.Metadata;
 import io.grpc.Status;
+import io.grpc.Status.Code;
 import io.grpc.stub.ClientResponseObserver;
 import io.grpc.stub.StreamObserver;
 import io.opencensus.trace.AttributeValue;
@@ -243,6 +244,24 @@ protected boolean isRequestRetryable() {
     return true;
   }
 
+  /** Read rows requests are retryable if the status is a rst stream error. */
+  @Override
+  protected boolean isStatusRetryable(Status status) {
+    return retryOptions.isRetryable(status.getCode()) || isRstStream(status);
+  }
+
+  private boolean isRstStream(Status status) {
+    if (status.getCode() == Code.INTERNAL) {
+      String description = status.getDescription();
+      if (description != null) {
+        return description.contains("Received Rst stream")
+            || description.contains("RST_STREAM closed stream")
+            || description.contains("Received RST_STREAM");
+      }
+    }
+    return false;
+  }
+
   /** {@inheritDoc} */
   @Override
   protected boolean onOK(Metadata trailers) {
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/scanner/RetryingReadRowsOperationTest.java
@@ -513,6 +513,26 @@ public Void answer(InvocationOnMock invocation) throws Throwable {
     Assert.assertTrue(underTest.getRowMerger().isComplete());
   }
 
+  @Test
+  public void testRetryRstStream() throws Exception {
+    RetryingReadRowsOperation underTest = createOperation();
+    start(underTest);
+
+    ByteString key1 = ByteString.copyFrom("SomeKey1", "UTF-8");
+    ByteString key2 = ByteString.copyFrom("SomeKey2", "UTF-8");
+    underTest.onMessage(buildResponse(key1));
+    underTest.onClose(
+        Status.INTERNAL.withDescription("HTTP/2 error code: INTERNAL_ERROR\nReceived Rst stream"),
+        null);
+    Assert.assertFalse(underTest.getRowMerger().isComplete());
+    underTest.onMessage(buildResponse(key2));
+    verify(mockFlatRowObserver, times(2)).onNext(any(FlatRow.class));
+    checkRetryRequest(underTest, key2, 8);
+    verify(mockClientCall, times(4)).request(eq(1));
+
+    finishOK(underTest, 1);
+  }
+
   protected void performTimeout(RetryingReadRowsOperation underTest) {
     underTest.onClose(
         Status.CANCELLED.withCause(
diff --git a/pom.xml b/pom.xml
@@ -56,7 +56,7 @@ limitations under the License.
     <compileSource.1.8>1.8</compileSource.1.8>
 
     <!-- core dependency versions -->
-    <bigtable.version>1.22.0-sp.1</bigtable.version>
+    <bigtable.version>1.22.0-sp.2</bigtable.version>
     <dropwizard.metrics.version>3.2.6</dropwizard.metrics.version>
     <!--  keeping at this version to align with hbase-->
     <slf4j.version>1.7.25</slf4j.version>
