test: Add ResumableStreamIterator tests for RESOURCE_EXHAUSTED error code (#3039)

* feat: Add ResumableStreamIterator tests for RESOURCE_EXHAUSTED error
code.

* chore: simplify test a bit

* test: add test to Gax retryer test

Add test cases to Gax retryer test file to ensure that only
RESOURCE_EXHAUSTED errors with RetryInfo are actually retried end-to-end
by the Spanner client library and the Gax configuration.

---------

Co-authored-by: Knut Olav Løite <[email protected]>

Author: Vizerai

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

@@ -26,6 +26,7 @@
 import com.google.api.client.util.BackOff;
 import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
 import com.google.common.collect.AbstractIterator;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Duration;
@@ -52,11 +53,13 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
 import org.mockito.Mockito;
 
 /** Unit tests for {@link ResumableStreamIterator}. */
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public class ResumableStreamIteratorTest {
   interface Starter {
     AbstractResultSet.CloseableIterator<PartialResultSet> startStream(
@@ -69,6 +72,14 @@ interface ResultSetStream {
     void close();
   }
 
+  @Parameter(0)
+  public ErrorCode errorCodeParameter;
+
+  @Parameters(name = "errorCodeParameter = {0}")
+  public static List<ErrorCode> data() {
+    return ImmutableList.of(ErrorCode.UNAVAILABLE, ErrorCode.RESOURCE_EXHAUSTED);
+  }
+
   private static StatusRuntimeException statusWithRetryInfo(ErrorCode code) {
     Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
     Metadata trailers = new Metadata();
@@ -223,7 +234,7 @@ public void restart() {
     Mockito.when(s1.next())
         .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "a"))
         .thenReturn(resultSet(ByteString.copyFromUtf8("r2"), "b"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r2")))
@@ -244,7 +255,7 @@ public void restartWithHoldBack() {
         .thenReturn(resultSet(ByteString.copyFromUtf8("r2"), "b"))
         .thenReturn(resultSet(null, "X"))
         .thenReturn(resultSet(null, "X"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r2")))
@@ -265,7 +276,7 @@ public void restartWithHoldBackMidStream() {
         .thenReturn(resultSet(null, "b"))
         .thenReturn(resultSet(null, "c"))
         .thenReturn(resultSet(ByteString.copyFromUtf8("r2"), "d"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r2")))
@@ -360,7 +371,7 @@ public void bufferLimitRestart() {
     Mockito.when(s1.next())
         .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "a"))
         .thenReturn(resultSet(ByteString.copyFromUtf8("r2"), "b"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r2")))
@@ -380,7 +391,7 @@ public void bufferLimitRestartWithinLimitAtStartOfResults() {
     Mockito.when(starter.startStream(null)).thenReturn(new ResultSetIterator(s1));
     Mockito.when(s1.next())
         .thenReturn(resultSet(null, "XXXXXX"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(null)).thenReturn(new ResultSetIterator(s2));
@@ -400,7 +411,7 @@ public void bufferLimitRestartWithinLimitMidResults() {
     Mockito.when(s1.next())
         .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "a"))
         .thenReturn(resultSet(null, "XXXXXX"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r1")))
@@ -422,11 +433,11 @@ public void bufferLimitMissingTokensUnsafeToRetry() {
         .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "a"))
         .thenReturn(resultSet(null, "b"))
         .thenReturn(resultSet(null, "c"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     assertThat(consumeAtMost(3, resumableStreamIterator)).containsExactly("a", "b", "c").inOrder();
     SpannerException e = assertThrows(SpannerException.class, () -> resumableStreamIterator.next());
-    assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNAVAILABLE);
+    assertThat(e.getErrorCode()).isEqualTo(errorCodeParameter);
   }
 
   @Test
@@ -439,7 +450,7 @@ public void bufferLimitMissingTokensSafeToRetry() {
         .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "a"))
         .thenReturn(resultSet(null, "b"))
         .thenReturn(resultSet(ByteString.copyFromUtf8("r3"), "c"))
-        .thenThrow(new RetryableException(ErrorCode.UNAVAILABLE, "failed by test"));
+        .thenThrow(new RetryableException(errorCodeParameter, "failed by test"));
 
     ResultSetStream s2 = Mockito.mock(ResultSetStream.class);
     Mockito.when(starter.startStream(ByteString.copyFromUtf8("r3")))

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

@@ -21,22 +21,28 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import com.google.api.gax.grpc.testing.LocalChannelProvider;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.protobuf.ListValue;
+import com.google.rpc.RetryInfo;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.TypeCode;
+import io.grpc.Metadata;
 import io.grpc.Server;
+import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.protobuf.ProtoUtils;
 import java.io.IOException;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.After;
 import org.junit.AfterClass;
@@ -84,6 +90,14 @@ public class SpannerGaxRetryTest {
       SimulatedExecutionTime.ofMinimumAndRandomTime(1000, 0);
   private static final StatusRuntimeException UNAVAILABLE =
       io.grpc.Status.UNAVAILABLE.withDescription("Retryable test exception.").asRuntimeException();
+  private static final StatusRuntimeException RESOURCE_EXHAUSTED_NON_RETRYABLE =
+      Status.RESOURCE_EXHAUSTED
+          .withDescription("Non-retryable test exception.")
+          .asRuntimeException();
+  private static final StatusRuntimeException RESOURCE_EXHAUSTED_RETRYABLE =
+      Status.RESOURCE_EXHAUSTED
+          .withDescription("Retryable test exception.")
+          .asRuntimeException(createRetryInfo());
   private static final StatusRuntimeException FAILED_PRECONDITION =
       io.grpc.Status.FAILED_PRECONDITION
           .withDescription("Non-retryable test exception.")
@@ -192,6 +206,19 @@ public void tearDown() {
     spanner.close();
   }
 
+  static Metadata createRetryInfo() {
+    Metadata trailers = new Metadata();
+    RetryInfo retryInfo =
+        RetryInfo.newBuilder()
+            .setRetryDelay(
+                com.google.protobuf.Duration.newBuilder()
+                    .setNanos((int) TimeUnit.MILLISECONDS.toNanos(1L))
+                    .setSeconds(0L))
+            .build();
+    trailers.put(ProtoUtils.keyForProto(RetryInfo.getDefaultInstance()), retryInfo);
+    return trailers;
+  }
+
   private void warmUpSessionPool(DatabaseClient client) {
     for (int i = 0; i < 10; i++) {
       int retryCount = 0;
@@ -229,9 +256,29 @@ public void singleUseTimeout() {
 
   @Test
   public void singleUseUnavailable() {
-    mockSpanner.addException(UNAVAILABLE);
-    try (ResultSet rs = client.singleUse().executeQuery(SELECT1AND2)) {
-      while (rs.next()) {}
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(UNAVAILABLE));
+    try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1AND2)) {
+      assertTrue(resultSet.next());
+    }
+  }
+
+  @Test
+  public void singleUseResourceExhausted_nonRetryable() {
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(RESOURCE_EXHAUSTED_NON_RETRYABLE));
+    try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1AND2)) {
+      SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
+      assertEquals(ErrorCode.RESOURCE_EXHAUSTED, exception.getErrorCode());
+    }
+  }
+
+  @Test
+  public void singleUseResourceExhausted_retryable() {
+    mockSpanner.setExecuteStreamingSqlExecutionTime(
+        SimulatedExecutionTime.ofException(RESOURCE_EXHAUSTED_RETRYABLE));
+    try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1AND2)) {
+      assertTrue(resultSet.next());
     }
   }