fix(gax): adjust RetryAlgorithm logic for timeout v. exception

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(

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

@@ -30,10 +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"})
@@ -113,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);
@@ -201,4 +244,19 @@ void testShouldRetryWithContext_noPreviousSettings() {
 
     assertFalse(algorithm.shouldRetry(context, previousThrowable, previousResult, null));
   }
+
+  @Test
+  void testShouldRetry_resultFalseAndThrowableNull_shortCircuits() {
+    ResultRetryAlgorithm<Object> resultAlgorithm = mock(ResultRetryAlgorithm.class);
+    TimedRetryAlgorithm timedAlgorithm = mock(TimedRetryAlgorithm.class);
+    RetryAlgorithm<Object> algorithm = new RetryAlgorithm<>(resultAlgorithm, timedAlgorithm);
+    Object previousResult = new Object();
+    TimedAttemptSettings previousSettings = mock(TimedAttemptSettings.class);
+    when(resultAlgorithm.shouldRetry(null, previousResult)).thenReturn(false);
+
+    boolean shouldRetry = algorithm.shouldRetry(null, previousResult, previousSettings);
+
+    assertFalse(shouldRetry);
+    verify(timedAlgorithm, never()).shouldRetry(previousSettings);
+  }
 }

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();
   }