Defer to timing-related exception in createNextAttempt instead of shouldRetry

Author: whowes

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetryAlgorithm.java

@@ -153,19 +153,25 @@ public TimedAttemptSettings createNextAttempt(
       Throwable previousThrowable,
       ResponseT previousResponse,
       TimedAttemptSettings previousSettings) {
-    // a small optimization that avoids calling relatively heavy methods
-    // like timedAlgorithm.createNextAttempt(), when it is not necessary.
-    if (!shouldRetryBasedOnResult(context, previousThrowable, previousResponse)) {
-      return null;
+    if (shouldRetryBasedOnResult(context, previousThrowable, previousResponse)) {
+      TimedAttemptSettings newSettings =
+          createNextAttemptBasedOnResult(
+              context, previousThrowable, previousResponse, previousSettings);
+      if (newSettings == null) {
+        newSettings = createNextAttemptBasedOnTiming(context, previousSettings);
+      }
+      return newSettings;
     }
-
-    TimedAttemptSettings newSettings =
-        createNextAttemptBasedOnResult(
-            context, previousThrowable, previousResponse, previousSettings);
-    if (newSettings == null) {
-      newSettings = createNextAttemptBasedOnTiming(context, previousSettings);
+    // If we shouldn't retry based on result and the last attempt failed with an
+    // exception, defer to any exception thrown by shouldRetryBasedOnTiming.
+    // This lets a timeout-related exception get propagated when justified
+    // rather than e.g. exceptions caused by very short RPC deadlines on a
+    // final polling attempt.
+    if (previousThrowable != null) {
+      TimedAttemptSettings nextSettings = createNextAttemptBasedOnTiming(context, previousSettings);
+      shouldRetryBasedOnTiming(context, nextSettings);
     }
-    return newSettings;
+    return null;
   }
 
   private TimedAttemptSettings createNextAttemptBasedOnResult(
@@ -232,18 +238,8 @@ public boolean shouldRetry(
       ResponseT previousResponse,
       TimedAttemptSettings nextAttemptSettings)
       throws CancellationException {
-    boolean retryBasedOnResult =
-        shouldRetryBasedOnResult(context, previousThrowable, previousResponse);
-
-    // Short-circuit when operation has succeeded to avoid erroneously throwing
-    // CancellationException below
-    if (!retryBasedOnResult && previousThrowable == null) {
-      return false;
-    }
-
-    // throws CancellationException
-    boolean retryBasedOnTiming = shouldRetryBasedOnTiming(context, nextAttemptSettings);
-    return retryBasedOnResult && retryBasedOnTiming;
+    return shouldRetryBasedOnResult(context, previousThrowable, previousResponse)
+        && shouldRetryBasedOnTiming(context, nextAttemptSettings);
   }
 
   boolean shouldRetryBasedOnResult(

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/retrying/RetryAlgorithmTest.java

@@ -30,11 +30,14 @@
 package com.google.api.gax.retrying;
 
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.util.concurrent.CancellationException;
 import org.junit.jupiter.api.Test;
 
 @SuppressWarnings({"unchecked", "deprecation"})
@@ -114,6 +117,45 @@ void testNextAttemptWithContext() {
     verify(resultAlgorithm).shouldRetry(context, previousThrowable, previousResult);
   }
 
+  @Test
+  void testCreateNextAttempt_exceptionAndTimeout_defersToTimingException() {
+    ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);
+    TimedRetryAlgorithm timedAlgorithm = mock(TimedRetryAlgorithm.class);
+    RetryAlgorithm<Object> algorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm);
+    Throwable throwable = new Throwable();
+    Object result = new Object();
+    TimedAttemptSettings previousSettings = mock(TimedAttemptSettings.class);
+    TimedAttemptSettings nextSettings = mock(TimedAttemptSettings.class);
+
+    when(resultAlgorithm.shouldRetry(throwable, result)).thenReturn(false);
+    when(timedAlgorithm.createNextAttempt(previousSettings)).thenReturn(nextSettings);
+    when(timedAlgorithm.shouldRetry(nextSettings)).thenThrow(new CancellationException());
+
+    assertThrows(
+        CancellationException.class,
+        () -> algorithm.createNextAttempt(throwable, result, previousSettings));
+  }
+
+  @Test
+  void testCreateNextAttempt_exceptionAndNoTimeout_returnsNull() {
+    ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);
+    TimedRetryAlgorithm timedAlgorithm = mock(TimedRetryAlgorithm.class);
+    RetryAlgorithm<Object> algorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm);
+    Throwable throwable = new Throwable();
+    Object result = new Object();
+    TimedAttemptSettings previousSettings = mock(TimedAttemptSettings.class);
+    TimedAttemptSettings nextSettings = mock(TimedAttemptSettings.class);
+
+    when(resultAlgorithm.shouldRetry(throwable, result)).thenReturn(false);
+    when(timedAlgorithm.createNextAttempt(previousSettings)).thenReturn(nextSettings);
+    when(timedAlgorithm.shouldRetry(nextSettings)).thenReturn(true);
+
+    TimedAttemptSettings nextAttemptSettings =
+        algorithm.createNextAttempt(throwable, result, previousSettings);
+
+    assertNull(nextAttemptSettings);
+  }
+
   @Test
   void testShouldRetry() {
     ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);
@@ -203,23 +245,6 @@ void testShouldRetryWithContext_noPreviousSettings() {
     assertFalse(algorithm.shouldRetry(context, previousThrowable, previousResult, null));
   }
 
-  @Test
-  void testShouldRetry_resultFalseAndThrowableNotNull_callsTimed() {
-    ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);
-    TimedRetryAlgorithm timedAlgorithm = mock(TimedRetryAlgorithm.class);
-    RetryAlgorithm<Object> algorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm);
-    Throwable previousThrowable = new Throwable();
-    Object previousResult = new Object();
-    TimedAttemptSettings previousSettings = mock(TimedAttemptSettings.class);
-    when(resultAlgorithm.shouldRetry(previousThrowable, previousResult)).thenReturn(false);
-
-    boolean shouldRetry =
-        algorithm.shouldRetry(previousThrowable, previousResult, previousSettings);
-
-    assertFalse(shouldRetry);
-    verify(timedAlgorithm).shouldRetry(previousSettings);
-  }
-
   @Test
   void testShouldRetry_resultFalseAndThrowableNull_shortCircuits() {
     ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/StreamingRetryAlgorithmTest.java

@@ -116,12 +116,17 @@ void testNextAttemptReturnsNullWhenShouldNotRetry() {
     StreamingRetryAlgorithm<String> algorithm =
         new StreamingRetryAlgorithm<>(resultAlgorithm, timedAlgorithm);
 
+    TimedAttemptSettings previousSettings = mock(TimedAttemptSettings.class);
+    when(previousSettings.getGlobalSettings()).thenReturn(DEFAULT_RETRY_SETTINGS);
+    when(previousSettings.getRpcTimeoutDuration())
+        .thenReturn(DEFAULT_RETRY_SETTINGS.getInitialRpcTimeoutDuration());
+
     TimedAttemptSettings attempt =
-        algorithm.createNextAttempt(context, exception, null, mock(TimedAttemptSettings.class));
+        algorithm.createNextAttempt(context, exception, null, previousSettings);
     assertThat(attempt).isNull();
 
     TimedAttemptSettings attemptWithoutContext =
-        algorithm.createNextAttempt(exception, null, mock(TimedAttemptSettings.class));
+        algorithm.createNextAttempt(exception, null, previousSettings);
     assertThat(attemptWithoutContext).isNull();
   }