Treat RESOURCE_EXHAUSTED with retry information as aborted.

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

@@ -16,11 +16,6 @@
 
 package com.google.cloud.spanner;
 
-import com.google.protobuf.util.Durations;
-import com.google.rpc.RetryInfo;
-import io.grpc.Metadata;
-import io.grpc.Status;
-import io.grpc.protobuf.ProtoUtils;
 import javax.annotation.Nullable;
 
 /**
@@ -29,8 +24,6 @@
  * other types of errors, most typically by retrying the transaction.
  */
 public class AbortedException extends SpannerException {
-  private static final Metadata.Key<RetryInfo> KEY_RETRY_INFO =
-      ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
 
   /**
    * Abort is not retryable per se: the operation request needs to change (generally to reflect a
@@ -43,21 +36,4 @@ public class AbortedException extends SpannerException {
       DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) {
     super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause);
   }
-
-  /**
-   * Return the retry delay for transaction in milliseconds. Return -1 if this does not specify any
-   * retry delay. In that case, clients should fall back to a locally computed retry delay.
-   */
-  public long getRetryDelayInMillis() {
-    if (this.getCause() != null) {
-      Metadata trailers = Status.trailersFromThrowable(this.getCause());
-      if (trailers != null && trailers.containsKey(KEY_RETRY_INFO)) {
-        RetryInfo retryInfo = trailers.get(KEY_RETRY_INFO);
-        if (retryInfo.hasRetryDelay()) {
-          return Durations.toMillis(retryInfo.getRetryDelay());
-        }
-      }
-    }
-    return -1L;
-  }
 }

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

@@ -18,11 +18,18 @@
 
 import com.google.cloud.grpc.BaseGrpcServiceException;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.util.Durations;
+import com.google.rpc.RetryInfo;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.protobuf.ProtoUtils;
 import javax.annotation.Nullable;
 
 /** Base exception type for all exceptions produced by the Cloud Spanner service. */
 public class SpannerException extends BaseGrpcServiceException {
   private static final long serialVersionUID = 20150916L;
+  private static final Metadata.Key<RetryInfo> KEY_RETRY_INFO =
+      ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
 
   private final ErrorCode code;
 
@@ -48,4 +55,26 @@ public ErrorCode getErrorCode() {
   enum DoNotConstructDirectly {
     ALLOWED
   }
+
+  /**
+   * Return the retry delay for operation in milliseconds. Return -1 if this does not specify any
+   * retry delay.
+   */
+  public long getRetryDelayInMillis() {
+    return extractRetryDelay(this.getCause());
+  }
+
+  static long extractRetryDelay(Throwable cause) {
+    if (cause != null) {
+      Metadata trailers = Status.trailersFromThrowable(cause);
+      if (trailers != null && trailers.containsKey(KEY_RETRY_INFO)) {
+        RetryInfo retryInfo = trailers.get(KEY_RETRY_INFO);
+        if (retryInfo.hasRetryDelay()) {
+          return Durations.toMillis(retryInfo.getRetryDelay());
+        }
+      }
+    }
+    return -1L;
+  }
+
 }

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

@@ -126,6 +126,8 @@ private static boolean isRetryable(ErrorCode code, @Nullable Throwable cause) {
         return hasCauseMatching(cause, Matchers.isRetryableInternalError);
       case UNAVAILABLE:
         return true;
+      case RESOURCE_EXHAUSTED:
+        return SpannerException.extractRetryDelay(cause) > 0;
       default:
         return false;
     }

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

@@ -231,7 +231,12 @@ static <T> T runWithRetries(Callable<T> callable) {
           throw e;
         }
         logger.log(Level.FINE, "Retryable exception, will sleep and retry", e);
-        backoffSleep(context, backOff);
+        long delay = e.getRetryDelayInMillis();
+        if (delay != -1) {
+          backoffSleep(context, delay);
+        } else {
+          backoffSleep(context, backOff);
+        }
       } catch (Exception e) {
         Throwables.throwIfUnchecked(e);
         throw newSpannerException(ErrorCode.INTERNAL, "Unexpected exception thrown", e);
@@ -2396,7 +2401,12 @@ protected PartialResultSet computeNext() {
             assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
             stream = null;
             try (Scope s = tracer.withSpan(span)) {
-              backoffSleep(context, backOff);
+              long delay = e.getRetryDelayInMillis();
+              if (delay != -1) {
+                backoffSleep(context, delay);
+              } else {
+                backoffSleep(context, backOff);
+              }
             }
             continue;
           }

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

@@ -51,6 +51,34 @@ public void connectionClosedIsRetryable() {
     assertThat(e.isRetryable()).isTrue();
   }
 
+  @Test
+  public void resourceExhausted() {
+    Status status =
+        Status.fromCodeValue(Status.Code.RESOURCE_EXHAUSTED.value())
+            .withDescription("Memory pushback");
+    SpannerException e =
+        SpannerExceptionFactory.newSpannerException(new StatusRuntimeException(status));
+    assertThat(e.isRetryable()).isFalse();
+  }
+
+  @Test
+  public void resourceExhaustedWithBackoff() {
+    Status status =
+        Status.fromCodeValue(Status.Code.RESOURCE_EXHAUSTED.value())
+            .withDescription("Memory pushback");
+    Metadata trailers = new Metadata();
+    Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
+    RetryInfo retryInfo =
+        RetryInfo.newBuilder()
+            .setRetryDelay(Duration.newBuilder().setNanos(1000000).setSeconds(1L))
+            .build();
+    trailers.put(key, retryInfo);
+    SpannerException e =
+        SpannerExceptionFactory.newSpannerException(new StatusRuntimeException(status, trailers));
+    assertThat(e.isRetryable()).isTrue();
+    assertThat(e.getRetryDelayInMillis()).isEqualTo(1001);
+  }
+
   @Test
   public void abortWithRetryInfo() {
     Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());

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

@@ -64,7 +64,7 @@ public void setUp() throws Exception {
 
   @Test
   public void runAbort() {
-    runTransaction(createRetryException());
+    runTransaction(createRetryException(Status.Code.ABORTED));
     ArgumentCaptor<Long> backoffMillis = ArgumentCaptor.forClass(Long.class);
     verify(sleeper, times(1)).backoffSleep(Mockito.<Context>any(), backoffMillis.capture());
     assertThat(backoffMillis.getValue()).isEqualTo(1001L);
@@ -81,7 +81,7 @@ public void runAbortNoRetryInfo() {
   @Test
   public void commitAbort() {
     final SpannerException error =
-        SpannerExceptionFactory.newSpannerException(createRetryException());
+        SpannerExceptionFactory.newSpannerException(createRetryException(Status.Code.ABORTED));
     when(rpc.commit(Mockito.<CommitRequest>any(), Mockito.<Map<Option, ?>>any()))
         .thenThrow(error)
         .thenReturn(
@@ -106,9 +106,15 @@ public Void run(TransactionContext transaction) throws Exception {
     assertThat(backoffMillis.getValue()).isEqualTo(1001L);
   }
 
-  private StatusRuntimeException createRetryException() {
+  @Test(expected = SpannerException.class)
+  public void runResourceExhaustedNoRetry() throws Exception {
+    runTransaction(
+        new StatusRuntimeException(Status.fromCodeValue(Status.Code.RESOURCE_EXHAUSTED.value())));
+  }
+
+  private StatusRuntimeException createRetryException(Status.Code code) {
     Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
-    Status status = Status.fromCodeValue(Status.Code.ABORTED.value());
+    Status status = Status.fromCodeValue(code.value());
     Metadata trailers = new Metadata();
     RetryInfo retryInfo =
         RetryInfo.newBuilder()