fix: retry on RST_STREAM internal error

Author: rahul2393

google-cloud-spanner/src/main/java/com/google/cloud/spanner/IsRetryableInternalError.java

@@ -29,6 +29,8 @@ public class IsRetryableInternalError implements Predicate<Throwable> {
   private static final String EOS_ERROR_MESSAGE =
       "Received unexpected EOS on DATA frame from server";
 
+  private static final String RST_STREAM_ERROR_MESSAGE = "stream terminated by RST_STREAM";
+
   @Override
   public boolean apply(Throwable cause) {
     if (isInternalError(cause)) {
@@ -38,6 +40,8 @@ public boolean apply(Throwable cause) {
         return true;
       } else if (cause.getMessage().contains(EOS_ERROR_MESSAGE)) {
         return true;
+      } else if (cause.getMessage().contains(RST_STREAM_ERROR_MESSAGE)) {
+        return true;
       }
     }
     return false;

google-cloud-spanner/src/test/java/com/google/cloud/spanner/IsRetryableInternalErrorTest.java

@@ -114,6 +114,18 @@ public void genericInternalStatusRuntimeExceptionIsRetryable() {
     assertThat(predicate.apply(e)).isFalse();
   }
 
+  @Test
+  public void rstStreamInternalExceptionIsRetryable() {
+    final InternalException e =
+        new InternalException(
+            "INTERNAL: stream terminated by RST_STREAM.",
+            null,
+            GrpcStatusCode.of(Code.INTERNAL),
+            false);
+
+    assertThat(predicate.apply(e)).isTrue();
+  }
+
   @Test
   public void genericInternalExceptionIsNotRetryable() {
     final InternalException e =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java

@@ -344,6 +344,45 @@ public void testExecuteStreamingPartitionedUpdateUnexpectedEOS() {
             Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class));
   }
 
+  @Test
+  public void testExecuteStreamingPartitionedUpdateRSTstream() {
+    ResultSetStats stats = ResultSetStats.newBuilder().setRowCountLowerBound(1000L).build();
+    PartialResultSet p1 = PartialResultSet.newBuilder().setResumeToken(resumeToken).build();
+    PartialResultSet p2 = PartialResultSet.newBuilder().setStats(stats).build();
+    ServerStream<PartialResultSet> stream1 = mock(ServerStream.class);
+    Iterator<PartialResultSet> iterator = mock(Iterator.class);
+    when(iterator.hasNext()).thenReturn(true, true, false);
+    when(iterator.next())
+        .thenReturn(p1)
+        .thenThrow(
+            new InternalException(
+                "INTERNAL: stream terminated by RST_STREAM.",
+                null,
+                GrpcStatusCode.of(Code.INTERNAL),
+                true));
+    when(stream1.iterator()).thenReturn(iterator);
+    ServerStream<PartialResultSet> stream2 = mock(ServerStream.class);
+    when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator());
+    when(rpc.executeStreamingPartitionedDml(
+            Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)))
+        .thenReturn(stream1);
+    when(rpc.executeStreamingPartitionedDml(
+            Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)))
+        .thenReturn(stream2);
+
+    PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker);
+    long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10));
+
+    assertThat(count).isEqualTo(1000L);
+    verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap());
+    verify(rpc)
+        .executeStreamingPartitionedDml(
+            Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class));
+    verify(rpc)
+        .executeStreamingPartitionedDml(
+            Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class));
+  }
+
   @Test
   public void testExecuteStreamingPartitionedUpdateGenericInternalException() {
     PartialResultSet p1 = PartialResultSet.newBuilder().setResumeToken(resumeToken).build();