chore: Fix testCancelOuterFutureBeforeStart flaky test (#1583)

* chore: Fix testCancelOuterFutureBeforeStart flaky test

* chore: Fix assertions

* chore: Allow FailableCallable to have access to ExternalFuture

* chore: Fix format issues

* chore: Use MockCallable for this flaky test

* chore: Update comment

* chore: Update to the latest mockito version

* chore: Address PR comments

* chore: Add externalfuture for FailingCallable

* chore: Add comments to document the externalFuture

* chore: Move tracer verification to the end

Author: lqiu96

gax-java/gax/src/test/java/com/google/api/gax/retrying/AbstractRetryingExecutorTest.java

@@ -120,6 +120,7 @@ public void testSuccess() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureSuccess(future);
@@ -136,6 +137,7 @@ public void testSuccessWithFailures() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureSuccess(future);
@@ -153,6 +155,7 @@ public void testSuccessWithFailuresPeekGetAttempt() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
 
     assertNull(future.peekAttemptResult());
     assertSame(future.peekAttemptResult(), future.peekAttemptResult());
@@ -179,6 +182,7 @@ public void testMaxRetriesExceeded() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureFail(future, CustomException.class);
@@ -211,6 +215,7 @@ public void testTotalTimeoutExceeded() throws Exception {
       context = FakeCallContext.createDefault().withTracer(tracer);
     }
     RetryingFuture<String> future = executor.createFuture(callable, context);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureFail(future, CustomException.class);
@@ -222,28 +227,29 @@ public void testTotalTimeoutExceeded() throws Exception {
   }
 
   @Test
-  public void testCancelOuterFutureBeforeStart() throws Exception {
-    FailingCallable callable = new FailingCallable(4, "request", "SUCCESS", tracer);
+  public void testCancelOuterFutureBeforeStart() {
+    FailingCallable callable = new FailingCallable(0, "request", "FAILURE", tracer);
 
     RetrySettings retrySettings =
         FAST_RETRY_SETTINGS
             .toBuilder()
             .setInitialRetryDelay(Duration.ofMillis(1_000L))
             .setMaxRetryDelay(Duration.ofMillis(1_000L))
-            .setTotalTimeout(Duration.ofMillis(10_0000L))
+            .setTotalTimeout(Duration.ofMillis(10_000L))
             .build();
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(retrySettings, 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     boolean res = future.cancel(false);
-
     assertTrue(res);
+    assertTrue(future.isCancelled());
 
     future.setAttemptFuture(executor.submit(future));
-
-    assertFutureCancel(future);
     assertEquals(0, future.getAttemptSettings().getAttemptCount());
 
+    // Tracer should have 0 interactions as the externalFuture has already been cancelled
+    // FailingCallable should exit immediately and no tracer logic should be invoked
     verifyNoMoreInteractions(tracer);
   }
 
@@ -253,6 +259,7 @@ public void testCancelByRetryingAlgorithm() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 5, new CancellationException()));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureCancel(future);
@@ -272,6 +279,7 @@ public void testUnexpectedExceptionFromRetryAlgorithm() throws Exception {
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 5, new RuntimeException()));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureFail(future, RuntimeException.class);
@@ -311,6 +319,7 @@ public void testPollExceptionByPollAlgorithm() throws Exception {
       context = FakeCallContext.createDefault().withTracer(tracer);
     }
     RetryingFuture<String> future = executor.createFuture(callable, context);
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     assertFutureFail(future, PollException.class);

gax-java/gax/src/test/java/com/google/api/gax/retrying/FailingCallable.java

@@ -30,6 +30,7 @@
 package com.google.api.gax.retrying;
 
 import com.google.api.gax.tracing.ApiTracer;
+import com.google.common.base.Preconditions;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -66,20 +67,36 @@ class FailingCallable implements Callable<String> {
   private final String result;
   private final CountDownLatch firstAttemptFinished = new CountDownLatch(1);
 
+  private volatile RetryingFuture<String> externalFuture;
+
   FailingCallable(int expectedFailuresCount, String request, String result, ApiTracer tracer) {
     this.request = request;
     this.tracer = tracer;
     this.expectedFailuresCount = expectedFailuresCount;
     this.result = result;
   }
 
+  // ExternalFuture should always be passed in when using the FailingCallable
+  // to accurately mimic the behavior of AttemptCallable. We use the external
+  // future to check that the future is done and that none of the callable's
+  // logic is run.
+  public void setExternalFuture(RetryingFuture<String> externalFuture) {
+    this.externalFuture = Preconditions.checkNotNull(externalFuture);
+  }
+
   CountDownLatch getFirstAttemptFinishedLatch() {
     return firstAttemptFinished;
   }
 
   @Override
   public String call() throws Exception {
     try {
+      // Assumption is that externalFuture is always passed in.
+      // No null check to confirm so that the test will fail if
+      // the externalFuture is not passed in
+      if (externalFuture.isDone()) {
+        return null;
+      }
       int attemptNumber = attemptsCount.getAndIncrement();
 
       tracer.attemptStarted(request, attemptNumber);

gax-java/gax/src/test/java/com/google/api/gax/retrying/ScheduledRetryingExecutorTest.java

@@ -101,6 +101,7 @@ public void testSuccessWithFailuresPeekAttempt() throws Exception {
           getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
       RetryingFuture<String> future =
           executor.createFuture(callable, FakeCallContext.createDefault().withTracer(tracer));
+      callable.setExternalFuture(future);
 
       assertNull(future.peekAttemptResult());
       assertSame(future.peekAttemptResult(), future.peekAttemptResult());
@@ -151,6 +152,7 @@ public void testSuccessWithFailuresGetAttempt() throws Exception {
           getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
       RetryingFuture<String> future =
           executor.createFuture(callable, FakeCallContext.createDefault().withTracer(tracer));
+      callable.setExternalFuture(future);
 
       assertNull(future.peekAttemptResult());
       assertSame(future.getAttemptResult(), future.getAttemptResult());
@@ -203,6 +205,7 @@ public void testCancelGetAttempt() throws Exception {
       RetryingExecutorWithContext<String> executor =
           getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
       RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
+      callable.setExternalFuture(future);
 
       assertNull(future.peekAttemptResult());
       assertSame(future.getAttemptResult(), future.getAttemptResult());
@@ -261,6 +264,7 @@ public void testCancelOuterFutureAfterStart() throws Exception {
           getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
       RetryingFuture<String> future =
           executor.createFuture(callable, FakeCallContext.createDefault().withTracer(tracer));
+      callable.setExternalFuture(future);
       future.setAttemptFuture(executor.submit(future));
 
       Thread.sleep(30L);
@@ -288,6 +292,7 @@ public void testCancelIsTraced() throws Exception {
         getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
     RetryingFuture<String> future =
         executor.createFuture(callable, FakeCallContext.createDefault().withTracer(tracer));
+    callable.setExternalFuture(future);
     future.setAttemptFuture(executor.submit(future));
 
     Thread.sleep(30L);
@@ -317,6 +322,7 @@ public void testCancelProxiedFutureAfterStart() throws Exception {
           getRetryingExecutor(getAlgorithm(retrySettings, 0, null), localExecutor);
       RetryingFuture<String> future =
           executor.createFuture(callable, FakeCallContext.createDefault().withTracer(tracer));
+      callable.setExternalFuture(future);
       future.setAttemptFuture(executor.submit(future));
 
       Thread.sleep(50L);